//
//  CPULinearAttention.cpp
//  MNN
//
//  Created by MNN on 2026/02/10.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#ifdef MNN_SUPPORT_TRANSFORMER_FUSE

#include <limits>
#include <cmath>
#include <vector>
#include <algorithm>
#include "CPULinearAttention.hpp"
#include "CPUBackend.hpp"
#include "core/MNNFileUtils.h"
#include "compute/CommonOptFunction.h"
#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"

namespace MNN {

// ─── Byte-aware element access helpers ───
static inline float _readElement(const int8_t* ptr, int index, int bytes) {
#ifdef __aarch64__
    if (bytes == 2) return (float)((const __fp16*)ptr)[index];
#endif
    return ((const float*)ptr)[index];
}

static inline void _writeElement(int8_t* ptr, int index, float val, int bytes) {
#ifdef __aarch64__
    if (bytes == 2) { ((__fp16*)ptr)[index] = (__fp16)val; return; }
#endif
    ((float*)ptr)[index] = val;
}

// Snapshot the post-prefix recurrent state (lazy-allocated) for eraseHistory
// rollback. Allocation failure leaves mSnapshotValid=false; rollback then
// falls back to zeroing.
static void snapshotPrefixState(StateCache* cache, Backend* backend) {
    if (cache == nullptr || cache->mConvState.get() == nullptr) {
        return;
    }
    int convStateBytes = cache->mConvState->elementSize();
    if (cache->mConvStateSnapshot.get() == nullptr) {
        cache->mConvStateSnapshot.reset(Tensor::createDevice<int8_t>({convStateBytes}));
        if (!backend->onAcquireBuffer(cache->mConvStateSnapshot.get(), Backend::STATIC)) {
            cache->mConvStateSnapshot.reset();
            return;
        }
    }
    ::memcpy(cache->mConvStateSnapshot->host<int8_t>(), cache->mConvState->host<int8_t>(), convStateBytes);
    if (cache->mRecurrentState.get() != nullptr) {
        int rnnBytes = cache->mRecurrentState->elementSize();
        if (cache->mRecurrentStateSnapshot.get() == nullptr) {
            cache->mRecurrentStateSnapshot.reset(Tensor::createDevice<int8_t>({rnnBytes}));
            if (!backend->onAcquireBuffer(cache->mRecurrentStateSnapshot.get(), Backend::STATIC)) {
                cache->mRecurrentStateSnapshot.reset();
                return;
            }
        }
        ::memcpy(cache->mRecurrentStateSnapshot->host<int8_t>(), cache->mRecurrentState->host<int8_t>(), rnnBytes);
    }
    cache->mSnapshotValid = true;
}

ErrorCode CPULinearAttention::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto qkv = inputs[0];
    auto convWeight = inputs[3];

    int batch      = qkv->length(0);
    int convDim    = qkv->length(1);    // D (total projection dim)
    int seqLen     = qkv->length(2);    // L
    int kernelSize = convWeight->length(2);
    int convStateSize = kernelSize - 1;

    // ─── Per-type parameters ───
    int convChannels = convDim;
    bool needRecurrentState = false;

    if (mAttentionType == "short_conv") {
        convChannels = mHeadVDim;
    } else if (mAttentionType == "gated_delta_rule") {
        needRecurrentState = true;
    }

    // ─── Persistent state buffers (STATIC): allocate once, shared via onClone ───
    if (mStateCache->mConvState.get() == nullptr) {
        mStateCache->mConvState.reset(Tensor::createDevice<int8_t>({batch * convChannels * convStateSize * mBytes}));
        bool success = backend()->onAcquireBuffer(mStateCache->mConvState.get(), Backend::STATIC);
        if (!success) return OUT_OF_MEMORY;
        ::memset(mStateCache->mConvState->host<int8_t>(), 0, batch * convChannels * convStateSize * mBytes);

        if (needRecurrentState) {
            int H = mNumVHeads, dk = mHeadKDim, dv = mHeadVDim;
            mStateCache->mRecurrentState.reset(Tensor::createDevice<int8_t>({batch * H * dk * dv * mBytes}));
            success = backend()->onAcquireBuffer(mStateCache->mRecurrentState.get(), Backend::STATIC);
            if (!success) return OUT_OF_MEMORY;
            ::memset(mStateCache->mRecurrentState->host<int8_t>(), 0, batch * H * dk * dv * mBytes);
        }
    } else if (seqLen > 1) {
        // Prefill: decide keep/restore/reset from meta. LA state isn't
        // token-indexed, so eraseHistory triggers a snapshot restore rather
        // than a truncation.
        bool loadingFromDisk = (mMeta != nullptr && mMeta->file_flag == KVMeta::PendingRead && mMeta->file_name.size() > 0);
        bool isExplicitRollback = (mMeta != nullptr && mMeta->remove > 0);
        bool isFreshPrefill = (mMeta == nullptr || mMeta->previous == 0);
        int convStateBytes = batch * convChannels * convStateSize * mBytes;
        int rnnBytes = 0;
        if (mStateCache->mRecurrentState.get() != nullptr) {
            int H = mNumVHeads, dk = mHeadKDim, dv = mHeadVDim;
            rnnBytes = batch * H * dk * dv * mBytes;
        }
        if (loadingFromDisk) {
            // onExecute will mmap-load the prefix state and snapshot it.
        } else if (isExplicitRollback) {
            // eraseHistory(): roll back to the saved post-prefix snapshot. If no
            // snapshot exists (rollback before any prefix prefill ran), zero out.
            if (mStateCache->mSnapshotValid && mStateCache->mConvStateSnapshot.get() != nullptr) {
                ::memcpy(mStateCache->mConvState->host<int8_t>(), mStateCache->mConvStateSnapshot->host<int8_t>(),
                         convStateBytes);
                if (mStateCache->mRecurrentState.get() != nullptr &&
                    mStateCache->mRecurrentStateSnapshot.get() != nullptr) {
                    ::memcpy(mStateCache->mRecurrentState->host<int8_t>(),
                             mStateCache->mRecurrentStateSnapshot->host<int8_t>(), rnnBytes);
                }
            } else {
                ::memset(mStateCache->mConvState->host<int8_t>(), 0, convStateBytes);
                if (mStateCache->mRecurrentState.get() != nullptr) {
                    ::memset(mStateCache->mRecurrentState->host<int8_t>(), 0, rnnBytes);
                }
            }
        } else if (isFreshPrefill) {
            // Fresh sequence: zero the state and drop any stale snapshot.
            ::memset(mStateCache->mConvState->host<int8_t>(), 0, convStateBytes);
            if (mStateCache->mRecurrentState.get() != nullptr) {
                ::memset(mStateCache->mRecurrentState->host<int8_t>(), 0, rnnBytes);
            }
            mStateCache->mSnapshotValid = false;
        }
        // Else (mMeta->previous > 0 && mMeta->remove == 0): reuse_kv continuation.
        // Keep the live state so the new prefill extends from it.
    }

