//
//  CPUAttention.cpp
//  MNN
//
//  Created by MNN on 2024/03/19.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#ifdef MNN_SUPPORT_TRANSFORMER_FUSE

#include <limits>
#include "CPUAttention.hpp"
#include "CPUBackend.hpp"
#include "compute/CommonOptFunction.h"
#include "compute/TurboQuant.hpp"
#include "core/Macro.h"
#include "core/Concurrency.h"
#include "core/BufferAllocator.hpp"
#include "core/TensorUtils.hpp"
#include "core/OpCommonUtils.hpp"
#include "core/BufferAllocator.hpp"
#include "compute/ConvolutionTiledExecutor.hpp"

#if defined(__aarch64__)
#define FLOAT16_T __fp16
#else
#define FLOAT16_T float
#endif

namespace MNN {

template <typename T>
static void _maskQK(float* qkPacked, const float* scale, size_t seqLen, size_t processedKvSeq, int pack, int kvSeqLen,
                    int kvoffset, int padKvSeqLen, const float* sinksPtr, const Tensor* mask, bool quantKey,
                    bool isLowerTriangular) {
    /*
     * FIGURE 1: mask->elementSize() == seqLen * maskStride
     * Context: Cross Attention or Prefill stage (Full Context).
     * Logic:   gapLen = 0. The mask tensor dimensions match the logical QK matrix exactly.
     *          Direct access: mask[row * stride + col]
     * Row\Col   0   1   2   3
     *
     *   0       0   X   X   X    (Can only see Col 0)
     *
     *   1       0   0   X   X    (Can see Col 0, 1)
     *
     *   2       0   0   0   X    (Can see Col 0, 1, 2)
     *
     *   3       0   0   0   0    (Fully visible)
     *
     * Legend:
     *   '0' : Visible (Value = Scale * QK)
     *   'X' : Masked  (Value = -inf)
     */

    /*
     * FIGURE 2: mask->elementSize() != seqLen * maskStride
     * Context: Self-Attention Inference (Decoding stage).
     * Logic:   gapLen = maskStride - seqLen (Right Alignment).
     *          The "Gap" represents History KV Cache, which is implicitly visible.
     *          The Mask Tensor only covers the current sequence window.
     *
     * Example: maskStride (Total KV) = 6
     *          seqLen (Current Q)    = 4
     *          gapLen                = 6 - 4 = 2
     *
     * Structure:
     *   - Cols [0, 1]: "Gap" / History region. Code logic: `if (col < gapLen) continue;`.
     *                  No mask is added, so they remain Visible ('0').
     *   - Cols [2-5]:  "Current" region. Code logic: `mask[col - gapLen]`.
     *
     * Row\Col   0   1   |   2   3   4   5
     *          (Gap)    |   (Mask Tensor Region)
     *
     *   0       0   0   |   0   X   X   X    <-- Mask row 0 applies to Col 2~5
     *                   |
     *   1       0   0   |   0   0   X   X    <-- Mask row 1 applies to Col 2~5
     *                   |
     *   2       0   0   |   0   0   0   X    <-- Mask row 2 applies to Col 2~5
     *                   |
     *   3       0   0   |   0   0   0   0    <-- Mask row 3 applies to Col 2~5
     *
     * Legend:
     *   '0' (Left)  : History KV, implicitly visible (code skips mask addition).
     *   '0' (Right) : Current KV, visible according to Mask Tensor.
     *   'X'         : Masked by Mask Tensor (-inf).
     */

    if (isLowerTriangular && quantKey) {
        return;
    }
    constexpr float NEG_INF = -std::numeric_limits<float>::infinity();
    auto source = (T*)qkPacked;
    float scaleVal = scale[0];

    auto kvBlockCount = UP_DIV(processedKvSeq, pack);
    auto qkSize = ROUND_UP(processedKvSeq, pack) * seqLen;

    if (isLowerTriangular) {
        for (int i = 0; i < qkSize; ++i) {
            source[i] *= scaleVal;
        }
        return;
    }

    if (mask == nullptr) {
        return;
    }

    int gapLen = (mask->elementSize() == (seqLen + padKvSeqLen) * (kvSeqLen + padKvSeqLen))
                     ? 0
                     : static_cast<int>(kvSeqLen - seqLen);
    auto maskPtr = mask->host<T>();
    auto maskCols = (mask->elementSize() == (seqLen + padKvSeqLen) * (kvSeqLen + padKvSeqLen)) ? kvSeqLen + padKvSeqLen
                                                                                               : seqLen + padKvSeqLen;
    for (int i = 0; i < kvBlockCount; ++i) {
        T* blockDataPtr = source + (i * seqLen * pack);

        for (int j = 0; j < seqLen; ++j) {
            T* dataPtr = blockDataPtr + (j * pack);
            const T* currentMaskRow = maskPtr + j * maskCols;

            for (int k = 0; k < pack; ++k) {
                float val = (float)dataPtr[k];
                if (!quantKey) {
                    val *= scaleVal;
                    dataPtr[k] = (T)val;
                }
                int currentKvSeqIndx = kvoffset + i * pack + k; // kvoffset=i*mBlockKv

                if (currentKvSeqIndx < gapLen) {
                    continue;
                }
                if (currentKvSeqIndx - gapLen >= maskCols) {
                    break;
                }

                val += (float)currentMaskRow[currentKvSeqIndx - gapLen];
                dataPtr[k] = (T)val;
            }
        }
    }
}

ErrorCode CPUAttention::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto gcore = static_cast<CPUBackend*>(backend())->functions();
    auto core = static_cast<CPUBackend*>(backend())->int8Functions();
    gcore->MNNGetMatMulPackMode(&eP, &lP, &hP);
    mThreadNum = ((CPUBackend*)backend())->threadNumber();
    mPack = gcore->pack;
    mBytes = gcore->bytes;
    int attentionOption = static_cast<CPUBackend*>(backend())->getRuntime()->hint().attentionOption;
    mUseFlashAttention = (attentionOption / 8 == 1);

    // attentionOption % 8:
    // 0: no quant, 1: K int8, 2: K+V int8, 3: K TQ3, 4: K+V TQ3, 5: K TQ4, 6: K+V TQ4
    int quantMode = attentionOption % 8;
    mKeyQuantMode = KVQuantMode::None;
    mValueQuantMode = KVQuantMode::None;
    if (inputs.size() < 5) {
        switch (quantMode) {
            case 1:
                mKeyQuantMode = KVQuantMode::Int8;
                break;
            case 2:
                mKeyQuantMode = KVQuantMode::Int8;
                mValueQuantMode = KVQuantMode::Int8;
                break;
            case 3:
                mKeyQuantMode = KVQuantMode::TQ3;
                break;
            case 4:
                mKeyQuantMode = KVQuantMode::TQ3;
                mValueQuantMode = KVQuantMode::TQ3;
                break;
            case 5:
                mKeyQuantMode = KVQuantMode::TQ4;
                break;
            case 6:
                mKeyQuantMode = KVQuantMode::TQ4;
                mValueQuantMode = KVQuantMode::TQ4;
                break;
            default:
                break;
        }
        if (mValueQuantMode == KVQuantMode::Int8 && !mUseFlashAttention) {
            mValueQuantMode = KVQuantMode::None;
        }
    }
    static_cast<CPUBackend*>(backend())->int8Functions()->MNNGetGemmUnit(&hP8, &lP8, &eP8);

    auto query = inputs[0];
    auto key = inputs[1];
    int seqLen = query->length(1);
    int mBlockNum = 1;
    mNumHead = query->length(2);
    mHeadDim = query->length(3);
    mKvNumHead = key->length(2);
    if (!mIsKVShared) {
        mKVCacheManager->setKVQuantMode(mUseFlashAttention, mKeyQuantMode, mValueQuantMode);
        mKVCacheManager->onResize(mKvNumHead, mHeadDim);
    }

