wehub-resource-sync
74 lines
3.1 KiB
C++
74 lines
3.1 KiB
C++
//
|
|
// CPULinearAttention.hpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/10.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
|
|
|
#ifndef CPULINEARATTENTION_HPP
|
|
#define CPULINEARATTENTION_HPP
|
|
|
|
#include <functional>
|
|
#include "core/Execution.hpp"
|
|
#include "core/OpCommonUtils.hpp"
|
|
#include "CPUKVCacheManager.hpp"
|
|
#include "MNN/ErrorCode.hpp"
|
|
|
|
namespace MNN {
|
|
|
|
// shared_ptr-shared across prefill/decode clones (see onClone). All tensors
|
|
// are Backend::STATIC and freed with the backend — no per-Execution release.
|
|
struct StateCache {
|
|
std::shared_ptr<Tensor> mConvState; // Conv1D padding state: [B, D, kernel_size - 1]
|
|
std::shared_ptr<Tensor> mRecurrentState; // Gated Delta Rule recurrent state S: [B, H, d_k, d_v]
|
|
// Post-prefix snapshot. LA state is not token-indexed, so eraseHistory
|
|
// can't truncate per-token; the next prefill restores from here instead.
|
|
std::shared_ptr<Tensor> mConvStateSnapshot;
|
|
std::shared_ptr<Tensor> mRecurrentStateSnapshot;
|
|
bool mSnapshotValid = false;
|
|
// Prefix-cache file index captured once per session (previous == remove);
|
|
// chunks 2..N reuse it instead of re-advancing mMeta->layer_index, which
|
|
// would drift past Full Attention layers and cause SIGBUS in hybrid models.
|
|
// Sentinel -1 = not captured.
|
|
int mPrefixLayerIndex = -1;
|
|
};
|
|
|
|
class CPULinearAttention : public Execution {
|
|
public:
|
|
CPULinearAttention(Backend *backend, const MNN::Op* op);
|
|
virtual ~CPULinearAttention();
|
|
virtual ErrorCode onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) override;
|
|
virtual ErrorCode onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) override;
|
|
virtual bool onClone(Backend* bn, const Op* op, Execution** dst) override;
|
|
void gated_delta_rule_ref(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs);
|
|
void gated_delta_rule_mnn(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs);
|
|
void gated_delta_rule_decode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) const;
|
|
void short_conv(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs);
|
|
private:
|
|
std::string mAttentionType;
|
|
int mHeadKDim;
|
|
int mHeadVDim;
|
|
int mNumKHeads;
|
|
int mNumVHeads;
|
|
bool mUseQKL2Norm;
|
|
int mBytes; // 4 for fp32, 2 for fp16 (Arm82)
|
|
std::shared_ptr<StateCache> mStateCache;
|
|
KVMeta* mMeta;
|
|
std::string mPrefixCacheDir;
|
|
|
|
// Temporary buffers for MNN-optimized path (per-Execution, DYNAMIC)
|
|
std::shared_ptr<Tensor> mConvPadded; // Padded conv input: [B, D, convStateSize + L]
|
|
std::shared_ptr<Tensor> mConvOut; // Conv output after SiLU: [B, D, L]
|
|
std::shared_ptr<Tensor> mThreadLocalBuf; // Per-thread q/k/v/vpred/delta: [threadNum, 2*d_k + 3*d_v]
|
|
std::shared_ptr<Tensor> mDecayBuf; // Pre-computed exp(gate): [B*L*H]
|
|
std::shared_ptr<Tensor> mConvFp32Buf; // fp16 path: per-thread fp32 temp for Conv1D+SiLu
|
|
};
|
|
|
|
} // namespace MNN
|
|
|
|
#endif // CPULINEARATTENTION_HPP
|
|
|
|
#endif // MNN_SUPPORT_TRANSFORMER_FUSE
|