    // ─── Temporary buffers (DYNAMIC) ───
    int totalLen = convStateSize + seqLen;
    mConvPadded.reset(Tensor::createDevice<int8_t>({batch * convChannels * totalLen * mBytes}));
    bool success = backend()->onAcquireBuffer(mConvPadded.get(), Backend::DYNAMIC);
    if (!success) return OUT_OF_MEMORY;

    mConvOut.reset(Tensor::createDevice<int8_t>({batch * convChannels * seqLen * mBytes}));
    success = backend()->onAcquireBuffer(mConvOut.get(), Backend::DYNAMIC);
    if (!success) return OUT_OF_MEMORY;

    if (needRecurrentState) {
        int dk = mHeadKDim, dv = mHeadVDim;
        int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();
        // Per-thread scratch holds q_local + k_local + v_local + vPred + delta.
        // Prefill uses MNNFusedGatedDelta (only needs first 2*dk+dv) but decode
        // falls back to the legacy two-call path (MNNDualMatVec + scalar
        // correction + MNNDecayRankOneUpdate) which needs the full 2*dk+3*dv:
        // the fused kernel regressed FP32 decode by ~3.5% on small L=1 shapes.
        int perThread = 2 * dk + 3 * dv;
        mThreadLocalBuf.reset(Tensor::createDevice<int8_t>({threadNum * perThread * mBytes}));
        success = backend()->onAcquireBuffer(mThreadLocalBuf.get(), Backend::DYNAMIC);
        if (!success) return OUT_OF_MEMORY;

        // Pre-computed decay buffer: exp(gate) for all [B, L, H]
        // Always fp32 — MNNExp requires fp32, decay is a scalar per timestep
        // Use int8_t with explicit byte count to avoid Arm82 backend halving the allocation
        mDecayBuf.reset(Tensor::createDevice<int8_t>({batch * seqLen * mNumVHeads * (int)sizeof(float)}));
        success = backend()->onAcquireBuffer(mDecayBuf.get(), Backend::DYNAMIC);
        if (!success) return OUT_OF_MEMORY;

        backend()->onReleaseBuffer(mDecayBuf.get(), Backend::DYNAMIC);
        backend()->onReleaseBuffer(mThreadLocalBuf.get(), Backend::DYNAMIC);
    }

    // fp16 path: per-thread fp32 temp buffer for Conv1D + SiLu (MNNSiLu requires fp32)
    if (mBytes == 2) {
        int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();
        // Need totalLen floats for padded input + L floats for SiLu output = (totalLen + seqLen) per thread
        mConvFp32Buf.reset(Tensor::createDevice<int8_t>({threadNum * (totalLen + seqLen) * (int)sizeof(float)}));
        success = backend()->onAcquireBuffer(mConvFp32Buf.get(), Backend::DYNAMIC);
        if (!success) return OUT_OF_MEMORY;
        backend()->onReleaseBuffer(mConvFp32Buf.get(), Backend::DYNAMIC);
    }

    backend()->onReleaseBuffer(mConvPadded.get(), Backend::DYNAMIC);
    backend()->onReleaseBuffer(mConvOut.get(), Backend::DYNAMIC);

    return NO_ERROR;
}