    // Common buffer allocated
    auto bufferAlloc = static_cast<CPUBackend*>(backend())->getBufferAllocator();
    mPackQKV.reset(Tensor::createDevice<int8_t>({mThreadNum, UP_DIV(mHeadDim, mPack), seqLen, mPack * mBytes}));
    backend()->onAcquireBuffer(mPackQKV.get(), Backend::DYNAMIC);
    if (inputs.size() > 4 || mUseFlashAttention) { // needed by flash attention and sliding attention with sink
        mRunningMax.reset(Tensor::createDevice<int8_t>({mThreadNum, seqLen * 4}));
        mRunningSum.reset(Tensor::createDevice<int8_t>({mThreadNum, seqLen * 4}));
        backend()->onAcquireBuffer(mRunningMax.get(), Backend::DYNAMIC);
        backend()->onAcquireBuffer(mRunningSum.get(), Backend::DYNAMIC);
    }
    if (mUseFlashAttention) { // extra buffer need by flash attention
        mExpfDiffMax.reset(Tensor::createDevice<int8_t>({mThreadNum, seqLen * 4}));
        mTempOut.reset(Tensor::createDevice<int8_t>({mThreadNum, UP_DIV(mHeadDim, mPack), seqLen, mPack * mBytes}));
        backend()->onAcquireBuffer(mExpfDiffMax.get(), Backend::DYNAMIC);
        backend()->onAcquireBuffer(mTempOut.get(), Backend::DYNAMIC);
    }
    if (mKeyQuantMode == KVQuantMode::TQ3 || mKeyQuantMode == KVQuantMode::TQ4 || mValueQuantMode == KVQuantMode::TQ3 ||
        mValueQuantMode == KVQuantMode::TQ4) {
        // Vec_dot fusion buffers (per thread, shared by TQ3/TQ4):
        // Q_rotated: seqLen * headDim floats (WHT_forward of scaled Q)
        // V_acc_rotated: headDim floats (accumulator in rotated domain)
        // weights: blockKV floats (extracted softmax weights for one query)
        int blockKV = mUseFlashAttention ? MNN_FLASH_ATTENTION_BLOCK_SIZE : (seqLen + 64);
        int qRotatedSize = seqLen * mHeadDim * sizeof(float);
        int vAccSize = mHeadDim * sizeof(float);
        int weightsSize = blockKV * sizeof(float);
        mTQ3DequantBuf.reset(Tensor::createDevice<int8_t>({mThreadNum, qRotatedSize + vAccSize + weightsSize}));
        backend()->onAcquireBuffer(mTQ3DequantBuf.get(), Backend::DYNAMIC);
    }
    if (mKeyQuantMode == KVQuantMode::Int8) {
        int outterSeqLen = UP_DIV(seqLen, eP8);
        int outterHeadDim = UP_DIV(mHeadDim, lP8);

        size_t packedQSize = 0;
        if (outterSeqLen > 0) {
            int fullSeqBlocks = (seqLen / eP8);
            packedQSize += (size_t)fullSeqBlocks * outterHeadDim * eP8 * lP8;

            int lastEUnit = seqLen % eP8;
            if (lastEUnit != 0) {
                packedQSize += (size_t)outterHeadDim * lastEUnit * lP8;
            }
        }
        mPackQ.reset(Tensor::createDevice<int8_t>({mNumHead, (int32_t)packedQSize}));
        backend()->onAcquireBuffer(mPackQ.get(), Backend::DYNAMIC);

        mSumQ = bufferAlloc->alloc(mThreadNum * ROUND_UP(seqLen, eP8) * mBlockNum * sizeof(int32_t));
        mQueryScale = bufferAlloc->alloc(mNumHead * seqLen * mBlockNum * QUANT_INFO_BYTES);
        mQueryZeroPoint = bufferAlloc->alloc(mNumHead * seqLen * mBlockNum * QUANT_INFO_BYTES);
        mQueryQuantZero = bufferAlloc->alloc(mNumHead * seqLen * mBlockNum * QUANT_INFO_BYTES);
        mQueryQuantScale = bufferAlloc->alloc(mNumHead * seqLen * mBlockNum * QUANT_INFO_BYTES);
        mQuantQuery = bufferAlloc->alloc(seqLen * mNumHead * UP_DIV(mHeadDim, gcore->pack) * gcore->pack);

        if (mBlockNum > 1) {
            mAccumBuffer = bufferAlloc->alloc(eP8 * hP8 * mThreadNum * QUANT_INFO_BYTES);
            if (mAccumBuffer.invalid()) {
                return OUT_OF_MEMORY;
            }
        }

        if (mSumQ.invalid() || mQueryScale.invalid() || mQueryQuantZero.invalid() || mQueryZeroPoint.invalid() ||
            mQueryQuantScale.invalid() || mQuantQuery.invalid()) {
            return OUT_OF_MEMORY;
        }

        // post parameters for int8 gemm
        mGemmRelu.reset(2 * sizeof(int32_t));
        if (!mGemmRelu.get()) {
            MNN_ERROR("Allocate mGemmRelu buffer failed in CPU Attention");
            return OUT_OF_MEMORY;
        }
        ((float*)mGemmRelu.get())[0] = -std::numeric_limits<float>().max();
        ((float*)mGemmRelu.get())[1] = std::numeric_limits<float>().max();
        if (mBytes == 2) {
            gcore->MNNFp32ToLowp((float*)mGemmRelu.get(), reinterpret_cast<int16_t*>(mGemmRelu.get()), 2);
        }

        // GemmInt8 kernels
        if (mBytes == 4) {
            mInt8GemmKernel = core->Int8GemmKernel;
        } else {
            mInt8GemmKernel = core->MNNGemmInt8AddBiasScale_Unit_FP16;
        }

        if (mValueQuantMode == KVQuantMode::Int8) {
            mQuantQK = bufferAlloc->alloc(mThreadNum * eP8 * ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, mPack));
            mQKScale = bufferAlloc->alloc(eP8 * QUANT_INFO_BYTES);
            mQKBias = bufferAlloc->alloc(eP8 * QUANT_INFO_BYTES);
            mSumQK = bufferAlloc->alloc(mThreadNum * eP8 * QUANT_INFO_BYTES);

            if (mQuantQK.invalid() || mQKScale.invalid() || mQKBias.invalid() || mSumQK.invalid()) {
                return OUT_OF_MEMORY;
            }
        }
    } else {
        mPackQ.reset(
            Tensor::createDevice<int8_t>({mThreadNum, UP_DIV(seqLen, eP), ROUND_UP(mHeadDim, lP), eP * mBytes}));
        backend()->onAcquireBuffer(mPackQ.get(), Backend::DYNAMIC);
        backend()->onAcquireBuffer(mPackQKV.get(), Backend::DYNAMIC);
    }

    // release tensor
    backend()->onReleaseBuffer(mPackQ.get(), Backend::DYNAMIC);
    backend()->onReleaseBuffer(mPackQKV.get(), Backend::DYNAMIC);

    if (inputs.size() > 4 || mUseFlashAttention) {
        backend()->onReleaseBuffer(mRunningMax.get(), Backend::DYNAMIC);
        backend()->onReleaseBuffer(mRunningSum.get(), Backend::DYNAMIC);
    }
    if (mUseFlashAttention) {
        backend()->onReleaseBuffer(mExpfDiffMax.get(), Backend::DYNAMIC);
        backend()->onReleaseBuffer(mTempOut.get(), Backend::DYNAMIC);
    }
    if (mKeyQuantMode == KVQuantMode::TQ3 || mKeyQuantMode == KVQuantMode::TQ4 || mValueQuantMode == KVQuantMode::TQ3 ||
        mValueQuantMode == KVQuantMode::TQ4) {
        backend()->onReleaseBuffer(mTQ3DequantBuf.get(), Backend::DYNAMIC);
    }

    // release memchunk
    if (mKeyQuantMode == KVQuantMode::Int8) {
        bufferAlloc->free(mSumQ);
        bufferAlloc->free(mQueryScale);
        bufferAlloc->free(mQueryZeroPoint);
        bufferAlloc->free(mQueryQuantScale);
        bufferAlloc->free(mQueryQuantZero);
        bufferAlloc->free(mQuantQuery);
        if (mBlockNum > 1) {
            bufferAlloc->free(mAccumBuffer);
        }
        if (mValueQuantMode == KVQuantMode::Int8) {
            bufferAlloc->free(mQuantQK);
            bufferAlloc->free(mQKScale);
            bufferAlloc->free(mQKBias);
            bufferAlloc->free(mSumQK);
        }
    }

    // Only allocated for quantized Q&K
    if (mKeyQuantMode == KVQuantMode::Int8) {
        if (mBytes == 4) {
            mQuantFunc = core->MNNFloat2Int8;
        } else {
            mQuantFunc = core->DynamicQuanInput_ARM82;
        }
    }
    return NO_ERROR;
}

