chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,435 @@
|
||||
//
|
||||
// VulkanLinearAttention.cpp
|
||||
// MNN
|
||||
//
|
||||
// Created by MNN on 2026/02/12.
|
||||
// Copyright © 2018, Alibaba Group Holding Limited
|
||||
//
|
||||
|
||||
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
||||
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
#include "VulkanLinearAttention.hpp"
|
||||
#include "backend/vulkan/vulkan/vulkan_wrapper.h"
|
||||
#include "core/Macro.h"
|
||||
#include "core/OpCommonUtils.hpp"
|
||||
|
||||
namespace MNN {
|
||||
|
||||
namespace {
|
||||
|
||||
static bool _supportLinearAttentionSubgroup(const VulkanDevice& device) {
|
||||
const auto& subgroup = device.getSubgroupInfo();
|
||||
if (0 == subgroup.size) {
|
||||
return false;
|
||||
}
|
||||
if (0 == (subgroup.stages & VK_SHADER_STAGE_COMPUTE_BIT)) {
|
||||
return false;
|
||||
}
|
||||
const VkSubgroupFeatureFlags required = VK_SUBGROUP_FEATURE_BASIC_BIT | VK_SUBGROUP_FEATURE_ARITHMETIC_BIT;
|
||||
return (subgroup.ops & required) == required;
|
||||
}
|
||||
|
||||
struct LinearAttnConvSiluParams {
|
||||
ivec4 size0; // batch, conv_dim, seq_len, kernel_size
|
||||
ivec4 size1; // conv_state_size, total, 0, 0
|
||||
};
|
||||
|
||||
struct LinearAttnConvStateUpdateParams {
|
||||
ivec4 size0; // batch, conv_dim, seq_len, conv_state_size
|
||||
ivec4 size1; // total, 0, 0, 0
|
||||
};
|
||||
|
||||
struct LinearAttnQKVPrepParams {
|
||||
ivec4 size0; // batch, conv_dim, seq_len, num_k_heads
|
||||
ivec4 size1; // num_v_heads, head_k_dim, head_v_dim, key_dim
|
||||
ivec4 size2; // val_dim, gqa_factor, use_l2norm, total
|
||||
vec4 size3; // q_scale, 0, 0, 0
|
||||
};
|
||||
|
||||
struct LinearAttnRecurrentParams {
|
||||
ivec4 size0; // batch, seq_len, num_v_heads, head_k_dim
|
||||
ivec4 size1; // head_v_dim, total_rows, 0, 0
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
VulkanLinearAttention::VulkanLinearAttention(const MNN::Op* op, Backend* backend)
|
||||
: VulkanBasicExecution(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();
|
||||
mStateCache.reset(new VulkanLinearAttentionState);
|
||||
|
||||
auto vkBn = static_cast<VulkanBackend*>(backend);
|
||||
mUseFP16 = vkBn->useFP16();
|
||||
mSubgroupSize = vkBn->getDevice().getSubgroupSize();
|
||||
mUseSubgroup = _supportLinearAttentionSubgroup(vkBn->getDevice());
|
||||
mLaneCount = mSubgroupSize > 0 ? mSubgroupSize : 32;
|
||||
mMeta = reinterpret_cast<KVMeta*>(vkBn->getMetaPtr());
|
||||
|
||||
auto shaderKey = [this](const char* base) {
|
||||
std::string key = base;
|
||||
if (mUseFP16) {
|
||||
key += "_FP16";
|
||||
}
|
||||
key += "_comp";
|
||||
return key;
|
||||
};
|
||||
|
||||
{
|
||||
std::vector<VkDescriptorType> types{
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
||||
};
|
||||
mConvSiluPipeline = vkBn->getPipeline(shaderKey("glsl_linear_attn_conv_silu"), types);
|
||||
MNN_ASSERT(nullptr != mConvSiluPipeline);
|
||||
mConvSiluDesSet.reset(mConvSiluPipeline->createSet());
|
||||
mConvSiluParam = vkBn->allocUniform();
|
||||
}
|
||||
{
|
||||
std::vector<VkDescriptorType> types{
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
||||
};
|
||||
mConvStateUpdatePipeline = vkBn->getPipeline(shaderKey("glsl_linear_attn_conv_state_update"), types);
|
||||
MNN_ASSERT(nullptr != mConvStateUpdatePipeline);
|
||||
mConvStateUpdateDesSet.reset(mConvStateUpdatePipeline->createSet());
|
||||
mConvStateUpdateParam = vkBn->allocUniform();
|
||||
}
|
||||
{
|
||||
std::vector<VkDescriptorType> types{
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
||||
};
|
||||
mQKVPrepPipeline = vkBn->getPipeline(shaderKey("glsl_linear_attn_qkv_prep"), types);
|
||||
MNN_ASSERT(nullptr != mQKVPrepPipeline);
|
||||
mQKVPrepDesSet.reset(mQKVPrepPipeline->createSet());
|
||||
mQKVPrepParam = vkBn->allocUniform();
|
||||
}
|
||||
{
|
||||
std::vector<VkDescriptorType> types{
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
||||
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
|
||||
};
|
||||
const char* prefillBase = mUseSubgroup ? "glsl_linear_attn_gated_delta_rule_prefill"
|
||||
: "glsl_linear_attn_gated_delta_rule_prefill_nosubgroup";
|
||||
const char* decodeBase = mUseSubgroup ? "glsl_linear_attn_gated_delta_rule_decode"
|
||||
: "glsl_linear_attn_gated_delta_rule_decode_nosubgroup";
|
||||
mPrefillPipeline = vkBn->getPipeline(shaderKey(prefillBase), types,
|
||||
{mLaneCount, mSubgroupsPerWorkgroup, 1});
|
||||
mDecodePipeline = vkBn->getPipeline(shaderKey(decodeBase), types,
|
||||
{mLaneCount, mSubgroupsPerWorkgroup, 1});
|
||||
#ifdef MNN_VULKAN_LINEAR_ATTN_VERBOSE
|
||||
MNN_PRINT("[VulkanLinearAttention] path=%s, laneCount=%u, rowsPerGroup=%u\n",
|
||||
mUseSubgroup ? "subgroup" : "shared_memory", mLaneCount, mSubgroupsPerWorkgroup);
|
||||
#endif
|
||||
MNN_ASSERT(nullptr != mPrefillPipeline);
|
||||
MNN_ASSERT(nullptr != mDecodePipeline);
|
||||
mPrefillDesSet.reset(mPrefillPipeline->createSet());
|
||||
mDecodeDesSet.reset(mDecodePipeline->createSet());
|
||||
mPrefillParam = vkBn->allocUniform();
|
||||
mDecodeParam = vkBn->allocUniform();
|
||||
}
|
||||
}
|
||||
|
||||
VulkanLinearAttention::~VulkanLinearAttention() {
|
||||
auto vkBn = static_cast<VulkanBackend*>(backend());
|
||||
vkBn->recycleUniform(mConvSiluParam);
|
||||
vkBn->recycleUniform(mConvStateUpdateParam);
|
||||
vkBn->recycleUniform(mQKVPrepParam);
|
||||
vkBn->recycleUniform(mPrefillParam);
|
||||
vkBn->recycleUniform(mDecodeParam);
|
||||
}
|
||||
|
||||
ErrorCode VulkanLinearAttention::ensurePersistentState(VulkanBackend* vkBn, int batch, int convDim, int convStateSize) {
|
||||
const int recurrentSize = batch * mNumVHeads * mHeadVDim * mHeadKDim;
|
||||
const int convSize = batch * convDim * convStateSize;
|
||||
const bool needRealloc = nullptr == mStateCache->mRecurrentState.get() || mStateCache->mBatch != batch ||
|
||||
mStateCache->mConvDim != convDim || mStateCache->mConvStateSize != convStateSize ||
|
||||
mStateCache->mNumVHeads != mNumVHeads || mStateCache->mHeadKDim != mHeadKDim ||
|
||||
mStateCache->mHeadVDim != mHeadVDim;
|
||||
if (!needRealloc) {
|
||||
return NO_ERROR;
|
||||
}
|
||||
|
||||
mStateCache->mConvState.reset();
|
||||
mStateCache->mRecurrentState.reset();
|
||||
|
||||
if (convStateSize > 0) {
|
||||
mStateCache->mConvState.reset(Tensor::createDevice<float>({convSize}));
|
||||
if (!backend()->onAcquireBuffer(mStateCache->mConvState.get(), Backend::STATIC)) {
|
||||
return OUT_OF_MEMORY;
|
||||
}
|
||||
}
|
||||
|
||||
mStateCache->mRecurrentState.reset(Tensor::createDevice<float>({recurrentSize}));
|
||||
if (!backend()->onAcquireBuffer(mStateCache->mRecurrentState.get(), Backend::STATIC)) {
|
||||
return OUT_OF_MEMORY;
|
||||
}
|
||||
|
||||
mStateCache->mBatch = batch;
|
||||
mStateCache->mConvDim = convDim;
|
||||
mStateCache->mConvStateSize = convStateSize;
|
||||
mStateCache->mNumVHeads = mNumVHeads;
|
||||
mStateCache->mHeadKDim = mHeadKDim;
|
||||
mStateCache->mHeadVDim = mHeadVDim;
|
||||
return resetPersistentState(vkBn);
|
||||
}
|
||||
|
||||
ErrorCode VulkanLinearAttention::resetPersistentState(VulkanBackend* vkBn) {
|
||||
if (mStateCache->mConvState.get() != nullptr) {
|
||||
std::vector<uint8_t> zeros(vkBn->getTensorSize(mStateCache->mConvState.get()), 0);
|
||||
auto buf = vkBn->getBuffer(mStateCache->mConvState.get());
|
||||
vkBn->copyToGPUBuffer(zeros.data(), std::get<0>(buf), zeros.size(), std::get<2>(buf));
|
||||
}
|
||||
if (mStateCache->mRecurrentState.get() != nullptr) {
|
||||
std::vector<uint8_t> zeros(vkBn->getTensorSize(mStateCache->mRecurrentState.get()), 0);
|
||||
auto buf = vkBn->getBuffer(mStateCache->mRecurrentState.get());
|
||||
vkBn->copyToGPUBuffer(zeros.data(), std::get<0>(buf), zeros.size(), std::get<2>(buf));
|
||||
}
|
||||
return NO_ERROR;
|
||||
}
|
||||
|
||||
ErrorCode VulkanLinearAttention::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||||
const VulkanCommandPool::Buffer* cmdBuffer) {
|
||||
auto vkBn = static_cast<VulkanBackend*>(backend());
|
||||
auto cmd = cmdBuffer->get();
|
||||
|
||||
MNN_ASSERT(inputs.size() >= 4);
|
||||
MNN_ASSERT(outputs.size() >= 1);
|
||||
|
||||
auto qkv = inputs[0];
|
||||
const int batch = qkv->length(0);
|
||||
const int convDim = qkv->length(1);
|
||||
const int seqLen = qkv->length(2);
|
||||
const int kernelSize = inputs[3]->length(2);
|
||||
const int convStateSize = kernelSize - 1;
|
||||
const int keyDim = mNumKHeads * mHeadKDim;
|
||||
const int valDim = mNumVHeads * mHeadVDim;
|
||||
const int gqaFactor = (mNumVHeads > mNumKHeads) ? (mNumVHeads / mNumKHeads) : 1;
|
||||
const float qScale = 1.0f / ::sqrtf((float)mHeadKDim);
|
||||
|
||||
auto code = ensurePersistentState(vkBn, batch, convDim, convStateSize);
|
||||
if (NO_ERROR != code) {
|
||||
return code;
|
||||
}
|
||||
const bool reusingKV = (nullptr != mMeta && mMeta->previous != mMeta->remove);
|
||||
const bool loadingFromDisk = (mMeta != nullptr && mMeta->file_flag == KVMeta::PendingRead && mMeta->file_name.size() > 0);
|
||||
if (seqLen > 1 && !reusingKV && !loadingFromDisk) {
|
||||
code = resetPersistentState(vkBn);
|
||||
if (NO_ERROR != code) {
|
||||
return code;
|
||||
}
|
||||
}
|
||||
|
||||
const int convOutSize = batch * convDim * seqLen;
|
||||
const int qSize = batch * seqLen * mNumVHeads * mHeadKDim;
|
||||
const int vSize = batch * seqLen * mNumVHeads * mHeadVDim;
|
||||
mConvOut.reset(Tensor::createDevice<float>({convOutSize}));
|
||||
mQ.reset(Tensor::createDevice<float>({qSize}));
|
||||
mK.reset(Tensor::createDevice<float>({qSize}));
|
||||
mV.reset(Tensor::createDevice<float>({vSize}));
|
||||
bool success = backend()->onAcquireBuffer(mConvOut.get(), Backend::DYNAMIC);
|
||||
success = success && backend()->onAcquireBuffer(mQ.get(), Backend::DYNAMIC);
|
||||
success = success && backend()->onAcquireBuffer(mK.get(), Backend::DYNAMIC);
|
||||
success = success && backend()->onAcquireBuffer(mV.get(), Backend::DYNAMIC);
|
||||
if (!success) {
|
||||
return OUT_OF_MEMORY;
|
||||
}
|
||||
backend()->onReleaseBuffer(mV.get(), Backend::DYNAMIC);
|
||||
backend()->onReleaseBuffer(mK.get(), Backend::DYNAMIC);
|
||||
backend()->onReleaseBuffer(mQ.get(), Backend::DYNAMIC);
|
||||
backend()->onReleaseBuffer(mConvOut.get(), Backend::DYNAMIC);
|
||||
|
||||
#ifdef ENABLE_VULKAN_TIME_PROFILE
|
||||
auto dispatchWithProfile = [&](const char* name, const VulkanPipeline* pipeline,
|
||||
const std::shared_ptr<VulkanLayout::DescriptorSet>& set, uint32_t x, uint32_t y,
|
||||
uint32_t z) {
|
||||
auto* profiler = vkBn->timeProfiler();
|
||||
if (nullptr != profiler) {
|
||||
VulkanTimeProfileScope scope(profiler, cmd, name, VulkanTimeProfiler::Kind::Shader);
|
||||
pipeline->bind(cmd, set->get());
|
||||
vkCmdDispatch(cmd, x, y, z);
|
||||
return;
|
||||
}
|
||||
pipeline->bind(cmd, set->get());
|
||||
vkCmdDispatch(cmd, x, y, z);
|
||||
};
|
||||
#else
|
||||
auto dispatchWithProfile = [&](const char*, const VulkanPipeline* pipeline,
|
||||
const std::shared_ptr<VulkanLayout::DescriptorSet>& set, uint32_t x, uint32_t y,
|
||||
uint32_t z) {
|
||||
pipeline->bind(cmd, set->get());
|
||||
vkCmdDispatch(cmd, x, y, z);
|
||||
};
|
||||
#endif
|
||||
|
||||
{
|
||||
LinearAttnConvSiluParams params;
|
||||
params.size0[0] = batch;
|
||||
params.size0[1] = convDim;
|
||||
params.size0[2] = seqLen;
|
||||
params.size0[3] = kernelSize;
|
||||
params.size1[0] = convStateSize;
|
||||
params.size1[1] = batch * convDim * seqLen;
|
||||
params.size1[2] = 0;
|
||||
params.size1[3] = 0;
|
||||
::memcpy(mConvSiluParam->map(), ¶ms, sizeof(params));
|
||||
mConvSiluParam->unmap();
|
||||
|
||||
mConvSiluDesSet->writeBuffer(vkBn->getBuffer(inputs[0]), 0);
|
||||
if (mStateCache->mConvState.get() != nullptr) {
|
||||
mConvSiluDesSet->writeBuffer(vkBn->getBuffer(mStateCache->mConvState.get()), 1);
|
||||
} else {
|
||||
mConvSiluDesSet->writeBuffer(vkBn->getBuffer(mConvOut.get()), 1);
|
||||
}
|
||||
mConvSiluDesSet->writeBuffer(vkBn->getBuffer(inputs[3]), 2);
|
||||
mConvSiluDesSet->writeBuffer(vkBn->getBuffer(mConvOut.get()), 3);
|
||||
mConvSiluDesSet->writeBuffer(mConvSiluParam->buffer(), 4, mConvSiluParam->size());
|
||||
|
||||
dispatchWithProfile("linear_attn_conv_silu", mConvSiluPipeline, mConvSiluDesSet,
|
||||
UP_DIV((uint32_t)(batch * convDim * seqLen), 256), 1, 1);
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mConvOut.get()));
|
||||
}
|
||||
|
||||
if (convStateSize > 0) {
|
||||
LinearAttnConvStateUpdateParams params;
|
||||
params.size0[0] = batch;
|
||||
params.size0[1] = convDim;
|
||||
params.size0[2] = seqLen;
|
||||
params.size0[3] = convStateSize;
|
||||
params.size1[0] = batch * convDim * convStateSize;
|
||||
params.size1[1] = 0;
|
||||
params.size1[2] = 0;
|
||||
params.size1[3] = 0;
|
||||
::memcpy(mConvStateUpdateParam->map(), ¶ms, sizeof(params));
|
||||
mConvStateUpdateParam->unmap();
|
||||
|
||||
mConvStateUpdateDesSet->writeBuffer(vkBn->getBuffer(inputs[0]), 0);
|
||||
mConvStateUpdateDesSet->writeBuffer(vkBn->getBuffer(mStateCache->mConvState.get()), 1);
|
||||
mConvStateUpdateDesSet->writeBuffer(mConvStateUpdateParam->buffer(), 2, mConvStateUpdateParam->size());
|
||||
|
||||
dispatchWithProfile("linear_attn_conv_state_update", mConvStateUpdatePipeline, mConvStateUpdateDesSet,
|
||||
UP_DIV((uint32_t)(batch * convDim * convStateSize), 256), 1, 1);
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mStateCache->mConvState.get()));
|
||||
}
|
||||
|
||||
{
|
||||
LinearAttnQKVPrepParams params;
|
||||
params.size0[0] = batch;
|
||||
params.size0[1] = convDim;
|
||||
params.size0[2] = seqLen;
|
||||
params.size0[3] = mNumKHeads;
|
||||
params.size1[0] = mNumVHeads;
|
||||
params.size1[1] = mHeadKDim;
|
||||
params.size1[2] = mHeadVDim;
|
||||
params.size1[3] = keyDim;
|
||||
params.size2[0] = valDim;
|
||||
params.size2[1] = gqaFactor;
|
||||
params.size2[2] = mUseQKL2Norm ? 1 : 0;
|
||||
params.size2[3] = batch * seqLen * mNumVHeads;
|
||||
params.size3[0] = qScale;
|
||||
params.size3[1] = 0.0f;
|
||||
params.size3[2] = 0.0f;
|
||||
params.size3[3] = 0.0f;
|
||||
::memcpy(mQKVPrepParam->map(), ¶ms, sizeof(params));
|
||||
mQKVPrepParam->unmap();
|
||||
|
||||
mQKVPrepDesSet->writeBuffer(vkBn->getBuffer(mConvOut.get()), 0);
|
||||
mQKVPrepDesSet->writeBuffer(vkBn->getBuffer(mQ.get()), 1);
|
||||
mQKVPrepDesSet->writeBuffer(vkBn->getBuffer(mK.get()), 2);
|
||||
mQKVPrepDesSet->writeBuffer(vkBn->getBuffer(mV.get()), 3);
|
||||
mQKVPrepDesSet->writeBuffer(mQKVPrepParam->buffer(), 4, mQKVPrepParam->size());
|
||||
|
||||
dispatchWithProfile("linear_attn_qkv_prep", mQKVPrepPipeline, mQKVPrepDesSet,
|
||||
UP_DIV((uint32_t)(batch * seqLen * mNumVHeads), 256), 1, 1);
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mQ.get()));
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mK.get()));
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mV.get()));
|
||||
}
|
||||
|
||||
{
|
||||
LinearAttnRecurrentParams params;
|
||||
params.size0[0] = batch;
|
||||
params.size0[1] = seqLen;
|
||||
params.size0[2] = mNumVHeads;
|
||||
params.size0[3] = mHeadKDim;
|
||||
params.size1[0] = mHeadVDim;
|
||||
params.size1[1] = batch * mNumVHeads * mHeadVDim;
|
||||
params.size1[2] = 0;
|
||||
params.size1[3] = 0;
|
||||
|
||||
auto recurrentParam = seqLen == 1 ? mDecodeParam : mPrefillParam;
|
||||
::memcpy(recurrentParam->map(), ¶ms, sizeof(params));
|
||||
recurrentParam->unmap();
|
||||
|
||||
auto recurrentSet = seqLen == 1 ? mDecodeDesSet : mPrefillDesSet;
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(mQ.get()), 0);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(mK.get()), 1);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(mV.get()), 2);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(inputs[1]), 3);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(inputs[2]), 4);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(mStateCache->mRecurrentState.get()), 5);
|
||||
recurrentSet->writeBuffer(vkBn->getBuffer(outputs[0]), 6);
|
||||
recurrentSet->writeBuffer(recurrentParam->buffer(), 7, recurrentParam->size());
|
||||
|
||||
auto recurrentPipeline = seqLen == 1 ? mDecodePipeline : mPrefillPipeline;
|
||||
const uint32_t groupsX = UP_DIV((uint32_t)(batch * mNumVHeads * mHeadVDim), mSubgroupsPerWorkgroup);
|
||||
dispatchWithProfile(seqLen == 1 ? "linear_attn_gated_delta_rule_decode" : "linear_attn_gated_delta_rule_prefill",
|
||||
recurrentPipeline, recurrentSet, groupsX, 1, 1);
|
||||
cmdBuffer->barrierSource(vkBn->getBuffer(mStateCache->mRecurrentState.get()));
|
||||
}
|
||||
|
||||
return NO_ERROR;
|
||||
}
|
||||
|
||||
bool VulkanLinearAttention::onClone(Backend* bn, const Op* op, VulkanBasicExecution** dst) {
|
||||
if (nullptr == dst) {
|
||||
return true;
|
||||
}
|
||||
auto res = new VulkanLinearAttention(op, bn);
|
||||
res->mStateCache = mStateCache;
|
||||
*dst = res;
|
||||
return true;
|
||||
}
|
||||
|
||||
class VulkanLinearAttentionCreator : public VulkanBackend::Creator {
|
||||
public:
|
||||
virtual VulkanBasicExecution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||||
const MNN::Op* op, Backend* backend) const override {
|
||||
auto param = op->main_as_LinearAttentionParam();
|
||||
if (nullptr == param || nullptr == param->attn_type() || param->attn_type()->str() != "gated_delta_rule") {
|
||||
return nullptr;
|
||||
}
|
||||
return new VulkanLinearAttention(op, backend);
|
||||
}
|
||||
};
|
||||
|
||||
static bool gResistor = []() {
|
||||
VulkanBackend::addCreator(OpType_LinearAttention, new VulkanLinearAttentionCreator);
|
||||
return true;
|
||||
}();
|
||||
|
||||
} // namespace MNN
|
||||
|
||||
#endif /* MNN_SUPPORT_TRANSFORMER_FUSE */
|
||||
Reference in New Issue
Block a user