void CPULinearAttention::gated_delta_rule_ref(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    // Reference implementation (fp32 only, for correctness verification)
    auto qkvTensor    = inputs[0];
    auto gateTensor   = inputs[1];
    auto betaTensor   = inputs[2];
    auto convWTensor  = inputs[3];
    auto outTensor    = outputs[0];

    const float* qkvPtr   = qkvTensor->host<float>();
    const float* gatePtr  = gateTensor->host<float>();
    const float* betaPtr  = betaTensor->host<float>();
    const float* convWPtr = convWTensor->host<float>();
    float* outPtr         = outTensor->host<float>();

    const int B       = qkvTensor->length(0);
    const int D       = qkvTensor->length(1);
    const int L       = qkvTensor->length(2);
    const int H_k     = mNumKHeads;
    const int H_v     = mNumVHeads;
    const int d_k     = mHeadKDim;
    const int d_v     = mHeadVDim;
    const int key_dim = H_k * d_k;
    const int val_dim = H_v * d_v;
    const int K_conv  = convWTensor->length(2);
    const int convStateSize = K_conv - 1;
    const bool useL2Norm = mUseQKL2Norm;
    const int gqa_factor = (H_v > H_k) ? (H_v / H_k) : 1;
    const int H = H_v;

    // Step 1: Depthwise Conv1D + SiLU
    const int totalLen = convStateSize + L;
    std::vector<float> convInput(B * D * totalLen, 0.0f);
    float* convStatePtr = mStateCache->mConvState->host<float>();

    for (int b = 0; b < B; ++b) {
        for (int d = 0; d < D; ++d) {
            float* dst = convInput.data() + b * D * totalLen + d * totalLen;
            const float* stateChannel = convStatePtr + b * D * convStateSize + d * convStateSize;
            ::memcpy(dst, stateChannel, convStateSize * sizeof(float));
            const float* inputChannel = qkvPtr + b * D * L + d * L;
            ::memcpy(dst + convStateSize, inputChannel, L * sizeof(float));
        }
    }

    std::vector<float> convOut(B * D * L, 0.0f);
    for (int b = 0; b < B; ++b) {
        for (int d = 0; d < D; ++d) {
            const float* src    = convInput.data() + b * D * totalLen + d * totalLen;
            const float* weight = convWPtr + d * K_conv;
            float* out          = convOut.data() + b * D * L + d * L;
            for (int l = 0; l < L; ++l) {
                float sum = 0.0f;
                for (int k = 0; k < K_conv; ++k) {
                    sum += src[l + k] * weight[k];
                }
                float sigmoid_val = 1.0f / (1.0f + expf(-sum));
                out[l] = sum * sigmoid_val;
            }
        }
    }

    for (int b = 0; b < B; ++b) {
        for (int d = 0; d < D; ++d) {
            const float* src = convInput.data() + b * D * totalLen + d * totalLen + (totalLen - convStateSize);
            float* dst       = convStatePtr + b * D * convStateSize + d * convStateSize;
            ::memcpy(dst, src, convStateSize * sizeof(float));
        }
    }

    // Step 2: Split Q, K, V
    std::vector<float> Q(B * L * H * d_k, 0.0f);
    std::vector<float> K(B * L * H * d_k, 0.0f);
    std::vector<float> V(B * L * H * d_v, 0.0f);

    for (int b = 0; b < B; ++b) {
        for (int l = 0; l < L; ++l) {
            for (int h = 0; h < H_k; ++h) {
                for (int dk = 0; dk < d_k; ++dk) {
                    int srcChannel = h * d_k + dk;
                    float val = convOut[b * D * L + srcChannel * L + l];
                    for (int r = 0; r < gqa_factor; ++r) {
                        int dstHead = h * gqa_factor + r;
                        Q[(b * L + l) * H * d_k + dstHead * d_k + dk] = val;
                    }
                }
            }
            for (int h = 0; h < H_k; ++h) {
                for (int dk = 0; dk < d_k; ++dk) {
                    int srcChannel = key_dim + h * d_k + dk;
                    float val = convOut[b * D * L + srcChannel * L + l];
                    for (int r = 0; r < gqa_factor; ++r) {
                        int dstHead = h * gqa_factor + r;
                        K[(b * L + l) * H * d_k + dstHead * d_k + dk] = val;
                    }
                }
            }
            for (int h = 0; h < H_v; ++h) {
                for (int dv = 0; dv < d_v; ++dv) {
                    int srcChannel = 2 * key_dim + h * d_v + dv;
                    float val = convOut[b * D * L + srcChannel * L + l];
                    V[(b * L + l) * H * d_v + h * d_v + dv] = val;
                }
            }
        }
    }

    // Step 3: Optional L2 Normalization
    if (useL2Norm) {
        const float eps = 1e-6f;
        for (int i = 0; i < B * L * H; ++i) {
            float* qHead = Q.data() + i * d_k;
            float sumSq = 0.0f;
            for (int dk = 0; dk < d_k; ++dk) sumSq += qHead[dk] * qHead[dk];
            float invNorm = 1.0f / sqrtf(sumSq + eps);
            for (int dk = 0; dk < d_k; ++dk) qHead[dk] *= invNorm;

            float* kHead = K.data() + i * d_k;
            sumSq = 0.0f;
            for (int dk = 0; dk < d_k; ++dk) sumSq += kHead[dk] * kHead[dk];
            invNorm = 1.0f / sqrtf(sumSq + eps);
            for (int dk = 0; dk < d_k; ++dk) kHead[dk] *= invNorm;
        }
    }

    // Step 4: Scale Q
    const float qScale = 1.0f / sqrtf((float)d_k);
    for (int i = 0; i < B * L * H * d_k; ++i) Q[i] *= qScale;

    // Step 5: Gated Delta Rule
    float* rnnStatePtr = mStateCache->mRecurrentState->host<float>();
    for (int b = 0; b < B; ++b) {
        for (int t = 0; t < L; ++t) {
            for (int h = 0; h < H; ++h) {
                float* state = rnnStatePtr + (b * H + h) * d_k * d_v;
                const float* q_t = Q.data() + (b * L + t) * H * d_k + h * d_k;
                const float* k_t = K.data() + (b * L + t) * H * d_k + h * d_k;
                const float* v_t = V.data() + (b * L + t) * H * d_v + h * d_v;
                float g_t    = gatePtr[b * L * H + t * H + h];
                float beta_t = betaPtr[b * L * H + t * H + h];

                float decay = expf(g_t);
                for (int i = 0; i < d_k * d_v; ++i) state[i] *= decay;

                std::vector<float> v_pred(d_v, 0.0f);
                for (int dk = 0; dk < d_k; ++dk)
                    for (int dv = 0; dv < d_v; ++dv)
                        v_pred[dv] += state[dk * d_v + dv] * k_t[dk];

                std::vector<float> delta(d_v);
                for (int dv = 0; dv < d_v; ++dv)
                    delta[dv] = beta_t * (v_t[dv] - v_pred[dv]);

                for (int dk = 0; dk < d_k; ++dk)
                    for (int dv = 0; dv < d_v; ++dv)
                        state[dk * d_v + dv] += k_t[dk] * delta[dv];

                float* o_t = outPtr + (b * L + t) * H * d_v + h * d_v;
                for (int dv = 0; dv < d_v; ++dv) {
                    float sum = 0.0f;
                    for (int dk = 0; dk < d_k; ++dk)
                        sum += state[dk * d_v + dv] * q_t[dk];
                    o_t[dv] = sum;
                }
            }
        }
    }
}