ErrorCode CPUAttention::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto gcore = static_cast<CPUBackend*>(backend())->functions();
    auto core = static_cast<CPUBackend*>(backend())->int8Functions();
    bool outputC4 = TensorUtils::getDescribe(outputs[0])->dimensionFormat == MNN_DATA_FORMAT_NC4HW4;
    auto query = inputs[0];
    auto key = inputs[1];
    auto value = inputs[2];
    int seqLen = query->length(1);
    const Tensor* mask = nullptr;
    if (inputs.size() > 3) {
        mask = inputs[3];
    }
    const Tensor* sinks = nullptr;
    if (inputs.size() > 4) {
        sinks = inputs[4];
        MNN_ASSERT(sinks != nullptr);
        MNN_ASSERT(sinks->elementSize() == mNumHead)
    }
    int numHeadDiv = UP_DIV(mNumHead, mThreadNum);
    int group_size = mNumHead / mKvNumHead;
    // reduce the value of 'query' to avoid fp16 overflow
    float mScale = (mMeta && mMeta->attn_scale > 0) ? mMeta->attn_scale : (1.0 / sqrt(mHeadDim));
    float q_scale = 1.0;
    if (mBytes == 2 && mKeyQuantMode != KVQuantMode::Int8) {
        // reduce the value of 'query' to 'query * FP16_QSCALE', avoid fp16 overflow
        FLOAT16_T minValue;
        FLOAT16_T maxValue;
        gcore->MNNCountMaxMinValue(query->host<float>(), (float*)(&minValue), (float*)(&maxValue),
                                   query->elementSize());
        float maxV = maxValue;
        float minV = minValue;
        float absMax = ALIMAX(fabsf(maxV), fabsf(minV));
        if (absMax > 1.0f) {
            q_scale = 1.0f / absMax;
        }
        mScale /= q_scale;
    }
    int insertLen = seqLen;

    if (!mIsKVShared) {
        if (mKVCache && mMeta != nullptr) {
            if (mMeta->previous == mMeta->remove) {
                mKVCacheManager->onClear();
                mKVCacheManager->onAlloc(mMeta, seqLen);
            } else {
                MNN_ASSERT(mMeta->previous == mKVCacheManager->kvLength());
                mKVCacheManager->onRealloc(mMeta);
            }
            insertLen = (int)mMeta->add;
        } else {
            mKVCacheManager->onClear();
            mKVCacheManager->onAlloc(mMeta, seqLen);
        }
        // Add the new kv to the kvcache
        mKVCacheManager->onUpdateKV(key, value, (int)insertLen);
    } else {
        // Shared layer: KV cache is shared via onClone, skip KV update
        insertLen = (int)mMeta->add;
    }

    if (mUseFlashAttention) {
        mBlockKV = ALIMIN(MNN_FLASH_ATTENTION_BLOCK_SIZE, mKVCacheManager->kvLength());
    } else {
        mBlockKV = mKVCacheManager->kvLength();
    }

    // Constant Initialization
    auto padSeqLength = seqLen - insertLen;
    seqLen = insertLen;
    int kvSeqLen = mKVCacheManager->kvLength();
    int maxLen = mKVCacheManager->maxLength();
    int32_t units[2] = {eP, lP};
    const float* sinksPtr = sinks ? sinks->host<float>() : nullptr;
    int kvValidOffset = kvSeqLen - seqLen; // reuse_kv=true or decode, kvValidOffset>0

    // Temporary tensors for intermediate results
    std::shared_ptr<Tensor> unpackQK(Tensor::createDevice<int32_t>({mThreadNum, seqLen, mBlockKV}));
    std::shared_ptr<Tensor> softmMaxQ(Tensor::createDevice<int32_t>(
        {mThreadNum, seqLen, ROUND_UP(mBlockKV, mPack)})); // [mBlockKV/mPack, seqLen, mPack ]
    std::shared_ptr<Tensor> newPackQK;
    if (mValueQuantMode != KVQuantMode::Int8) {
        newPackQK.reset(Tensor::createDevice<int8_t>({mThreadNum, eP * ROUND_UP(mBlockKV, lP) * mBytes}));
    } else {
        newPackQK.reset(
            Tensor::createDevice<int8_t>({mThreadNum, eP8 * ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP8)}));
    }
    std::shared_ptr<Tensor> mTempQKBlock(
        Tensor::createDevice<int8_t>({mThreadNum, UP_DIV(mBlockKV, mPack), seqLen, mPack * mBytes}));
    backend()->onAcquireBuffer(unpackQK.get(), Backend::STATIC);
    backend()->onAcquireBuffer(softmMaxQ.get(), Backend::STATIC);
    backend()->onAcquireBuffer(newPackQK.get(), Backend::STATIC);
    backend()->onAcquireBuffer(mTempQKBlock.get(), Backend::STATIC);

    // Quantize Q and initialize bias 0
    if (mKeyQuantMode == KVQuantMode::Int8) {
        mGemmBias.reset(ROUND_UP(ALIMAX(mBlockKV, mHeadDim), hP8) * QUANT_INFO_BYTES);
        if (!mGemmBias.get()) {
            MNN_ERROR("Allocate bias buffer failed in CPU Attention\n");
            return OUT_OF_MEMORY;
        }
        memset(mGemmBias.get(), 0, ROUND_UP(ALIMAX(mBlockKV, mHeadDim), hP8) * QUANT_INFO_BYTES);

        // Q: [seqLen,numHead,headDim]
        // maxQ, minQ: [seqLen,numHead]
        // scaleQ, zeroQ: [numHead, seqLen]
        // quantQ: [seqLen,numHead,headDim]
        auto queryPtr = query->host<int8_t>();
        int divPart = UP_DIV(seqLen * mNumHead, mThreadNum);
        MNN_CONCURRENCY_BEGIN(tId, mThreadNum) {
            size_t info[9] = {1, (size_t)mHeadDim, 1, 1, 1, 1, 1, 1, 0};
            auto remainLu = seqLen * mNumHead - tId * divPart;
            if (remainLu > 0) {
                remainLu = ALIMIN(divPart, remainLu);
                for (int i = tId * divPart; i < tId * divPart + remainLu; ++i) {
                    // address
                    auto srcFloatPtr = (float*)(queryPtr + i * mHeadDim * mBytes);
                    auto dstInt8Ptr = (int8_t*)(mQuantQuery.ptr() + i * mHeadDim);
                    auto quantScalePtr = (float*)(mQueryQuantScale.ptr() + i * QUANT_INFO_BYTES);
                    auto quantZeroPtr = (float*)(mQueryQuantZero.ptr() + i * QUANT_INFO_BYTES);

                    // scaleQ, zeroQ, [seqLen,numHead]->[numHead,seqLen]
                    int indexQ = (i / mNumHead) + (i % mNumHead) * seqLen;
                    auto scalePtr = (float*)(mQueryScale.ptr() + indexQ * QUANT_INFO_BYTES);
                    auto zeroPtr = (float*)(mQueryZeroPoint.ptr() + indexQ * QUANT_INFO_BYTES);

                    // compute the quant/dequant scale/bias
                    gcore->MNNAsyQuantInfo(scalePtr, zeroPtr, quantScalePtr, quantZeroPtr, nullptr, nullptr,
                                           srcFloatPtr, info);
                    scalePtr[0] *= mScale;
                    zeroPtr[0] *= mScale;

                    // quantize the float query to int8_t query
                    mQuantFunc(srcFloatPtr, dstInt8Ptr, UP_DIV(mHeadDim, gcore->pack), quantScalePtr, -128, 127,
                               quantZeroPtr, 0);
                }
            }
        }
        MNN_CONCURRENCY_END();

        // source int8_t query: [seqLen,numHead,headDim]
        // dest int8_t query: [numHead,seqLen/eP,headDim/lP,eP,lP]

        int outterSeqLen = UP_DIV(seqLen, eP8);
        int outterHeadDim = UP_DIV(mHeadDim, lP8);
        size_t outputOffset = 0;

        const int8_t* src_base_ptr = (const int8_t*)mQuantQuery.ptr();
        int8_t* dst_base_ptr = mPackQ->host<int8_t>();

        for (int h = 0; h < mNumHead; ++h) {
            for (int seqBlock = 0; seqBlock < outterSeqLen; ++seqBlock) {
                int seqBase = seqBlock * eP8;
                int eunit = std::min(eP8, seqLen - seqBase);
                size_t currentSeqBlockSize = (size_t)outterHeadDim * eunit * lP8;

                for (int dimBlock = 0; dimBlock < outterHeadDim; ++dimBlock) {
                    int dimBase = dimBlock * lP8;
                    int headDimRemain = mHeadDim - dimBase;
                    int copyLen = std::min(lP8, headDimRemain);

                    if (copyLen <= 0) {
                        continue;
                    }

                    int8_t* dst_block_ptr = dst_base_ptr + outputOffset + (size_t)dimBlock * (eunit * lP8);

                    const size_t src_row_stride = (size_t)mNumHead * mHeadDim;

                    for (int seqLocal = 0; seqLocal < eunit; ++seqLocal) {
                        int innerSeq = seqBase + seqLocal;

                        const int8_t* src_row_ptr =
                            src_base_ptr + (size_t)innerSeq * src_row_stride + (size_t)h * mHeadDim + dimBase;

                        int8_t* dst_row_ptr = dst_block_ptr + seqLocal * lP8;

                        std::memcpy(dst_row_ptr, src_row_ptr, copyLen);
                    }
                    if (copyLen < lP8) {
                        for (int seqLocal = 0; seqLocal < eunit; ++seqLocal) {
                            int8_t* dst_pad_ptr = dst_block_ptr + seqLocal * lP8 + copyLen;
                            std::memset(dst_pad_ptr, 0, lP8 - copyLen);
                        }
                    }
                }
                outputOffset += currentSeqBlockSize;
            }
        } // Finish quantize Q

        if (mValueQuantMode == KVQuantMode::Int8) {
            auto scalePtr = (float*)(mQKScale.ptr());
            auto zeroPtr = (float*)(mQKBias.ptr());
            for (int k = 0; k < eP8; ++k) {
                scalePtr[k] = 1.f / 255.f;
#ifdef MNN_USE_SSE
                zeroPtr[k] = 0;
#else
                zeroPtr[k] = 128.f / 255.f;
#endif
            }
        }
    }

    std::function<void(int)> mCompute = [=](int tId) {
        int8_t* qReordered = nullptr;
        auto qkPacked = mTempQKBlock->host<int8_t>() + tId * mTempQKBlock->stride(0);
        auto qkFlatten = unpackQK->host<float>() + tId * unpackQK->stride(0);
        auto qkSoftmax = softmMaxQ->host<float>() + tId * softmMaxQ->stride(0);
        auto qkReordered = newPackQK->host<int8_t>() + tId * newPackQK->stride(0);
        auto qkvPacked = mPackQKV->host<int8_t>() + tId * mPackQKV->stride(0);
        int headIndex = tId * numHeadDiv;
        int headsToCompute = ALIMIN(numHeadDiv, mNumHead - headIndex);

        // Flash Attention
        auto runningMax = mRunningMax ? (float*)(mRunningMax->host<int8_t>() + tId * mRunningMax->stride(0)) : nullptr;
        auto runningSum = mRunningSum ? (float*)(mRunningSum->host<int8_t>() + tId * mRunningSum->stride(0)) : nullptr;
        auto diffScale =
            mExpfDiffMax ? (float*)(mExpfDiffMax->host<int8_t>() + tId * mExpfDiffMax->stride(0)) : nullptr;
        auto outputPacked = mTempOut ? mTempOut->host<int8_t>() + tId * mTempOut->stride(0) : qkvPacked;

        int kvBlocks = UP_DIV(kvSeqLen, mBlockKV);

        bool isLowerTriangular = (mask == nullptr);
        if (mask != nullptr && mask->shape().empty()) {
            if (mBytes == 2) {
                auto maskPtr = mask->host<FLOAT16_T>();
                if (maskPtr[0] < 1e-6) {
                    isLowerTriangular = true;
                }
            } else {
                auto maskPtr = mask->host<float>();
                if (maskPtr[0] < 1e-6f) {
                    isLowerTriangular = true;
                }
            }
        }
        bool useMaskInSoftmax = (isLowerTriangular && sinksPtr == nullptr);

        QuanPostTreatParameters gemmParam4QxK, gemmParam4QKxV; // used by int8 gemm, allocated per thread.
        SumByAxisParams sumParams4QxK, sumParams4QKxV = {};
        float* qSumAddr = nullptr;
        float* qScale = nullptr;
        float* qBias = nullptr;
        float* accumbuff = nullptr;
        int32_t unitColBufferSize = 0;
        if (mKeyQuantMode == KVQuantMode::Int8) {
            // parameters shared by all mBlockKV
            gemmParam4QxK.blockNum = mBlockNum;
            gemmParam4QxK.biasFloat = reinterpret_cast<float*>(mGemmBias.get());
            gemmParam4QxK.useInt8 = 0;
            gemmParam4QxK.fp32minmax = reinterpret_cast<float*>(mGemmRelu.get());

            sumParams4QxK.oneScale = 0;
            sumParams4QxK.SRC_UNIT = lP8;
            sumParams4QxK.blockNum = mBlockNum;
            sumParams4QxK.DST_XUNIT = eP8;
            sumParams4QxK.inputBlock = 0;
            sumParams4QxK.kernelxy = 1;
            // fixed
            sumParams4QxK.LU = UP_DIV(mHeadDim, lP8);
            sumParams4QxK.unitColBufferSize = ROUND_UP(mHeadDim, lP8) * eP8;
            sumParams4QxK.kernelCountUnitDouble = UP_DIV(mHeadDim, lP8);
            sumParams4QxK.valid = mHeadDim % lP8;

            if (mBlockNum > 1) {
                accumbuff = (float*)(mAccumBuffer.ptr() + tId * eP8 * hP8 * QUANT_INFO_BYTES);
            }
            unitColBufferSize = eP8 * ROUND_UP(mHeadDim, lP8);

            if (mValueQuantMode == KVQuantMode::Int8) {
                gemmParam4QKxV.blockNum = mBlockNum;
                gemmParam4QKxV.biasFloat = reinterpret_cast<float*>(mGemmBias.get());
                gemmParam4QKxV.useInt8 = 0;
                gemmParam4QKxV.fp32minmax = reinterpret_cast<float*>(mGemmRelu.get());
                gemmParam4QKxV.inputScale = (float*)mQKScale.ptr();
                gemmParam4QKxV.inputBias = (float*)mQKBias.ptr();
                gemmParam4QKxV.srcKernelSum = (float*)(mSumQK.ptr() + tId * eP8 * QUANT_INFO_BYTES);

                sumParams4QKxV.oneScale = 0;
                sumParams4QKxV.SRC_UNIT = lP8;
                sumParams4QKxV.blockNum = mBlockNum;
                sumParams4QKxV.DST_XUNIT = eP8;
                sumParams4QKxV.inputBlock = 0;
                sumParams4QKxV.kernelxy = 1;
                sumParams4QKxV.unitColBufferSize = ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP8) * eP8;
                sumParams4QKxV.kernelCountUnitDouble = UP_DIV(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP8);
            }
        }

        size_t vstride0 = ROUND_UP(mHeadDim, hP) * ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP);
        if (mValueQuantMode == KVQuantMode::Int8) {
            vstride0 = (ROUND_UP(mHeadDim, hP8) * ROUND_UP(mKVCacheManager->getFlashAttentionBlockKv(), lP8) +
                        2 * QUANT_INFO_BYTES * mBlockNum * ROUND_UP(mHeadDim, hP8));
        }

        // use for V
        float const* srcPtr[1];
        // only used for quantized V
        float vQuantScale[1] = {255.f};
        float vQuantBias[1] = {-128.f};
        int32_t infoInt8V[5];
        infoInt8V[0] = 1; // number
        infoInt8V[2] = static_cast<int32_t>(sumParams4QKxV.unitColBufferSize);
        infoInt8V[3] = 1; // stride
        int32_t elInt8V[4] = {eP8, ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP8), 0, 0};

        // only used for float V
        int32_t infoFloatV[4];
        infoFloatV[0] = 1;      // number
        infoFloatV[1] = seqLen; // eReal
        infoFloatV[3] = 1;      // stride
        int32_t elFloatV[4] = {seqLen, ROUND_UP(kvSeqLen, lP), 0, 0};

        int offset[2] = {seqLen, mNumHead * mHeadDim};

        for (int h = headIndex; h < headIndex + headsToCompute; h++) {
            auto dstStep = mBytes * seqLen * mPack;
            if (outputC4) {
                outputPacked = outputs[0]->host<int8_t>() + h * mHeadDim * seqLen * mBytes;
                if (!mUseFlashAttention) {
                    qkvPacked = outputPacked;
                }
            }
            // Prepare for flash attention
            if (runningSum && runningMax) {
                if (sinksPtr == nullptr) {
                    memset(runningSum, 0, mRunningSum->stride(0));
                    for (int k = 0; k < seqLen; ++k) {
                        runningMax[k] = std::numeric_limits<float>::lowest();
                    }
                } else {
                    for (int k = 0; k < seqLen; ++k) {
                        runningSum[k] = 1.f; // exp(sink-sink)
                    }
                    float sinkVal;
                    if (mBytes == 2) {
                        sinkVal = ((FLOAT16_T*)sinksPtr)[h];
                    } else {
                        sinkVal = sinksPtr[h];
                    }
                    for (int k = 0; k < seqLen; ++k) {
                        runningMax[k] = sinkVal;
                    }
                }
            }

            // Compute the current addresses
            int kvHeadIndex = h / group_size;
            int8_t* keyAddr = mKVCacheManager->addrOfKey(kvHeadIndex);
            int8_t* keySum = mKVCacheManager->addrOfKeySum(kvHeadIndex);
            int8_t* valueAddr = mKVCacheManager->addrOfValue(kvHeadIndex);
            float* valueSum = (float*)mKVCacheManager->addrOfValueSum(kvHeadIndex);

            // Get packed Q
            if (mKeyQuantMode != KVQuantMode::Int8) {
                qReordered = mPackQ->host<int8_t>() + tId * mPackQ->stride(0);
                gcore->MNNAttenPackAndScaleSingleHead((float*)qReordered,
                                                      (float*)(query->host<int8_t>() + h * mHeadDim * mBytes),
                                                      mHeadDim * mNumHead, &q_scale, units, seqLen, mHeadDim);
            } else {
                qReordered = mPackQ->host<int8_t>() + h * mPackQ->stride(0);
                qSumAddr = (float*)(mSumQ.ptr() + tId * ROUND_UP(seqLen, eP8) * mBlockNum * QUANT_INFO_BYTES);
                qScale = (float*)(mQueryScale.ptr() + h * seqLen * mBlockNum * QUANT_INFO_BYTES);
                qBias = (float*)(mQueryZeroPoint.ptr() + h * seqLen * mBlockNum * QUANT_INFO_BYTES);
                gcore->MNNSumByAxisLForMatmul_A(qSumAddr, qReordered, qScale, seqLen, sumParams4QxK);
            }

            // Start computing
            for (int i = 0; i < kvBlocks; ++i) {
                int subKvSeqLen = ALIMIN(mBlockKV, kvSeqLen - i * mBlockKV);

                // 1. query @ key
                if (mKeyQuantMode == KVQuantMode::TQ3) {
                    // Vec_dot fusion: Q_rotated · TQ3_compressed_K directly (no dequant buffer)
                    // Q_rotated = WHT_forward(Q * q_scale) computed once per KV block iteration
                    int tq3BytesPerSeq = (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
                    int numBlocks = mHeadDim / TQ3_BLOCK_SIZE;
                    auto tq3Buf = mTQ3DequantBuf->host<int8_t>() + tId * mTQ3DequantBuf->stride(0);
                    auto qRotated = (float*)tq3Buf; // seqLen * headDim floats

                    // Pre-rotate Q vectors (only on first KV block)
                    if (i == 0) {
                        float qScale = 1.0f / sqrtf((float)mHeadDim);
                        auto queryBase = (float*)(query->host<int8_t>() + h * mHeadDim * mBytes);
                        int qStride = mHeadDim * mNumHead; // stride between seq positions
                        for (int q = 0; q < seqLen; q++) {
                            for (int b = 0; b < numBlocks; b++) {
                                float scaled[TQ3_BLOCK_SIZE];
                                if (mBytes == 2) {
                                    auto src16 = (FLOAT16_T*)(query->host<int8_t>() + h * mHeadDim * mBytes) +
                                                 q * mHeadDim * mNumHead;
                                    for (int d = 0; d < TQ3_BLOCK_SIZE; d++) {
                                        scaled[d] = (float)src16[b * TQ3_BLOCK_SIZE + d] * qScale;
                                    }
                                } else {
                                    auto srcF = queryBase + q * qStride;
                                    for (int d = 0; d < TQ3_BLOCK_SIZE; d++) {
                                        scaled[d] = srcF[b * TQ3_BLOCK_SIZE + d] * qScale;
                                    }
                                }
                                tq3_wht_forward_32(qRotated + q * mHeadDim + b * TQ3_BLOCK_SIZE, scaled);
                            }
                        }
                    }

                    // Compute QK scores directly: score[q][s] = Σ_b vec_dot_block(Q_rot, K_tq3)
                    // Output format: qkPacked [kvSeq/mPack, seqLen, mPack]
                    for (int s = 0; s < subKvSeqLen; s++) {
                        int seqIdx = i * mBlockKV + s;
                        auto kPtr = (uint8_t*)keyAddr + seqIdx * tq3BytesPerSeq;
                        for (int q = 0; q < seqLen; q++) {
                            float score = 0.0f;
                            auto qr = qRotated + q * mHeadDim;
                            for (int b = 0; b < numBlocks; b++) {
                                score += tq3_vec_dot_block(qr + b * TQ3_BLOCK_SIZE, kPtr + b * TQ3_BYTES_PER_BLOCK);
                            }
                            // Write to [kvSeq/mPack, seqLen, mPack] format
                            int packIdx = (s / mPack) * seqLen * mPack + q * mPack + s % mPack;
                            if (mBytes == 2) {
                                ((FLOAT16_T*)qkPacked)[packIdx] = (FLOAT16_T)score;
                            } else {
                                ((float*)qkPacked)[packIdx] = score;
                            }
                        }
                    }
                } else if (mKeyQuantMode == KVQuantMode::TQ4) {
                    // Vec_dot fusion for TQ4 (4-bit): same logic as TQ3, different bytesPerBlock + functions
                    int tq4BytesPerSeq = (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
                    int numBlocks = mHeadDim / TQ4_BLOCK_SIZE;
                    auto tq4Buf = mTQ3DequantBuf->host<int8_t>() + tId * mTQ3DequantBuf->stride(0);
                    auto qRotated = (float*)tq4Buf;

                    if (i == 0) {
                        float qScale = 1.0f / sqrtf((float)mHeadDim);
                        for (int q = 0; q < seqLen; q++) {
                            for (int b = 0; b < numBlocks; b++) {
                                float scaled[TQ4_BLOCK_SIZE];
                                if (mBytes == 2) {
                                    auto src16 = (FLOAT16_T*)(query->host<int8_t>() + h * mHeadDim * mBytes) +
                                                 q * mHeadDim * mNumHead;
                                    for (int d = 0; d < TQ4_BLOCK_SIZE; d++)
                                        scaled[d] = (float)src16[b * TQ4_BLOCK_SIZE + d] * qScale;
                                } else {
                                    auto srcF = (float*)(query->host<int8_t>() + h * mHeadDim * mBytes) +
                                                q * mHeadDim * mNumHead;
                                    for (int d = 0; d < TQ4_BLOCK_SIZE; d++)
                                        scaled[d] = srcF[b * TQ4_BLOCK_SIZE + d] * qScale;
                                }
                                tq3_wht_forward_32(qRotated + q * mHeadDim + b * TQ4_BLOCK_SIZE, scaled);
                            }
                        }
                    }

                    for (int s = 0; s < subKvSeqLen; s++) {
                        int seqIdx = i * mBlockKV + s;
                        auto kPtr = (uint8_t*)keyAddr + seqIdx * tq4BytesPerSeq;
                        for (int q = 0; q < seqLen; q++) {
                            float score = 0.0f;
                            auto qr = qRotated + q * mHeadDim;
                            for (int b = 0; b < numBlocks; b++) {
                                score += tq4_vec_dot_block(qr + b * TQ4_BLOCK_SIZE, kPtr + b * TQ4_BYTES_PER_BLOCK);
                            }
                            int packIdx = (s / mPack) * seqLen * mPack + q * mPack + s % mPack;
                            if (mBytes == 2) {
                                ((FLOAT16_T*)qkPacked)[packIdx] = (FLOAT16_T)score;
                            } else {
                                ((float*)qkPacked)[packIdx] = score;
                            }
                        }
                    }
                } else if (mKeyQuantMode != KVQuantMode::Int8) {
                    auto keyPtr = keyAddr + i * UP_DIV(mBlockKV, hP) * ROUND_UP(mHeadDim, lP) * hP * mBytes;
                    int loop_e = seqLen / eP;
                    int remain = seqLen % eP;
                    auto qStride0 = ROUND_UP(mHeadDim, lP) * eP * mBytes;
                    size_t shapeParameters[7] = {(size_t)eP * lP * mBytes,
                                                 ROUND_UP((size_t)mHeadDim, lP),
                                                 (size_t)subKvSeqLen,
                                                 (size_t)seqLen * mPack * mBytes,
                                                 0,
                                                 0,
                                                 0};
                    for (int ei = 0; ei < loop_e; ei++) {
                        gcore->MNNPackedMatMul((float*)(qkPacked + (ei * eP * mPack) * mBytes),
                                               (float*)(qReordered + ei * qStride0), (float*)keyPtr, shapeParameters,
                                               nullptr, nullptr, nullptr, nullptr);
                    }
                    if (remain > 0) {
                        gcore->MNNPackedMatMulRemain((float*)(qkPacked + (loop_e * eP * mPack) * mBytes),
                                                     (float*)(qReordered + loop_e * qStride0), (float*)keyPtr, remain,
                                                     shapeParameters, nullptr, nullptr, nullptr, nullptr);
                    }
                } else {
                    auto eRemain = seqLen;
                    auto srcInt8 = qReordered;
                    auto dstInt8 = qkPacked;
                    auto keyPtr = keyAddr + i * UP_DIV(mBlockKV, hP8) *
                                                (ROUND_UP(mHeadDim, lP8) * hP8 + 2 * hP8 * QUANT_INFO_BYTES);
                    gemmParam4QxK.weightKernelSum = (float*)(keySum + i * mBlockKV * QUANT_INFO_BYTES);
                    gemmParam4QxK.inputScale = qScale;
                    gemmParam4QxK.inputBias = qBias;
                    gemmParam4QxK.srcKernelSum = qSumAddr;
                    while (eRemain > 0) {
                        auto eSize = ALIMIN(eP8, eRemain);
                        mInt8GemmKernel(dstInt8, srcInt8, keyPtr, UP_DIV(mHeadDim, lP8), mBytes * seqLen * mPack,
                                        UP_DIV(subKvSeqLen, mPack), &gemmParam4QxK, eSize);
                        eRemain -= eP8;
                        gemmParam4QxK.inputScale += eP8;
                        gemmParam4QxK.inputBias += eP8;
                        gemmParam4QxK.srcKernelSum += eP8;
                        srcInt8 += unitColBufferSize;
                        dstInt8 += eP8 * mPack * mBytes;
                        if (mBlockNum > 1) {
                            memset(accumbuff, 0, eP8 * hP8 * QUANT_INFO_BYTES);
                            gemmParam4QxK.accumBuffer = accumbuff;
                        }
                    }
                }
                // 2. softmax scores, softmax src/dst shape: [kv_seq_len/mPack, seq_len, mPack]
                {
                    if (mKeyQuantMode != KVQuantMode::Int8 || isLowerTriangular == false || sinksPtr != nullptr) {
                        if (mBytes == 2) {
                            _maskQK<FLOAT16_T>((float*)qkPacked, &mScale, seqLen, subKvSeqLen, mPack, kvSeqLen,
                                               i * mBlockKV, padSeqLength, sinksPtr, mask,
                                               (mKeyQuantMode == KVQuantMode::Int8), isLowerTriangular);
                        } else {
                            _maskQK<float>((float*)qkPacked, &mScale, seqLen, subKvSeqLen, mPack, kvSeqLen,
                                           i * mBlockKV, padSeqLength, sinksPtr, mask,
                                           (mKeyQuantMode == KVQuantMode::Int8), isLowerTriangular);
                        }
                    }
                    gcore->MNNSoftmax(qkSoftmax, (float*)qkPacked, runningMax, runningSum, diffScale, seqLen,
                                      subKvSeqLen, i * mBlockKV, kvValidOffset, mPack, useMaskInSoftmax);
                }
                // 3. qk @ v
                auto qkStride0 = ROUND_UP(subKvSeqLen, lP) * eP * mBytes;
                auto rowStart =
                    (!isLowerTriangular || i * mBlockKV < kvValidOffset) ? 0 : (i * mBlockKV - kvValidOffset);

                if (mValueQuantMode == KVQuantMode::TQ3) {
                    // Vec_dot Value fusion: accumulate in rotated domain, WHT_inverse once
                    // qkSoftmax format: [kvSeq/mPack, seqLen, mPack], element (s,q) at (s/mPack)*seqLen*mPack + q*mPack
                    // + s%mPack
                    int tq3BytesPerSeq = (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
                    int numBlocks = mHeadDim / TQ3_BLOCK_SIZE;
                    auto tq3Buf = mTQ3DequantBuf->host<int8_t>() + tId * mTQ3DequantBuf->stride(0);
                    auto vAccRotated = (float*)(tq3Buf + seqLen * mHeadDim * sizeof(float));

                    auto weightsPtr = (float*)(tq3Buf + seqLen * mHeadDim * sizeof(float) + mHeadDim * sizeof(float));
                    for (int q = rowStart; q < seqLen; q++) {
                        // Extract softmax weights for this query position (float)
                        float* weights = weightsPtr;
                        for (int s = 0; s < subKvSeqLen; s++) {
                            int packIdx = (s / mPack) * seqLen * mPack + q * mPack + s % mPack;
                            if (mBytes == 2) {
                                weights[s] = (float)((FLOAT16_T*)qkSoftmax)[packIdx];
                            } else {
                                weights[s] = ((float*)qkSoftmax)[packIdx];
                            }
                        }

                        // For each dim block: accumulate weighted codebook values in rotated domain
                        for (int b = 0; b < numBlocks; b++) {
                            memset(vAccRotated, 0, TQ3_BLOCK_SIZE * sizeof(float));
                            for (int s = 0; s < subKvSeqLen; s++) {
                                int seqIdx = i * mBlockKV + s;
                                const uint8_t* block =
                                    (uint8_t*)valueAddr + seqIdx * tq3BytesPerSeq + b * TQ3_BYTES_PER_BLOCK;
                                uint16_t scaleFp16;
                                memcpy(&scaleFp16, block, 2);
                                float w = weights[s] * tq3_fp16_to_float(scaleFp16);
                                tq3_weighted_acc_block(vAccRotated, w, block + 2);
                            }
                            // WHT_inverse to get final output values
                            float reconstructed[TQ3_BLOCK_SIZE];
                            tq3_wht_inverse_32(reconstructed, vAccRotated);

                            // Write to qkvPacked: [headDim/mPack, seqLen, mPack]
                            for (int d = 0; d < TQ3_BLOCK_SIZE; d++) {
                                int dimIdx = b * TQ3_BLOCK_SIZE + d;
                                int outIdx = (dimIdx / mPack) * seqLen * mPack + q * mPack + dimIdx % mPack;
                                if (mBytes == 2) {
                                    ((FLOAT16_T*)qkvPacked)[outIdx] = (FLOAT16_T)reconstructed[d];
                                } else {
                                    ((float*)qkvPacked)[outIdx] = reconstructed[d];
                                }
                            }
                        }
                    }
                } else if (mValueQuantMode == KVQuantMode::TQ4) {
                    // Vec_dot Value fusion for TQ4: same structure as TQ3
                    int tq4BytesPerSeq = (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
                    int numBlocks = mHeadDim / TQ4_BLOCK_SIZE;
                    auto tqBuf = mTQ3DequantBuf->host<int8_t>() + tId * mTQ3DequantBuf->stride(0);
                    auto vAccRotated = (float*)(tqBuf + seqLen * mHeadDim * sizeof(float));
                    auto weightsPtr = (float*)(tqBuf + seqLen * mHeadDim * sizeof(float) + mHeadDim * sizeof(float));

                    for (int q = rowStart; q < seqLen; q++) {
                        float* weights = weightsPtr;
                        for (int s = 0; s < subKvSeqLen; s++) {
                            int packIdx = (s / mPack) * seqLen * mPack + q * mPack + s % mPack;
                            if (mBytes == 2) {
                                weights[s] = (float)((FLOAT16_T*)qkSoftmax)[packIdx];
                            } else {
                                weights[s] = ((float*)qkSoftmax)[packIdx];
                            }
                        }
                        for (int b = 0; b < numBlocks; b++) {
                            memset(vAccRotated, 0, TQ4_BLOCK_SIZE * sizeof(float));
                            for (int s = 0; s < subKvSeqLen; s++) {
                                int seqIdx = i * mBlockKV + s;
                                const uint8_t* block =
                                    (uint8_t*)valueAddr + seqIdx * tq4BytesPerSeq + b * TQ4_BYTES_PER_BLOCK;
                                uint16_t scaleFp16;
                                memcpy(&scaleFp16, block, 2);
                                float w = weights[s] * tq3_fp16_to_float(scaleFp16);
                                tq4_weighted_acc_block(vAccRotated, w, block + 2);
                            }
                            float reconstructed[TQ4_BLOCK_SIZE];
                            tq3_wht_inverse_32(reconstructed, vAccRotated);
                            for (int d = 0; d < TQ4_BLOCK_SIZE; d++) {
                                int dimIdx = b * TQ4_BLOCK_SIZE + d;
                                int outIdx = (dimIdx / mPack) * seqLen * mPack + q * mPack + dimIdx % mPack;
                                if (mBytes == 2) {
                                    ((FLOAT16_T*)qkvPacked)[outIdx] = (FLOAT16_T)reconstructed[d];
                                } else {
                                    ((float*)qkvPacked)[outIdx] = reconstructed[d];
                                }
                            }
                        }
                    }
                } else if (mValueQuantMode != KVQuantMode::Int8) {
                    auto valuePtr = valueAddr + i * vstride0 * mBytes;
                    size_t shapeParameters[7] = {(size_t)eP * lP * mBytes,
                                                 ROUND_UP((size_t)subKvSeqLen, lP),
                                                 (size_t)mHeadDim,
                                                 (size_t)dstStep,
                                                 0,
                                                 0,
                                                 0};
                    size_t bExtraStride =
                        (i < kvBlocks - 1)
                            ? 0
                            : (ROUND_UP(mKVCacheManager->getFlashAttentionBlockKv(), lP) - ROUND_UP(subKvSeqLen, lP)) *
                                  hP * mBytes;
                    shapeParameters[5] = bExtraStride;

                    int loop_e = (seqLen - rowStart) / eP;
                    int remain = (seqLen - rowStart) % eP;

                    int ei = 0;
                    elFloatV[0] = eP;
                    elFloatV[1] = ROUND_UP(subKvSeqLen, lP);
                    infoFloatV[2] = eP;
                    for (; ei < loop_e; ei++) {
                        srcPtr[0] = (float const*)((int8_t*)qkSoftmax + (ei * eP + rowStart) * mPack * mBytes);
                        gcore->MNNPackC4ForMatMul_A((float*)qkReordered, srcPtr, infoFloatV, elFloatV);
                        gcore->MNNPackedMatMul((float*)(qkvPacked + (ei * eP + rowStart) * mPack * mBytes),
                                               (float*)qkReordered, (float*)valuePtr, shapeParameters, nullptr, nullptr,
                                               nullptr, nullptr);
                    }
                    if (remain > 0) {
                        elFloatV[0] = remain;
                        infoFloatV[2] = remain;
                        srcPtr[0] = (float const*)((int8_t*)qkSoftmax + (loop_e * eP + rowStart) * mPack * mBytes);
                        shapeParameters[0] = remain * lP * mBytes;
                        gcore->MNNPackC4ForMatMul_A((float*)qkReordered, srcPtr, infoFloatV, elFloatV);
                        gcore->MNNPackedMatMulRemain((float*)(qkvPacked + (loop_e * eP + rowStart) * mPack * mBytes),
                                                     (float*)qkReordered, (float*)valuePtr, remain, shapeParameters,
                                                     nullptr, nullptr, nullptr, nullptr);
                    }
                } else { // use int8 kernel to compute qk@ v
                    auto valuePtr = valueAddr + i * vstride0;
                    auto eRemain = seqLen - rowStart;
                    auto qkPtr =
                        (int8_t*)(qkSoftmax) + rowStart * mPack * mBytes; // [UP_DIV(subKvSeqLen,pack),seqLen,pack]
                    auto qkvFloat = qkvPacked + rowStart * mPack * mBytes;
                    gemmParam4QKxV.weightKernelSum = valueSum + i * ROUND_UP(mHeadDim, hP8);
                    sumParams4QKxV.valid = subKvSeqLen % lP8;
                    sumParams4QKxV.LU = UP_DIV(subKvSeqLen, lP8);

                    auto dstInt8Ptr =
                        (int8_t*)mQuantQK.ptr() + tId * eP8 * ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, mPack);
                    srcPtr[0] = (const float*)(dstInt8Ptr);

                    while (eRemain > 0) {
                        auto eSize = ALIMIN(eRemain, eP8);

                        memset(dstInt8Ptr, 0, eP8 * ROUND_UP(MNN_FLASH_ATTENTION_BLOCK_SIZE, mPack));

                        infoInt8V[1] = eSize; // eReal
                        infoInt8V[4] = eSize; // e to process
                        elInt8V[0] = eSize;   // e to process

                        for (int qi = 0; qi < UP_DIV(subKvSeqLen, mPack); ++qi) {
                            mQuantFunc((float*)(qkPtr + qi * seqLen * mPack * mBytes), dstInt8Ptr + qi * eSize * mPack,
                                       eSize, vQuantScale, -128, 127, vQuantBias, 0);
                        }
                        core->MNNPackC4Int8ForMatMul_A(qkReordered, (int8_t const**)srcPtr, infoInt8V, elInt8V);
                        // mSumQK
                        gcore->MNNSumByAxisLForMatmul_A(gemmParam4QKxV.srcKernelSum, qkReordered,
                                                        (float*)mQKScale.ptr(), eSize, sumParams4QKxV);
                        mInt8GemmKernel(qkvFloat, qkReordered, valuePtr, UP_DIV(MNN_FLASH_ATTENTION_BLOCK_SIZE, lP8),
                                        dstStep, UP_DIV(mHeadDim, mPack), &gemmParam4QKxV, eSize);

                        eRemain -= eSize;
                        qkPtr += (eSize * mPack * mBytes);
                        qkvFloat += (eSize * mPack * mBytes);
                    }
                }

                // 4. flash attention, update each sub kvSeq's final results
                if (runningMax != nullptr && runningSum != nullptr && diffScale != nullptr) {
                    gcore->MNNFlashAttentionUpdateBlockOutput((float*)outputPacked, (float*)qkvPacked, diffScale,
                                                              runningSum, UP_DIV(mHeadDim, mPack), seqLen, mPack, i,
                                                              kvBlocks, mPackQKV->stride(0) / mBytes, mBytes, rowStart);
                }
            }

            // Final results writing: [head_dim/mPack, seq_len, mPack] -> [seq_len, num_head, head_dim]
            if (!outputC4) {
                auto dstPtr = outputs[0]->host<int8_t>() + h * mHeadDim * mBytes;
                // offset = {seqLen, mNumHead * mHeadDim};
                gcore->MNNUnpackCUnitTranspose((float*)dstPtr, (float*)outputPacked, seqLen, mHeadDim, offset);
            }
        }
    };

    MNN_CONCURRENCY_BEGIN(tId, mThreadNum) {
        mCompute((int)tId);
    }
    MNN_CONCURRENCY_END();

    backend()->onReleaseBuffer(unpackQK.get(), Backend::STATIC);
    backend()->onReleaseBuffer(softmMaxQ.get(), Backend::STATIC);
    backend()->onReleaseBuffer(newPackQK.get(), Backend::STATIC);
    backend()->onReleaseBuffer(mTempQKBlock.get(), Backend::STATIC);

    if (!mKVCache) {
        mKVCacheManager->onClear();
    }
    if (!outputC4) {
        auto ptr = outputs[0]->host<float>();
        if (seqLen < outputs[0]->length(1)) {
            ::memset(outputs[0]->host<uint8_t>() + seqLen * mHeadDim * mNumHead * mBytes, 0,
                     (outputs[0]->length(1) - seqLen) * mHeadDim * mNumHead * mBytes);
        }
    }
    return NO_ERROR;
}

bool CPUAttention::onClone(Backend* bn, const Op* op, Execution** dst) {
    if (nullptr == dst) {
        return true;
    }
    auto tmp = new CPUAttention(bn, mKVCache);
    // Share KV cache when cloning within the same session (same meta pointer)
    if (bn->getMetaPtr() == mMeta) {
        tmp->mKVCacheManager = mKVCacheManager;
        // Mark as KV-shared if the target op requests KV reuse
        auto param = op->main_as_AttentionParam();
        if (param && param->kv_shared_layer_index() >= 0) {
            tmp->mIsKVShared = true;
        }
    }
    *dst = tmp;
    return true;
}

CPUAttention::CPUAttention(Backend* backend, bool kv_cache) : Execution(backend), mKVCache(kv_cache) {
    mMeta = (KVMeta*)(backend->getMetaPtr());
    mPackQ.reset(Tensor::createDevice<float>({1, 1, 1, 1}));
    mPackQKV.reset(Tensor::createDevice<float>({1, 1, 1, 1}));
    MNN::KVCacheManager::KVCacheConfig kvconfig;

    // attentionOption % 8:
    // 0: Do not quantize
    // 1: Q,K: Int8, V: Float32
    // 2: Q,K,V: Int8
    // 3: K: TQ3, V: Float32
    // 4: K,V: TQ3
    // 5: K: TQ4, V: Float32
    // 6: K,V: TQ4

    // attentionOption / 8:
    // 0: do not use flash attention
    // 1: use flash attention
    kvconfig.mKVCacheDir = static_cast<CPUBackend*>(backend)->getRuntime()->hint().kvcacheDirPath;
    kvconfig.mPrefixCacheDir = static_cast<CPUBackend*>(backend)->getRuntime()->hint().prefixcacheDirPath;
    kvconfig.mExpandChunk = 64;
    kvconfig.mBlockNum = 1;
    mKVCacheManager.reset(new CPUKVCacheManager(backend, kvconfig));
}

class CPUAttentionCreator : public CPUBackend::Creator {
public:
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                                const MNN::Op* op, Backend* backend) const override {
        auto param = op->main_as_AttentionParam();
        return new CPUAttention(backend, param->kv_cache());
    }
};

REGISTER_CPU_OP_CREATOR_TRANSFORMER(CPUAttentionCreator, OpType_Attention);

} // namespace MNN

#endif // MNN_SUPPORT_TRANSFORMER_FUSE