ErrorCode CPULinearAttention::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    // onResize() may be skipped when shapes are unchanged. Ensure state is reset here too.
    int seqLen = inputs[0]->length(2);
    if (seqLen > 1 && mMeta != nullptr && mMeta->previous == mMeta->remove) {
        bool loadingFromDisk = (mMeta->file_flag == KVMeta::PendingRead && mMeta->file_name.size() > 0);
        if (!loadingFromDisk) {
            if (mStateCache->mConvState.get() != nullptr) {
                ::memset(mStateCache->mConvState->host<int8_t>(), 0, mStateCache->mConvState->elementSize());
            }
            if (mStateCache->mRecurrentState.get() != nullptr) {
                ::memset(mStateCache->mRecurrentState->host<int8_t>(), 0, mStateCache->mRecurrentState->elementSize());
            }
        }
    }

    // Capture layer_index once per prefix-cache session (chunk 1, marked by
    // previous == remove); chunks 2..N reuse it. Mirrors CPUKVCacheManager,
    // which only advances layer_index in onAlloc (chunk 1) and not in onRealloc
    // (chunks 2..N) — advancing on every chunk would drift LA past FA layers
    // and clobber FA's prefix files (SIGBUS in hybrid models).
    if (mMeta != nullptr && mMeta->file_name.size() > 0 &&
        (mMeta->file_flag == KVMeta::PendingWrite || mMeta->file_flag == KVMeta::PendingRead) &&
        mMeta->previous == mMeta->remove) {
        mStateCache->mPrefixLayerIndex = mMeta->layer_index;
        mMeta->layer_index = (mMeta->layer_index + 1) % mMeta->layer_nums;
    }

    // Load prefix cache from disk (PendingRead)
    if (mMeta != nullptr && mMeta->file_name.size() > 0 && mMeta->file_flag == KVMeta::PendingRead) {
        // Sentinel guard: capture above only fires on previous == remove.
        // On other paths (e.g. partial eraseHistory) the index stays -1; using
        // it would write/read "_-1.k" and corrupt the cache dir, so skip with
        // a diagnostic.
        if (mStateCache->mPrefixLayerIndex < 0) {
            MNN_ERROR(
                "CPULinearAttention: PendingRead skipped, no prefix-layer-index captured "
                "for this session (previous=%zu remove=%zu — capture predicate requires "
                "previous == remove)\n",
                mMeta->previous, mMeta->remove);
        } else {
            int layer_index = mStateCache->mPrefixLayerIndex;
            std::string basePath =
                MNNFilePathConcat(mPrefixCacheDir, mMeta->file_name) + "_" + std::to_string(layer_index);
            std::string pathk = basePath + ".k";
            std::string pathv = basePath + ".v";
            // Load conv state (.k file)
            auto kfd = MNNOpenFile(pathk.c_str(), MNN_FILE_READ);
            if (kfd != INVALID_FILE) {
                size_t kSize = MNNGetFileSize(kfd);
                if (kSize > 0 && kSize != INVALID_SIZE) {
                    void* kMap = MNNMmapFile(kfd, kSize, true);
                    if (kMap != nullptr) {
                        ::memcpy(mStateCache->mConvState->host<int8_t>(), kMap, kSize);
                        MNNUnmapFile(kMap, kSize);
                    }
                }
                MNNCloseFile(kfd);
            } else {
                MNN_PRINT("CPULinearAttention: Failed to open prefix cache file: %s\n", pathk.c_str());
            }
            // Load recurrent state (.v file)
            auto vfd = MNNOpenFile(pathv.c_str(), MNN_FILE_READ);
            if (vfd != INVALID_FILE) {
                size_t vSize = MNNGetFileSize(vfd);
                if (vSize > 0 && vSize != INVALID_SIZE && mStateCache->mRecurrentState.get() != nullptr) {
                    void* vMap = MNNMmapFile(vfd, vSize, true);
                    if (vMap != nullptr) {
                        ::memcpy(mStateCache->mRecurrentState->host<int8_t>(), vMap, vSize);
                        MNNUnmapFile(vMap, vSize);
                    }
                }
                MNNCloseFile(vfd);
            } else {
                MNN_PRINT("CPULinearAttention: Failed to open prefix cache file: %s\n", pathv.c_str());
            }
            // Snapshot the loaded state for in-memory eraseHistory rollback.
            snapshotPrefixState(mStateCache.get(), backend());
        }
    }

    // Normal execution
    if (mAttentionType == "short_conv") {
        short_conv(inputs, outputs);
    } else {
        gated_delta_rule_mnn(inputs, outputs);
    }

    // Save prefix cache to disk (PendingWrite)
    if (mMeta != nullptr && mMeta->file_name.size() > 0 && mMeta->file_flag == KVMeta::PendingWrite) {
        // Sentinel guard: same rationale as PendingRead above.
        if (mStateCache->mPrefixLayerIndex < 0) {
            MNN_ERROR(
                "CPULinearAttention: PendingWrite skipped, no prefix-layer-index captured "
                "for this session (previous=%zu remove=%zu — capture predicate requires "
                "previous == remove)\n",
                mMeta->previous, mMeta->remove);
        } else {
            MNNCreateDir(mPrefixCacheDir.c_str());
            int layer_index = mStateCache->mPrefixLayerIndex;
            std::string basePath =
                MNNFilePathConcat(mPrefixCacheDir, mMeta->file_name) + "_" + std::to_string(layer_index);
            std::string pathk = basePath + ".k";
            std::string pathv = basePath + ".v";
            // Save conv state (.k file)
            size_t convBytes = mStateCache->mConvState->elementSize();
            auto kfd = MNNCreateFile(pathk.c_str());
            if (kfd != INVALID_FILE) {
                MNNSetFileSize(kfd, convBytes);
                void* kMap = MNNMmapFile(kfd, convBytes);
                if (kMap != nullptr) {
                    ::memcpy(kMap, mStateCache->mConvState->host<int8_t>(), convBytes);
                    MNNUnmapFile(kMap, convBytes);
                }
                MNNCloseFile(kfd);
            } else {
                MNN_PRINT("CPULinearAttention: Failed to create prefix cache file: %s\n", pathk.c_str());
            }
            // Save recurrent state (.v file) — may be empty for short_conv
            size_t recurrentBytes =
                (mStateCache->mRecurrentState.get() != nullptr) ? mStateCache->mRecurrentState->elementSize() : 0;
            auto vfd = MNNCreateFile(pathv.c_str());
            if (vfd != INVALID_FILE) {
                if (recurrentBytes > 0) {
                    MNNSetFileSize(vfd, recurrentBytes);
                    void* vMap = MNNMmapFile(vfd, recurrentBytes);
                    if (vMap != nullptr) {
                        ::memcpy(vMap, mStateCache->mRecurrentState->host<int8_t>(), recurrentBytes);
                        MNNUnmapFile(vMap, recurrentBytes);
                    }
                }
                MNNCloseFile(vfd);
            } else {
                MNN_PRINT("CPULinearAttention: Failed to create prefix cache file: %s\n", pathv.c_str());
            }
            // Snapshot the written state for in-memory eraseHistory rollback.
            snapshotPrefixState(mStateCache.get(), backend());
        }
    }

    return NO_ERROR;
}

void CPULinearAttention::gated_delta_rule_mnn(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto qkvTensor    = inputs[0];
    auto gateTensor   = inputs[1];
    auto betaTensor   = inputs[2];
    auto convWTensor  = inputs[3];
    auto outTensor    = outputs[0];

    const int8_t* qkvPtr   = qkvTensor->host<int8_t>();
    const int8_t* gatePtr  = gateTensor->host<int8_t>();
    const int8_t* betaPtr  = betaTensor->host<int8_t>();
    const int8_t* convWPtr = convWTensor->host<int8_t>();
    int8_t* outPtr         = outTensor->host<int8_t>();

    const int B       = qkvTensor->length(0);
    const int D       = qkvTensor->length(1);
    const int L       = qkvTensor->length(2);

    // Decode fast path: L=1, skip decay buffer, stride=1 contiguous access
    if (L == 1) {
        gated_delta_rule_decode(inputs, outputs);
        return;
    }
    const int H_k     = mNumKHeads;
    const int H_v     = mNumVHeads;
    const int d_k     = mHeadKDim;
    const int d_v     = mHeadVDim;
    const int key_dim = H_k * d_k;
    const int val_dim = H_v * d_v;
    const int K_conv  = convWTensor->length(2);
    const int convStateSize = K_conv - 1;
    const bool useL2Norm = mUseQKL2Norm;
    const int gqa_factor = (H_v > H_k) ? (H_v / H_k) : 1;
    const int H = H_v;
    const int bytes = mBytes;

    // Get pre-allocated buffers
    int8_t* convPadded  = mConvPadded->host<int8_t>();
    int8_t* convOut     = mConvOut->host<int8_t>();
    int8_t* convStatePtr = mStateCache->mConvState->host<int8_t>();

    // ─── Step 1: Depthwise Conv1D + SiLU (multi-threaded across B×D channels) ───
    const int totalLen = convStateSize + L;
    const int totalChannels = B * D;
    int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();

    // fp16 path uses per-thread fp32 temp buffers for vectorized Conv1D + SiLu
    float* convFp32Base = (bytes == 2) ? mConvFp32Buf->host<float>() : nullptr;

    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        // Per-thread fp32 buffers (only used for fp16 path)
        float* fp32Padded = (bytes == 2) ? convFp32Base + (int)tId * (totalLen + L) : nullptr;
        float* fp32Out    = (bytes == 2) ? fp32Padded + totalLen : nullptr;

        for (int idx = (int)tId; idx < totalChannels; idx += threadNum) {
            int d = idx % D;

            // 1a. Build padded input: cat(convState, qkv) for this channel
            int8_t* padded = convPadded + idx * totalLen * bytes;
            const int8_t* stateChannel = convStatePtr + idx * convStateSize * bytes;
            ::memcpy(padded, stateChannel, convStateSize * bytes);
            const int8_t* inputChannel = qkvPtr + idx * L * bytes;
            ::memcpy(padded + convStateSize * bytes, inputChannel, L * bytes);

            // 1b. Save conv state first (before we overwrite padded)
            const int8_t* newState = padded + (totalLen - convStateSize) * bytes;
            int8_t* dstState = convStatePtr + idx * convStateSize * bytes;
            ::memcpy(dstState, newState, convStateSize * bytes);

            // 1c. Conv1D + SiLU
            const int8_t* weight = convWPtr + d * K_conv * bytes;
            int8_t* out = convOut + idx * L * bytes;

            if (bytes == 2) {
                // fp16 path: convert to fp32, compute Conv1D + SiLu vectorized, convert back
                auto coreFn = static_cast<CPUBackend*>(backend())->functions();
                coreFn->MNNLowpToFp32((const int16_t*)padded, fp32Padded, totalLen);
                float w0 = fp32Padded[0]; // dummy, will read weights separately
#ifdef __aarch64__
                w0 = (float)((__fp16*)weight)[0];
                float w1 = (float)((__fp16*)weight)[1];
                float w2 = (float)((__fp16*)weight)[2];
                float w3 = (float)((__fp16*)weight)[3];
#else
                float w1 = 0, w2 = 0, w3 = 0;
#endif
                if (K_conv == 4) {
                    int l = 0;
                    for (; l + 3 < L; l += 4) {
                        fp32Padded[l]   = fp32Padded[l]*w0 + fp32Padded[l+1]*w1 + fp32Padded[l+2]*w2 + fp32Padded[l+3]*w3;
                        fp32Padded[l+1] = fp32Padded[l+1]*w0 + fp32Padded[l+2]*w1 + fp32Padded[l+3]*w2 + fp32Padded[l+4]*w3;
                        fp32Padded[l+2] = fp32Padded[l+2]*w0 + fp32Padded[l+3]*w1 + fp32Padded[l+4]*w2 + fp32Padded[l+5]*w3;
                        fp32Padded[l+3] = fp32Padded[l+3]*w0 + fp32Padded[l+4]*w1 + fp32Padded[l+5]*w2 + fp32Padded[l+6]*w3;
                    }
                    for (; l < L; ++l) {
                        fp32Padded[l] = fp32Padded[l]*w0 + fp32Padded[l+1]*w1 + fp32Padded[l+2]*w2 + fp32Padded[l+3]*w3;
                    }
                } else {
                    for (int l = 0; l < L; ++l) {
                        float sum = 0.0f;
                        for (int k = 0; k < K_conv; ++k) {
                            float wk = _readElement(weight, k, bytes);
                            sum += fp32Padded[l + k] * wk;
                        }
                        fp32Padded[l] = sum;
                    }
                }
                MNNSiLu(fp32Out, fp32Padded, L);
                coreFn->MNNFp32ToLowp(fp32Out, (int16_t*)out, L);
            } else {
                // fp32 path: direct compute
                float* fPadded = (float*)padded;
                if (K_conv == 4) {
                    float w0 = ((float*)weight)[0], w1 = ((float*)weight)[1];
                    float w2 = ((float*)weight)[2], w3 = ((float*)weight)[3];
                    int l = 0;
                    for (; l + 3 < L; l += 4) {
                        fPadded[l]   = fPadded[l]*w0 + fPadded[l+1]*w1 + fPadded[l+2]*w2 + fPadded[l+3]*w3;
                        fPadded[l+1] = fPadded[l+1]*w0 + fPadded[l+2]*w1 + fPadded[l+3]*w2 + fPadded[l+4]*w3;
                        fPadded[l+2] = fPadded[l+2]*w0 + fPadded[l+3]*w1 + fPadded[l+4]*w2 + fPadded[l+5]*w3;
                        fPadded[l+3] = fPadded[l+3]*w0 + fPadded[l+4]*w1 + fPadded[l+5]*w2 + fPadded[l+6]*w3;
                    }
                    for (; l < L; ++l) {
                        fPadded[l] = fPadded[l]*w0 + fPadded[l+1]*w1 + fPadded[l+2]*w2 + fPadded[l+3]*w3;
                    }
                } else {
                    for (int l = 0; l < L; ++l) {
                        float sum = 0.0f;
                        for (int k = 0; k < K_conv; ++k) sum += fPadded[l + k] * ((float*)weight)[k];
                        fPadded[l] = sum;
                    }
                }
                MNNSiLu((float*)out, fPadded, L);
            }
        }
    }
    MNN_CONCURRENCY_END();

    // ─── Step 1.5: Batch exp(gate) ───
    // Decay buffer is always fp32. Convert gate to fp32 if needed, then MNNExp.
    float* decayPtr = mDecayBuf->host<float>();
    const int gateTotalSize = B * L * H;
    if (bytes == 4) {
        float expOffset[4] = {1.0f, 0.0f, 0.0f, 0.0f};
        MNNExp(decayPtr, (const float*)gatePtr, expOffset, gateTotalSize);
    } else {
        // fp16: compute exp per-element (gate is small: B*L*H)
        for (int i = 0; i < gateTotalSize; ++i) {
            decayPtr[i] = expf(_readElement(gatePtr, i, bytes));
        }
    }

    // ─── Steps 2-5 fused: Split + L2Norm + Scale + Gated Delta Rule ───
    const float qScale = 1.0f / sqrtf((float)d_k);
    auto gcore = static_cast<CPUBackend*>(backend())->functions();
    int8_t* rnnStatePtr = mStateCache->mRecurrentState->host<int8_t>();

    const int totalHeads = B * H;

    int8_t* threadBufBase = mThreadLocalBuf->host<int8_t>();
    // Prefill uses fused kernel (only first 2*dk+dv touched) but the per-thread
    // stride must match the larger allocation (decode's 2*dk+3*dv).
    const int perThread = 2 * d_k + 3 * d_v;

    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        int8_t* tBuf = threadBufBase + (int)tId * perThread * bytes;
        // Local buffers in native format (fp16 or fp32)
        int8_t* q_local = tBuf;
        int8_t* k_local = tBuf + d_k * bytes;
        int8_t* v_local = tBuf + 2 * d_k * bytes;

        for (int idx = (int)tId; idx < totalHeads; idx += threadNum) {
            int b = idx / H;
            int h = idx % H;
            int k_head = h / gqa_factor;

            int8_t* state = rnnStatePtr + idx * d_k * d_v * bytes;
            const int8_t* convBase = convOut + b * D * L * bytes;

            // Pre-compute base pointers for this head's channels
            const int8_t* qBase = convBase + k_head * d_k * L * bytes;
            const int8_t* kBase = convBase + (key_dim + k_head * d_k) * L * bytes;
            const int8_t* vBase = convBase + (2 * key_dim + h * d_v) * L * bytes;

            for (int t = 0; t < L; ++t) {
                // ── Step 2: Extract q_t, k_t, v_t from convOut (strided access) ──
                for (int i = 0; i < d_k; ++i) {
                    float qv = _readElement(qBase, i * L + t, bytes);
                    float kv = _readElement(kBase, i * L + t, bytes);
                    _writeElement(q_local, i, qv, bytes);
                    _writeElement(k_local, i, kv, bytes);
                }
                for (int i = 0; i < d_v; ++i) {
                    float vv = _readElement(vBase, i * L + t, bytes);
                    _writeElement(v_local, i, vv, bytes);
                }

                // ── Step 3+4: L2 Normalization + Scale (fused) ──
                if (useL2Norm) {
                    const float eps = 1e-6f;
                    float qSumSq = 0.0f, kSumSq = 0.0f;
                    for (int i = 0; i < d_k; ++i) {
                        float qi = _readElement(q_local, i, bytes);
                        float ki = _readElement(k_local, i, bytes);
                        qSumSq += qi * qi;
                        kSumSq += ki * ki;
                    }
                    float qNormScale = qScale / sqrtf(qSumSq + eps);
                    float kInvNorm = 1.0f / sqrtf(kSumSq + eps);
                    for (int i = 0; i < d_k; ++i) {
                        _writeElement(q_local, i, _readElement(q_local, i, bytes) * qNormScale, bytes);
                        _writeElement(k_local, i, _readElement(k_local, i, bytes) * kInvNorm, bytes);
                    }
                } else {
                    for (int i = 0; i < d_k; ++i) {
                        _writeElement(q_local, i, _readElement(q_local, i, bytes) * qScale, bytes);
                    }
                }

                // ── Step 5: Gated Delta Rule recurrence ──
                float decay  = decayPtr[b * L * H + t * H + h];
                float beta_t = _readElement(betaPtr, b * L * H + t * H + h, bytes);

                // dot(k, q) — small reduction in fp32 for precision.
                float kq = 0.0f;
                for (int i = 0; i < d_k; ++i) {
                    kq += _readElement(k_local, i, bytes) * _readElement(q_local, i, bytes);
                }

                // out_t is written; state S is updated in-place.
                int8_t* o_t = outPtr + ((b * L + t) * H * d_v + h * d_v) * bytes;
                gcore->MNNFusedGatedDelta((float*)state, (float*)k_local, (float*)q_local, (float*)v_local, (float*)o_t,
                                          decay, beta_t, kq, d_k, d_v);
            } // end timestep
        } // end head
    }
    MNN_CONCURRENCY_END();
}

void CPULinearAttention::gated_delta_rule_decode(const std::vector<Tensor*>& inputs,
                                                 const std::vector<Tensor*>& outputs) const {
    auto qkvTensor = inputs[0];
    auto gateTensor = inputs[1];
    auto betaTensor = inputs[2];
    auto convWTensor = inputs[3];
    auto outTensor = outputs[0];

    const int8_t* qkvPtr = qkvTensor->host<int8_t>();
    const int8_t* gatePtr = gateTensor->host<int8_t>();
    const int8_t* betaPtr = betaTensor->host<int8_t>();
    const int8_t* convWPtr = convWTensor->host<int8_t>();
    int8_t* outPtr = outTensor->host<int8_t>();

    const int B = qkvTensor->length(0);
    const int D = qkvTensor->length(1);
    // L == 1 guaranteed
    const int H_k = mNumKHeads;
    const int H_v = mNumVHeads;
    const int d_k = mHeadKDim;
    const int d_v = mHeadVDim;
    const int key_dim = H_k * d_k;
    const int K_conv = convWTensor->length(2);
    const int convStateSize = K_conv - 1;
    const bool useL2Norm = mUseQKL2Norm;
    const int gqa_factor = (H_v > H_k) ? (H_v / H_k) : 1;
    const int H = H_v;
    const int bytes = mBytes;

    auto* convOut = mConvOut->host<int8_t>();
    auto* convStatePtr = mStateCache->mConvState->host<int8_t>();

    const int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();

    // ─── Step 1: Conv1D + SiLU (L=1, one output per channel) ───
    // Each channel: dot product of [convState, input_val] with weight, then SiLU
    const int totalChannels = B * D;
    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        for (int idx = (int)tId; idx < totalChannels; idx += threadNum) {
            const int d = idx % D;

            // Read the single input value for this channel
            const float inputVal = _readElement(qkvPtr, idx, bytes);

            // Compute conv: dot(cat(state, input), weight)
            float sum = 0.0f;
            const int8_t* stateChannel = convStatePtr + idx * convStateSize * bytes;
            const int8_t* weight = convWPtr + d * K_conv * bytes;
            for (int k = 0; k < convStateSize; ++k) {
                sum += _readElement(stateChannel, k, bytes) * _readElement(weight, k, bytes);
            }
            sum += inputVal * _readElement(weight, convStateSize, bytes);

            // SiLU activation
            const float sigmoid_val = 1.0f / (1.0f + expf(-sum));
            const float convResult = sum * sigmoid_val;
            _writeElement(convOut, idx, convResult, bytes);

            // Update conv state: shift left by 1, append new input
            for (int k = 0; k < convStateSize - 1; ++k) {
                const float v = _readElement(stateChannel, k + 1, bytes);
                _writeElement(convStatePtr + idx * convStateSize * bytes, k, v, bytes);
            }
            _writeElement(convStatePtr + idx * convStateSize * bytes, convStateSize - 1, inputVal, bytes);
        }
    }
    MNN_CONCURRENCY_END();

    // ─── Steps 2-5 fused: QKV extraction + L2Norm + Scale + Gated Delta Rule ───
    const float qScale = 1.0f / sqrtf((float)d_k);
    const auto gcore = static_cast<CPUBackend*>(backend())->functions();
    auto* rnnStatePtr = mStateCache->mRecurrentState->host<int8_t>();

    const int totalHeads = B * H;
    auto* threadBufBase = mThreadLocalBuf->host<int8_t>();
    // Decode (L=1) keeps the legacy two-call path: the fused kernel regressed
    // FP32 decode by ~3.5% on this shape (small d_v, single timestep).
    const int perThread = 2 * d_k + 3 * d_v;

    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        int8_t* tBuf = threadBufBase + (int)tId * perThread * bytes;
        int8_t* q_local = tBuf;
        int8_t* k_local = tBuf + d_k * bytes;
        int8_t* v_local = tBuf + 2 * d_k * bytes;
        int8_t* localVPred = tBuf + (2 * d_k + d_v) * bytes;
        int8_t* localDelta = tBuf + (2 * d_k + 2 * d_v) * bytes;

        for (int idx = (int)tId; idx < totalHeads; idx += threadNum) {
            const int b = idx / H;
            const int h = idx % H;
            const int k_head = h / gqa_factor;

            int8_t* state = rnnStatePtr + idx * d_k * d_v * bytes;

            // L=1: conv_out is [B, D, 1], stride=1, contiguous read
            const int8_t* convBase = convOut + b * D * bytes;
            const int8_t* qBase = convBase + k_head * d_k * bytes;
            const int8_t* kBase = convBase + (key_dim + k_head * d_k) * bytes;
            const int8_t* vBase = convBase + (2 * key_dim + h * d_v) * bytes;

            // ── Step 2: Extract q, k, v (contiguous copy, stride=1) ──
            ::memcpy(q_local, qBase, d_k * bytes);
            ::memcpy(k_local, kBase, d_k * bytes);
            ::memcpy(v_local, vBase, d_v * bytes);

            // ── Step 3+4: L2 Normalization + Scale (fused) ──
            if (useL2Norm) {
                const float eps = 1e-6f;
                float qSumSq = 0.0f, kSumSq = 0.0f;
                for (int i = 0; i < d_k; ++i) {
                    const float qi = _readElement(q_local, i, bytes);
                    const float ki = _readElement(k_local, i, bytes);
                    qSumSq += qi * qi;
                    kSumSq += ki * ki;
                }
                const float qNormScale = qScale / sqrtf(qSumSq + eps);
                const float kInvNorm = 1.0f / sqrtf(kSumSq + eps);
                for (int i = 0; i < d_k; ++i) {
                    _writeElement(q_local, i, _readElement(q_local, i, bytes) * qNormScale, bytes);
                    _writeElement(k_local, i, _readElement(k_local, i, bytes) * kInvNorm, bytes);
                }
            } else {
                for (int i = 0; i < d_k; ++i) {
                    _writeElement(q_local, i, _readElement(q_local, i, bytes) * qScale, bytes);
                }
            }

            // ── Step 5: Gated Delta Rule (legacy two-call path) ──
            const float decay = expf(_readElement(gatePtr, b * H + h, bytes));
            const float beta_t = _readElement(betaPtr, b * H + h, bytes);

            // Pass 1 (read-only): out_k = S^T @ k → localVPred,
            //                     out_q = S^T @ q → o_t (overwritten by correction below).
            int8_t* o_t = outPtr + (b * H * d_v + h * d_v) * bytes;
            gcore->MNNDualMatVec((float*)state, (float*)k_local, (float*)q_local, (float*)localVPred, (float*)o_t, d_k,
                                 d_v);

            // Analytic correction: delta = beta * (v - decay * vPred);
            //                      out   = decay * out_q + dot(k,q) * delta.
            float kq = 0.0f;
            for (int i = 0; i < d_k; ++i) {
                kq += _readElement(k_local, i, bytes) * _readElement(q_local, i, bytes);
            }
            for (int i = 0; i < d_v; ++i) {
                const float vPred_i = decay * _readElement(localVPred, i, bytes);
                const float v_i = _readElement(v_local, i, bytes);
                const float delta_i = beta_t * (v_i - vPred_i);
                const float out_i = decay * _readElement(o_t, i, bytes) + kq * delta_i;
                _writeElement(localDelta, i, delta_i, bytes);
                _writeElement(o_t, i, out_i, bytes);
            }

            // Pass 2: S = decay * S + k ⊗ delta.
            gcore->MNNDecayRankOneUpdate((float*)state, (float*)k_local, (float*)localDelta, decay, d_k, d_v);
        }
    }
    MNN_CONCURRENCY_END();
}

void CPULinearAttention::short_conv(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto qkvTensor   = inputs[0];
    auto convWTensor = inputs[3];
    auto outTensor   = outputs[0];

    const int8_t* qkvPtr   = qkvTensor->host<int8_t>();
    const int8_t* convWPtr = convWTensor->host<int8_t>();
    int8_t* outPtr         = outTensor->host<int8_t>();

    const int B      = qkvTensor->length(0);
    const int D      = qkvTensor->length(1);   // 3H
    const int L      = qkvTensor->length(2);
    const int H      = D / 3;
    const int K_conv = convWTensor->length(2);
    const int convStateSize = K_conv - 1;
    const int bytes = mBytes;

    int8_t* convPadded   = mConvPadded->host<int8_t>();
    int8_t* convOut = mConvOut->host<int8_t>();
    int8_t* convStatePtr = mStateCache->mConvState->host<int8_t>();

    int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();
    const int totalLen = convStateSize + L;
    const int totalChannels = B * H;

    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        for (int idx = (int)tId; idx < totalChannels; idx += threadNum) {
            int b = idx / H;
            int h = idx % H;

            // 1a. Compute Bx = B_[b,h,:] * x_[b,h,:] and build padded input
            int8_t* padded = convPadded + idx * totalLen * bytes;
            const int8_t* stateChannel = convStatePtr + idx * convStateSize * bytes;
            ::memcpy(padded, stateChannel, convStateSize * bytes);

            for (int l = 0; l < L; ++l) {
                float b_val = _readElement(qkvPtr, b * D * L + h * L + l, bytes);
                float x_val = _readElement(qkvPtr, b * D * L + (2 * H + h) * L + l, bytes);
                _writeElement(padded, convStateSize + l, b_val * x_val, bytes);
            }

            // 1b. Depthwise Conv1D (no SiLU)
            int8_t* out = convOut + idx * L * bytes;
            for (int l = 0; l < L; ++l) {
                float sum = 0.0f;
                for (int k = 0; k < K_conv; ++k) {
                    sum += _readElement(padded, l + k, bytes) * _readElement(convWPtr, h * K_conv + k, bytes);
                }
                _writeElement(out, l, sum, bytes);
            }

            // 1c. Update conv state
            const int8_t* newState = padded + (totalLen - convStateSize) * bytes;
            int8_t* dstState = convStatePtr + idx * convStateSize * bytes;
            ::memcpy(dstState, newState, convStateSize * bytes);
        }
    }
    MNN_CONCURRENCY_END();

    // Step 2: y = C_ * conv_out, transpose to output [B, L, 1, H]
    MNN_CONCURRENCY_BEGIN(tId, threadNum) {
        for (int idx = (int)tId; idx < totalChannels; idx += threadNum) {
            int b = idx / H;
            int h = idx % H;

            for (int l = 0; l < L; ++l) {
                float c_val = _readElement(qkvPtr, b * D * L + (H + h) * L + l, bytes);
                float conv_val = _readElement(convOut, idx * L + l, bytes);
                _writeElement(outPtr, (b * L + l) * H + h, c_val * conv_val, bytes);
            }
        }
    }
    MNN_CONCURRENCY_END();
}

bool CPULinearAttention::onClone(Backend* bn, const Op* op, Execution** dst) {
    if (nullptr == dst) {
        return true;
    }
    auto tmp = new CPULinearAttention(bn, op);
    // Share persistent state buffers between prefill and decode Executions
    tmp->mStateCache = mStateCache;
    *dst = tmp;
    return true;
}

CPULinearAttention::CPULinearAttention(Backend *backend, const MNN::Op* op) : Execution(backend) {
    auto param = op->main_as_LinearAttentionParam();
    mAttentionType = param->attn_type()->str();
    mNumKHeads = param->num_k_heads();
    mNumVHeads = param->num_v_heads();
    mHeadKDim = param->head_k_dim();
    mHeadVDim = param->head_v_dim();
    mUseQKL2Norm = param->use_qk_l2norm();
    mBytes = static_cast<CPUBackend*>(backend)->functions()->bytes;
    mStateCache.reset(new StateCache);
    mMeta = (KVMeta*)(backend->getMetaPtr());
    mPrefixCacheDir = static_cast<CPUBackend*>(backend)->getRuntime()->hint().prefixcacheDirPath;
}

CPULinearAttention::~CPULinearAttention() {

}

class CPULinearAttentionCreator : 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 {
        return new CPULinearAttention(backend, op);
    }
};

REGISTER_CPU_OP_CREATOR_TRANSFORMER(CPULinearAttentionCreator, OpType_LinearAttention);

} // namespace MNN

#endif // MNN_SUPPORT_TRANSFORMER_FUSE