wehub-resource-sync
975 lines
45 KiB
C++
975 lines
45 KiB
C++
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
|
|
|
#include "VulkanAttention.hpp"
|
|
#include "core/Macro.h"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "backend/vulkan/vulkan/vulkan_wrapper.h"
|
|
#include <climits>
|
|
|
|
namespace MNN {
|
|
|
|
static inline float _invSqrt(float x) {
|
|
return 1.0f / ::sqrtf(x);
|
|
}
|
|
|
|
static uint32_t _selectSoftmaxLocalSize(int totalLen, uint32_t maxSizeX, uint32_t maxInvocations) {
|
|
if (totalLen <= 1) {
|
|
return 1;
|
|
}
|
|
uint32_t cap = 128;
|
|
cap = ALIMIN(cap, maxSizeX);
|
|
cap = ALIMIN(cap, maxInvocations);
|
|
cap = ALIMIN(cap, (uint32_t)totalLen);
|
|
uint32_t localSize = 1;
|
|
while ((localSize << 1) <= cap) {
|
|
localSize <<= 1;
|
|
}
|
|
return localSize;
|
|
}
|
|
|
|
static constexpr int kAttentionPrefillKBlock = 512;
|
|
|
|
static bool _supportDecodeQ1Subgroup(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;
|
|
if ((subgroup.ops & required) != required) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void VulkanAttention::KVCache::reset() {
|
|
maxLen = 0;
|
|
kvHeadNum = 0;
|
|
headDim = 0;
|
|
fp16 = false;
|
|
key = nullptr;
|
|
value = nullptr;
|
|
}
|
|
|
|
void VulkanAttention::KVCache::ensureCapacity(VulkanBackend* vkBn, int requiredLen, int kvH, int d, bool useFP16) {
|
|
MNN_ASSERT(requiredLen >= 0);
|
|
MNN_ASSERT(kvH > 0);
|
|
MNN_ASSERT(d > 0);
|
|
if (kvHeadNum != kvH || headDim != d || fp16 != useFP16 || nullptr == key || nullptr == value) {
|
|
reset();
|
|
kvHeadNum = kvH;
|
|
headDim = d;
|
|
fp16 = useFP16;
|
|
maxLen = requiredLen + expandChunk;
|
|
maxLen = ALIMAX(maxLen, expandChunk);
|
|
const size_t bytes = fp16 ? sizeof(uint16_t) : sizeof(float);
|
|
const size_t bufSize = (size_t)maxLen * (size_t)kvHeadNum * (size_t)headDim * bytes;
|
|
key.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, bufSize, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT));
|
|
value.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, bufSize, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT));
|
|
return;
|
|
}
|
|
if (requiredLen <= maxLen) {
|
|
return;
|
|
}
|
|
const int oldMaxLen = maxLen;
|
|
maxLen = requiredLen + expandChunk;
|
|
const size_t bytes = fp16 ? sizeof(uint16_t) : sizeof(float);
|
|
const size_t newSize = (size_t)maxLen * (size_t)kvHeadNum * (size_t)headDim * bytes;
|
|
std::shared_ptr<VulkanBuffer> newKey(new VulkanBuffer(vkBn->getMemoryPool(), false, newSize, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_SRC_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT));
|
|
std::shared_ptr<VulkanBuffer> newValue(new VulkanBuffer(vkBn->getMemoryPool(), false, newSize, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_SRC_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT));
|
|
// Preserve old content.
|
|
//
|
|
// cacheKey is packed as [kvHeadNum, headDim/4, maxLen, 4], so changing maxLen changes the row stride and we must repack.
|
|
// cacheValue is kvh-major as [kvHeadNum, maxLen, headDim], so changing maxLen changes the kvh stride and we must repack too.
|
|
const size_t oldSize = key->size();
|
|
if (oldSize > 0) {
|
|
// Value: repack kvh blocks with new stride.
|
|
{
|
|
const VkDeviceSize rowBytes = (VkDeviceSize)oldMaxLen * (VkDeviceSize)headDim * (VkDeviceSize)bytes;
|
|
const VkDeviceSize srcStride = rowBytes;
|
|
const VkDeviceSize dstStride = (VkDeviceSize)maxLen * (VkDeviceSize)headDim * (VkDeviceSize)bytes;
|
|
std::vector<VkBufferCopy> regions;
|
|
regions.reserve((size_t)kvHeadNum);
|
|
for (int kvh = 0; kvh < kvHeadNum; ++kvh) {
|
|
VkBufferCopy c;
|
|
c.srcOffset = (VkDeviceSize)kvh * srcStride;
|
|
c.dstOffset = (VkDeviceSize)kvh * dstStride;
|
|
c.size = rowBytes;
|
|
regions.emplace_back(c);
|
|
}
|
|
vkBn->copyGPUToGPUBufferRegions(value->buffer(), newValue->buffer(), regions.data(), (uint32_t)regions.size());
|
|
}
|
|
|
|
// Key: repack rows with new stride.
|
|
const int d4Size = headDim / 4;
|
|
MNN_ASSERT(d4Size > 0);
|
|
const uint32_t rowCount = (uint32_t)kvHeadNum * (uint32_t)d4Size;
|
|
const VkDeviceSize vec4Bytes = (VkDeviceSize)(4 * bytes);
|
|
const VkDeviceSize srcRowStride = (VkDeviceSize)oldMaxLen * vec4Bytes;
|
|
const VkDeviceSize dstRowStride = (VkDeviceSize)maxLen * vec4Bytes;
|
|
std::vector<VkBufferCopy> regions;
|
|
regions.reserve(rowCount);
|
|
for (uint32_t r = 0; r < rowCount; ++r) {
|
|
VkBufferCopy c;
|
|
c.srcOffset = (VkDeviceSize)r * srcRowStride;
|
|
c.dstOffset = (VkDeviceSize)r * dstRowStride;
|
|
c.size = srcRowStride;
|
|
regions.emplace_back(c);
|
|
}
|
|
vkBn->copyGPUToGPUBufferRegions(key->buffer(), newKey->buffer(), regions.data(), (uint32_t)regions.size());
|
|
}
|
|
key = newKey;
|
|
value = newValue;
|
|
}
|
|
|
|
VulkanAttention::VulkanAttention(const Op* op, Backend* bn) : VulkanBasicExecution(bn), mOp(op) {
|
|
auto vkBn = static_cast<VulkanBackend*>(bn);
|
|
mUseFP16 = vkBn->useFP16();
|
|
mMeta = reinterpret_cast<KVMeta*>(vkBn->getMetaPtr());
|
|
if (nullptr != op && nullptr != op->main_as_AttentionParam()) {
|
|
mNeedKvCache = op->main_as_AttentionParam()->kv_cache();
|
|
}
|
|
mKVCache.reset(new KVCache);
|
|
mParam = vkBn->allocUniform(nullptr, sizeof(GpuParam));
|
|
if (!mNeedKvCache) {
|
|
std::vector<VkDescriptorType> typesAttn{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // output
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // query
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // keyIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // valueIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheKey
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheValue
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // mask
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string attnName = "glsl_attention_fused_";
|
|
if (mUseFP16) {
|
|
attnName += "FP16_";
|
|
}
|
|
attnName += "comp";
|
|
mAttentionLegacyPipeline = vkBn->getPipeline(attnName, typesAttn);
|
|
MNN_ASSERT(nullptr != mAttentionLegacyPipeline);
|
|
mAttentionLegacySet.reset(mAttentionLegacyPipeline->createSet());
|
|
return;
|
|
}
|
|
|
|
// kv_cache=true path: pre-create update/prefill/decode pipelines to avoid resize cold-start.
|
|
{
|
|
std::vector<VkDescriptorType> typesUpdate{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // keyIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // valueIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheKey
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheValue
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string updateName = "glsl_attention_kvcache_update_";
|
|
if (mUseFP16) {
|
|
updateName += "FP16_";
|
|
}
|
|
updateName += "comp";
|
|
mUpdatePipeline = vkBn->getPipeline(updateName, typesUpdate);
|
|
MNN_ASSERT(nullptr != mUpdatePipeline);
|
|
mUpdateSet.reset(mUpdatePipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesRearrange{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // queryOut
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // queryIn
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string rqName = "glsl_attention_prefill_rearrange_q_";
|
|
if (mUseFP16) {
|
|
rqName += "FP16_";
|
|
}
|
|
rqName += "comp";
|
|
mRearrangeQPipeline = vkBn->getPipeline(rqName, typesRearrange);
|
|
MNN_ASSERT(nullptr != mRearrangeQPipeline);
|
|
mRearrangeQSet.reset(mRearrangeQPipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesInit{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // m
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // l
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // alpha
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // oAcc
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string initName = "glsl_attention_prefill_kblock_init_state_";
|
|
if (mUseFP16) {
|
|
initName += "FP16_";
|
|
}
|
|
initName += "comp";
|
|
mInitStatePipeline = vkBn->getPipeline(initName, typesInit);
|
|
MNN_ASSERT(nullptr != mInitStatePipeline);
|
|
mInitStateSet.reset(mInitStatePipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesQK{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // qk
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // query
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheKey
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // mask
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
|
|
std::string qkName = "glsl_attention_prefill_kblock_qk_";
|
|
if (mUseFP16) {
|
|
qkName += "FP16_";
|
|
}
|
|
qkName += "comp";
|
|
mQKBlockPipeline = vkBn->getPipeline(qkName, typesQK);
|
|
MNN_ASSERT(nullptr != mQKBlockPipeline);
|
|
mQKBlockSet.reset(mQKBlockPipeline->createSet());
|
|
|
|
std::string qkFullName = "glsl_attention_prefill_kblock_qk_full_";
|
|
if (mUseFP16) {
|
|
qkFullName += "FP16_";
|
|
}
|
|
qkFullName += "comp";
|
|
mQKBlockFullPipeline = vkBn->getPipeline(qkFullName, typesQK);
|
|
MNN_ASSERT(nullptr != mQKBlockFullPipeline);
|
|
mQKBlockFullSet.reset(mQKBlockFullPipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesSoftmax{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // w
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // qk
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // m
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // l
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // alpha
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string softmaxName = "glsl_attention_prefill_kblock_softmax_online_";
|
|
if (mUseFP16) {
|
|
softmaxName += "FP16_";
|
|
}
|
|
softmaxName += "comp";
|
|
const auto& limits = vkBn->getDevice().proty().limits;
|
|
const int kBlock4 = UP_DIV(kAttentionPrefillKBlock, 4) * 4;
|
|
const int maxK4 = UP_DIV(kBlock4, 4);
|
|
uint32_t localSize = _selectSoftmaxLocalSize(maxK4, (uint32_t)limits.maxComputeWorkGroupSize[0],
|
|
(uint32_t)limits.maxComputeWorkGroupInvocations);
|
|
mSoftmaxOnlinePipeline = vkBn->getPipeline(softmaxName, typesSoftmax, {localSize});
|
|
MNN_ASSERT(nullptr != mSoftmaxOnlinePipeline);
|
|
mSoftmaxOnlineSet.reset(mSoftmaxOnlinePipeline->createSet());
|
|
mSoftmaxOnlineLocalSize = localSize;
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesQKV{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // oAcc
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // w
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheValue
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // alpha
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string qkvName = "glsl_attention_prefill_kblock_qkv_acc_";
|
|
if (mUseFP16) {
|
|
qkvName += "FP16_";
|
|
}
|
|
qkvName += "comp";
|
|
mQKVAccPipeline = vkBn->getPipeline(qkvName, typesQKV);
|
|
MNN_ASSERT(nullptr != mQKVAccPipeline);
|
|
mQKVAccSet.reset(mQKVAccPipeline->createSet());
|
|
|
|
std::string qkvFullName = "glsl_attention_prefill_kblock_qkv_acc_full_";
|
|
if (mUseFP16) {
|
|
qkvFullName += "FP16_";
|
|
}
|
|
qkvFullName += "comp";
|
|
mQKVAccFullPipeline = vkBn->getPipeline(qkvFullName, typesQKV);
|
|
MNN_ASSERT(nullptr != mQKVAccFullPipeline);
|
|
mQKVAccFullSet.reset(mQKVAccFullPipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesFinal{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // output
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // oAcc
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // l
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string finalName = "glsl_attention_prefill_kblock_finalize_";
|
|
if (mUseFP16) {
|
|
finalName += "FP16_";
|
|
}
|
|
finalName += "comp";
|
|
mFinalizePipeline = vkBn->getPipeline(finalName, typesFinal);
|
|
MNN_ASSERT(nullptr != mFinalizePipeline);
|
|
mFinalizeSet.reset(mFinalizePipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> typesAttn{
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // output
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // query
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // keyIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // valueIn
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheKey
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // cacheValue
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, // mask
|
|
VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER // param
|
|
};
|
|
std::string attnName = "glsl_attention_fused_packed_";
|
|
if (mUseFP16) {
|
|
attnName += "FP16_";
|
|
}
|
|
attnName += "comp";
|
|
mAttentionPipeline = vkBn->getPipeline(attnName, typesAttn);
|
|
MNN_ASSERT(nullptr != mAttentionPipeline);
|
|
mAttentionSet.reset(mAttentionPipeline->createSet());
|
|
|
|
if (_supportDecodeQ1Subgroup(vkBn->getDevice())) {
|
|
mDecodeQ1SubgroupLocalSize = vkBn->getDevice().getSubgroupSize();
|
|
if (mDecodeQ1SubgroupLocalSize > 0) {
|
|
std::string decodeQ1Name = "glsl_attention_decode_q1_subgroup_";
|
|
if (mUseFP16) {
|
|
decodeQ1Name += "FP16_";
|
|
}
|
|
decodeQ1Name += "comp";
|
|
mDecodeQ1SubgroupPipeline = vkBn->getPipeline(decodeQ1Name, typesAttn, {mDecodeQ1SubgroupLocalSize});
|
|
if (nullptr != mDecodeQ1SubgroupPipeline) {
|
|
mDecodeQ1SubgroupSet.reset(mDecodeQ1SubgroupPipeline->createSet());
|
|
}
|
|
|
|
std::string decodeQ1HD128Name = "glsl_attention_decode_q1_subgroup_hd128_";
|
|
if (mUseFP16) {
|
|
decodeQ1HD128Name += "FP16_";
|
|
}
|
|
decodeQ1HD128Name += "comp";
|
|
mDecodeQ1SubgroupHD128Pipeline = vkBn->getPipeline(decodeQ1HD128Name, typesAttn, {mDecodeQ1SubgroupLocalSize});
|
|
if (nullptr != mDecodeQ1SubgroupHD128Pipeline) {
|
|
mDecodeQ1SubgroupHD128Set.reset(mDecodeQ1SubgroupHD128Pipeline->createSet());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
VulkanAttention::~VulkanAttention() {
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
if (mTempQuery) {
|
|
vkBn->onReleaseBuffer(mTempQuery.get(), Backend::DYNAMIC);
|
|
mTempQuery.reset();
|
|
}
|
|
if (mTempQKBlock) {
|
|
vkBn->onReleaseBuffer(mTempQKBlock.get(), Backend::DYNAMIC);
|
|
mTempQKBlock.reset();
|
|
}
|
|
if (mTempWBlock) {
|
|
vkBn->onReleaseBuffer(mTempWBlock.get(), Backend::DYNAMIC);
|
|
mTempWBlock.reset();
|
|
}
|
|
if (mTempM) {
|
|
vkBn->onReleaseBuffer(mTempM.get(), Backend::DYNAMIC);
|
|
mTempM.reset();
|
|
}
|
|
if (mTempL) {
|
|
vkBn->onReleaseBuffer(mTempL.get(), Backend::DYNAMIC);
|
|
mTempL.reset();
|
|
}
|
|
if (mTempAlpha) {
|
|
vkBn->onReleaseBuffer(mTempAlpha.get(), Backend::DYNAMIC);
|
|
mTempAlpha.reset();
|
|
}
|
|
if (mTempOAcc) {
|
|
vkBn->onReleaseBuffer(mTempOAcc.get(), Backend::DYNAMIC);
|
|
mTempOAcc.reset();
|
|
}
|
|
vkBn->recycleUniform(mParam);
|
|
}
|
|
|
|
bool VulkanAttention::onClone(Backend* bn, const Op* op, VulkanBasicExecution** dst) {
|
|
if (nullptr == dst) {
|
|
return true;
|
|
}
|
|
auto res = new VulkanAttention(op, bn);
|
|
res->mKVCache = mKVCache;
|
|
res->mMeta = mMeta;
|
|
*dst = res;
|
|
return true;
|
|
}
|
|
|
|
ErrorCode VulkanAttention::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) {
|
|
MNN_ASSERT(!inputs.empty());
|
|
MNN_ASSERT(!outputs.empty());
|
|
auto query = inputs[0];
|
|
auto key = inputs[1];
|
|
auto value = inputs[2];
|
|
MNN_ASSERT(nullptr != query && nullptr != key && nullptr != value);
|
|
MNN_ASSERT(query->dimensions() == 4);
|
|
MNN_ASSERT(key->dimensions() == 4);
|
|
MNN_ASSERT(value->dimensions() == 4);
|
|
MNN_ASSERT(query->length(0) == 1);
|
|
MNN_ASSERT(key->length(0) == 1);
|
|
MNN_ASSERT(value->length(0) == 1);
|
|
mQueryLen = query->length(1);
|
|
mKeyLen = key->length(1);
|
|
mHeadNum = query->length(2);
|
|
mHeadDim = query->length(3);
|
|
mKvHeadNum = key->length(2);
|
|
MNN_ASSERT(mHeadNum > 0 && mKvHeadNum > 0);
|
|
MNN_ASSERT(mHeadNum % mKvHeadNum == 0);
|
|
MNN_ASSERT(mHeadDim > 0);
|
|
MNN_ASSERT((mHeadDim & 3) == 0);
|
|
MNN_ASSERT(mHeadDim <= 256);
|
|
MNN_ASSERT(value->length(1) == mKeyLen);
|
|
MNN_ASSERT(value->length(2) == mKvHeadNum);
|
|
MNN_ASSERT(value->length(3) == mHeadDim);
|
|
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
auto cmd = cmdBuffer->get();
|
|
|
|
#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
|
|
|
|
const bool usePrefill = mNeedKvCache && mQueryLen > 1;
|
|
mUsePrefill = usePrefill;
|
|
|
|
if (mNeedKvCache) {
|
|
MNN_ASSERT(nullptr != mUpdatePipeline);
|
|
MNN_ASSERT(nullptr != mUpdateSet);
|
|
|
|
// Dispatch: KV update (x=dim/4, y=keyLen, z=kvHeadNum).
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_kvcache_update_FP16_comp" : "glsl_attention_kvcache_update_comp",
|
|
mUpdatePipeline, mUpdateSet, UP_DIV(mHeadDim / 4, 8), mKeyLen, mKvHeadNum);
|
|
// NOTE: KV cache buffers may be reallocated in onBeforeExecute (descriptor set updated there), so we must not
|
|
// record a VkBufferMemoryBarrier with a stale VkBuffer handle here. Use a global memory barrier instead.
|
|
{
|
|
VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER};
|
|
barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT;
|
|
barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_TRANSFER_READ_BIT;
|
|
vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT,
|
|
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 1, &barrier,
|
|
0, nullptr, 0, nullptr);
|
|
}
|
|
}
|
|
|
|
if (usePrefill) {
|
|
constexpr int K_BLOCK = kAttentionPrefillKBlock;
|
|
int pastLenForPrefill = 0;
|
|
if (mNeedKvCache) {
|
|
MNN_ASSERT(nullptr != mMeta);
|
|
MNN_ASSERT(mMeta->n_reserve == 0);
|
|
MNN_ASSERT(mMeta->computeReverseSize() == 0);
|
|
const int previous = (int)mMeta->previous;
|
|
const int remove = (int)mMeta->remove;
|
|
MNN_ASSERT(previous >= 0);
|
|
MNN_ASSERT(remove >= 0);
|
|
MNN_ASSERT(remove <= previous);
|
|
pastLenForPrefill = previous - remove;
|
|
}
|
|
mPrefillTotalLen = pastLenForPrefill + mKeyLen;
|
|
mQueryLen4 = UP_DIV(mQueryLen, 4) * 4;
|
|
MNN_ASSERT(mPrefillTotalLen > 0);
|
|
|
|
const int64_t queryElementsI64 = (int64_t)mHeadNum * (int64_t)mHeadDim * (int64_t)mQueryLen4;
|
|
MNN_ASSERT(queryElementsI64 > 0 && queryElementsI64 <= (int64_t)INT_MAX);
|
|
const int queryElements = (int)queryElementsI64;
|
|
|
|
const int kBlock4 = UP_DIV(K_BLOCK, 4) * 4;
|
|
const int64_t rowCountI64 = (int64_t)mQueryLen * (int64_t)mHeadNum;
|
|
MNN_ASSERT(rowCountI64 > 0 && rowCountI64 <= (int64_t)INT_MAX);
|
|
const int rowCount = (int)rowCountI64;
|
|
|
|
const int64_t qkElementsI64 = (int64_t)rowCount * (int64_t)kBlock4;
|
|
MNN_ASSERT(qkElementsI64 > 0 && qkElementsI64 <= (int64_t)INT_MAX);
|
|
const int qkElements = (int)qkElementsI64;
|
|
|
|
const int64_t oaccElementsI64 = (int64_t)rowCount * (int64_t)mHeadDim;
|
|
MNN_ASSERT(oaccElementsI64 > 0 && oaccElementsI64 <= (int64_t)INT_MAX);
|
|
const int oaccElements = (int)oaccElementsI64;
|
|
|
|
// Acquire workspace tensors fresh each onEncode (Backend::DYNAMIC). They are released at the end of this
|
|
// call so other ops within the same resize can reuse the pool chunks. Descriptor sets capture
|
|
// (VkBuffer, offset) below; the underlying GPU memory stays alive past release because the parent
|
|
// VkBuffer is owned by the pool, and command-buffer order + barriers serialize cross-op access.
|
|
// M / L / Alpha / OAcc must be FP32 even when the backend runs FP16 -> int tensor forces 4-byte storage.
|
|
auto acquireTemp = [&](std::shared_ptr<Tensor>& t, Tensor* dev) -> bool {
|
|
t.reset(dev);
|
|
return vkBn->onAcquireBuffer(t.get(), Backend::DYNAMIC);
|
|
};
|
|
if (!acquireTemp(mTempQuery, Tensor::createDevice<float>({queryElements}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempQKBlock, Tensor::createDevice<float>({qkElements}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempWBlock, Tensor::createDevice<float>({qkElements}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempM, Tensor::createDevice<int>({rowCount}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempL, Tensor::createDevice<int>({rowCount}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempAlpha, Tensor::createDevice<int>({rowCount}))) return OUT_OF_MEMORY;
|
|
if (!acquireTemp(mTempOAcc, Tensor::createDevice<int>({oaccElements}))) return OUT_OF_MEMORY;
|
|
|
|
MNN_ASSERT(nullptr != mRearrangeQPipeline);
|
|
MNN_ASSERT(nullptr != mRearrangeQSet);
|
|
MNN_ASSERT(nullptr != mInitStatePipeline);
|
|
MNN_ASSERT(nullptr != mInitStateSet);
|
|
MNN_ASSERT(nullptr != mQKBlockPipeline);
|
|
MNN_ASSERT(nullptr != mQKBlockSet);
|
|
MNN_ASSERT(nullptr != mQKBlockFullPipeline);
|
|
MNN_ASSERT(nullptr != mQKBlockFullSet);
|
|
MNN_ASSERT(nullptr != mSoftmaxOnlinePipeline);
|
|
MNN_ASSERT(nullptr != mSoftmaxOnlineSet);
|
|
MNN_ASSERT(nullptr != mQKVAccPipeline);
|
|
MNN_ASSERT(nullptr != mQKVAccSet);
|
|
MNN_ASSERT(nullptr != mQKVAccFullPipeline);
|
|
MNN_ASSERT(nullptr != mQKVAccFullSet);
|
|
MNN_ASSERT(nullptr != mFinalizePipeline);
|
|
MNN_ASSERT(nullptr != mFinalizeSet);
|
|
|
|
// Bind workspace + uniform descriptor slots here (their (VkBuffer, offset) is stable across executes).
|
|
// Cache (cacheKey/cacheValue) and external I/O (query/mask/output) are bound in onBeforeExecute because
|
|
// KV cache may be reallocated by ensureCapacity().
|
|
auto tqBuf = vkBn->getTensorBuffer(mTempQuery.get());
|
|
auto qkBuf = vkBn->getTensorBuffer(mTempQKBlock.get());
|
|
auto wBuf = vkBn->getTensorBuffer(mTempWBlock.get());
|
|
auto mBuf = vkBn->getTensorBuffer(mTempM.get());
|
|
auto lBuf = vkBn->getTensorBuffer(mTempL.get());
|
|
auto aBuf = vkBn->getTensorBuffer(mTempAlpha.get());
|
|
auto oBuf = vkBn->getTensorBuffer(mTempOAcc.get());
|
|
|
|
mRearrangeQSet->writeBuffer(tqBuf.first->buffer(), 0, vkBn->getTensorSize(mTempQuery.get()), tqBuf.second);
|
|
mRearrangeQSet->writeBuffer(mParam->buffer(), 2, mParam->size());
|
|
|
|
mInitStateSet->writeBuffer(mBuf.first->buffer(), 0, vkBn->getTensorSize(mTempM.get()), mBuf.second);
|
|
mInitStateSet->writeBuffer(lBuf.first->buffer(), 1, vkBn->getTensorSize(mTempL.get()), lBuf.second);
|
|
mInitStateSet->writeBuffer(aBuf.first->buffer(), 2, vkBn->getTensorSize(mTempAlpha.get()), aBuf.second);
|
|
mInitStateSet->writeBuffer(oBuf.first->buffer(), 3, vkBn->getTensorSize(mTempOAcc.get()), oBuf.second);
|
|
mInitStateSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
|
|
mQKBlockSet->writeBuffer(qkBuf.first->buffer(), 0, vkBn->getTensorSize(mTempQKBlock.get()), qkBuf.second);
|
|
mQKBlockSet->writeBuffer(tqBuf.first->buffer(), 1, vkBn->getTensorSize(mTempQuery.get()), tqBuf.second);
|
|
mQKBlockSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
|
|
mQKBlockFullSet->writeBuffer(qkBuf.first->buffer(), 0, vkBn->getTensorSize(mTempQKBlock.get()), qkBuf.second);
|
|
mQKBlockFullSet->writeBuffer(tqBuf.first->buffer(), 1, vkBn->getTensorSize(mTempQuery.get()), tqBuf.second);
|
|
mQKBlockFullSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
|
|
mSoftmaxOnlineSet->writeBuffer(wBuf.first->buffer(), 0, vkBn->getTensorSize(mTempWBlock.get()), wBuf.second);
|
|
mSoftmaxOnlineSet->writeBuffer(qkBuf.first->buffer(), 1, vkBn->getTensorSize(mTempQKBlock.get()), qkBuf.second);
|
|
mSoftmaxOnlineSet->writeBuffer(mBuf.first->buffer(), 2, vkBn->getTensorSize(mTempM.get()), mBuf.second);
|
|
mSoftmaxOnlineSet->writeBuffer(lBuf.first->buffer(), 3, vkBn->getTensorSize(mTempL.get()), lBuf.second);
|
|
mSoftmaxOnlineSet->writeBuffer(aBuf.first->buffer(), 4, vkBn->getTensorSize(mTempAlpha.get()), aBuf.second);
|
|
mSoftmaxOnlineSet->writeBuffer(mParam->buffer(), 5, mParam->size());
|
|
|
|
mQKVAccSet->writeBuffer(oBuf.first->buffer(), 0, vkBn->getTensorSize(mTempOAcc.get()), oBuf.second);
|
|
mQKVAccSet->writeBuffer(wBuf.first->buffer(), 1, vkBn->getTensorSize(mTempWBlock.get()), wBuf.second);
|
|
mQKVAccSet->writeBuffer(aBuf.first->buffer(), 3, vkBn->getTensorSize(mTempAlpha.get()), aBuf.second);
|
|
mQKVAccSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
|
|
mQKVAccFullSet->writeBuffer(oBuf.first->buffer(), 0, vkBn->getTensorSize(mTempOAcc.get()), oBuf.second);
|
|
mQKVAccFullSet->writeBuffer(wBuf.first->buffer(), 1, vkBn->getTensorSize(mTempWBlock.get()), wBuf.second);
|
|
mQKVAccFullSet->writeBuffer(aBuf.first->buffer(), 3, vkBn->getTensorSize(mTempAlpha.get()), aBuf.second);
|
|
mQKVAccFullSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
|
|
mFinalizeSet->writeBuffer(oBuf.first->buffer(), 1, vkBn->getTensorSize(mTempOAcc.get()), oBuf.second);
|
|
mFinalizeSet->writeBuffer(lBuf.first->buffer(), 2, vkBn->getTensorSize(mTempL.get()), lBuf.second);
|
|
mFinalizeSet->writeBuffer(mParam->buffer(), 3, mParam->size());
|
|
|
|
// 1) Rearrange Q to packed-D Qtmp: (x=qLen4, y=headDim/4, z=headNum)
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_prefill_rearrange_q_FP16_comp" : "glsl_attention_prefill_rearrange_q_comp",
|
|
mRearrangeQPipeline, mRearrangeQSet, UP_DIV(mQueryLen4, 8), UP_DIV(mHeadDim / 4, 8), mHeadNum);
|
|
cmdBuffer->barrierSource(tqBuf.first->buffer(), tqBuf.second, vkBn->getTensorSize(mTempQuery.get()));
|
|
|
|
// K-block prefill: online softmax in K dimension to avoid O(qLen*totalLen) intermediates.
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_prefill_kblock_init_state_FP16_comp" : "glsl_attention_prefill_kblock_init_state_comp",
|
|
mInitStatePipeline, mInitStateSet, UP_DIV((uint32_t)mQueryLen * (uint32_t)mHeadNum * (uint32_t)mHeadDim, 256),
|
|
1, 1);
|
|
cmdBuffer->barrierSource(mBuf.first->buffer(), mBuf.second, vkBn->getTensorSize(mTempM.get()));
|
|
cmdBuffer->barrierSource(lBuf.first->buffer(), lBuf.second, vkBn->getTensorSize(mTempL.get()));
|
|
cmdBuffer->barrierSource(aBuf.first->buffer(), aBuf.second, vkBn->getTensorSize(mTempAlpha.get()));
|
|
cmdBuffer->barrierSource(oBuf.first->buffer(), oBuf.second, vkBn->getTensorSize(mTempOAcc.get()));
|
|
|
|
struct QKPushConst {
|
|
uint32_t kStart;
|
|
uint32_t blockLen;
|
|
};
|
|
struct SoftmaxPushConst {
|
|
uint32_t blockLen;
|
|
};
|
|
|
|
auto dispatchWithPushConst = [&](const char* name, const VulkanPipeline* pipeline,
|
|
const std::shared_ptr<VulkanLayout::DescriptorSet>& set, uint32_t x, uint32_t y,
|
|
uint32_t z, const void* pcData, uint32_t pcSize) {
|
|
#ifdef ENABLE_VULKAN_TIME_PROFILE
|
|
auto* profiler = vkBn->timeProfiler();
|
|
if (nullptr != profiler) {
|
|
VulkanTimeProfileScope scope(profiler, cmd, name, VulkanTimeProfiler::Kind::Shader);
|
|
pipeline->bind(cmd, set->get());
|
|
vkCmdPushConstants(cmd, pipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, pcSize, pcData);
|
|
vkCmdDispatch(cmd, x, y, z);
|
|
return;
|
|
}
|
|
#endif
|
|
pipeline->bind(cmd, set->get());
|
|
vkCmdPushConstants(cmd, pipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, pcSize, pcData);
|
|
vkCmdDispatch(cmd, x, y, z);
|
|
};
|
|
|
|
const int totalLen = mPrefillTotalLen;
|
|
const int kBlock = K_BLOCK;
|
|
for (int kStart = 0; kStart < totalLen; kStart += kBlock) {
|
|
const int blockLen = ALIMIN(kBlock, totalLen - kStart);
|
|
const int blockLen4 = UP_DIV(blockLen, 4) * 4;
|
|
const int blockLen4_4 = UP_DIV(blockLen4, 4);
|
|
|
|
// 2) QK block: (x=blockLen4/4, y=qLen4/4, z=headNum)
|
|
QKPushConst pcQK{(uint32_t)kStart, (uint32_t)blockLen};
|
|
const bool fullBlock = (blockLen == kBlock) && (kStart + kBlock <= totalLen);
|
|
const VulkanPipeline* qkPipe = fullBlock ? mQKBlockFullPipeline : mQKBlockPipeline;
|
|
const std::shared_ptr<VulkanLayout::DescriptorSet>& qkSet = fullBlock ? mQKBlockFullSet : mQKBlockSet;
|
|
const char* qkName = nullptr;
|
|
if (fullBlock) {
|
|
qkName = mUseFP16 ? "glsl_attention_prefill_kblock_qk_full_FP16_comp" : "glsl_attention_prefill_kblock_qk_full_comp";
|
|
} else {
|
|
qkName = mUseFP16 ? "glsl_attention_prefill_kblock_qk_FP16_comp" : "glsl_attention_prefill_kblock_qk_comp";
|
|
}
|
|
dispatchWithPushConst(qkName, qkPipe, qkSet, UP_DIV((uint32_t)blockLen4_4, 8),
|
|
UP_DIV((uint32_t)UP_DIV(mQueryLen4, 4), 8), (uint32_t)mHeadNum, &pcQK, sizeof(pcQK));
|
|
cmdBuffer->barrierSource(qkBuf.first->buffer(), qkBuf.second, vkBn->getTensorSize(mTempQKBlock.get()));
|
|
|
|
// 3) Softmax online: updates m/l and writes unnormalized w (x=headNum, y=qLen)
|
|
SoftmaxPushConst pcSM{(uint32_t)blockLen};
|
|
dispatchWithPushConst(mUseFP16 ? "glsl_attention_prefill_kblock_softmax_online_FP16_comp"
|
|
: "glsl_attention_prefill_kblock_softmax_online_comp",
|
|
mSoftmaxOnlinePipeline, mSoftmaxOnlineSet, (uint32_t)mHeadNum, (uint32_t)mQueryLen, 1, &pcSM,
|
|
sizeof(pcSM));
|
|
cmdBuffer->barrierSource(wBuf.first->buffer(), wBuf.second, vkBn->getTensorSize(mTempWBlock.get()));
|
|
cmdBuffer->barrierSource(mBuf.first->buffer(), mBuf.second, vkBn->getTensorSize(mTempM.get()));
|
|
cmdBuffer->barrierSource(lBuf.first->buffer(), lBuf.second, vkBn->getTensorSize(mTempL.get()));
|
|
cmdBuffer->barrierSource(aBuf.first->buffer(), aBuf.second, vkBn->getTensorSize(mTempAlpha.get()));
|
|
|
|
// 4) QKV accumulate: (x=headDim/4, y=qLen/2, z=headNum)
|
|
const VulkanPipeline* qkvPipe = fullBlock ? mQKVAccFullPipeline : mQKVAccPipeline;
|
|
const std::shared_ptr<VulkanLayout::DescriptorSet>& qkvSet = fullBlock ? mQKVAccFullSet : mQKVAccSet;
|
|
const char* qkvName = nullptr;
|
|
if (fullBlock) {
|
|
qkvName = mUseFP16 ? "glsl_attention_prefill_kblock_qkv_acc_full_FP16_comp"
|
|
: "glsl_attention_prefill_kblock_qkv_acc_full_comp";
|
|
} else {
|
|
qkvName =
|
|
mUseFP16 ? "glsl_attention_prefill_kblock_qkv_acc_FP16_comp" : "glsl_attention_prefill_kblock_qkv_acc_comp";
|
|
}
|
|
dispatchWithPushConst(qkvName, qkvPipe, qkvSet, UP_DIV((uint32_t)(mHeadDim / 4), 8),
|
|
UP_DIV((uint32_t)UP_DIV(mQueryLen, 2), 8), (uint32_t)mHeadNum, &pcQK, sizeof(pcQK));
|
|
cmdBuffer->barrierSource(oBuf.first->buffer(), oBuf.second, vkBn->getTensorSize(mTempOAcc.get()));
|
|
}
|
|
|
|
// 5) Finalize: output = oAcc / l
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_prefill_kblock_finalize_FP16_comp" : "glsl_attention_prefill_kblock_finalize_comp",
|
|
mFinalizePipeline, mFinalizeSet, UP_DIV((uint32_t)(mHeadDim / 4), 8),
|
|
UP_DIV((uint32_t)UP_DIV(mQueryLen, 2), 8), (uint32_t)mHeadNum);
|
|
|
|
auto releaseTemp = [&](std::shared_ptr<Tensor>& t) {
|
|
if (t) {
|
|
vkBn->onReleaseBuffer(t.get(), Backend::DYNAMIC);
|
|
t.reset();
|
|
}
|
|
};
|
|
releaseTemp(mTempQuery);
|
|
releaseTemp(mTempQKBlock);
|
|
releaseTemp(mTempWBlock);
|
|
releaseTemp(mTempM);
|
|
releaseTemp(mTempL);
|
|
releaseTemp(mTempAlpha);
|
|
releaseTemp(mTempOAcc);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
// Decode (or kv_cache disabled): keep fused shader.
|
|
mQueryLen4 = 0;
|
|
if (mTempQuery) {
|
|
vkBn->onReleaseBuffer(mTempQuery.get(), Backend::DYNAMIC);
|
|
mTempQuery.reset();
|
|
}
|
|
if (mTempQKBlock) {
|
|
vkBn->onReleaseBuffer(mTempQKBlock.get(), Backend::DYNAMIC);
|
|
mTempQKBlock.reset();
|
|
}
|
|
if (mTempWBlock) {
|
|
vkBn->onReleaseBuffer(mTempWBlock.get(), Backend::DYNAMIC);
|
|
mTempWBlock.reset();
|
|
}
|
|
if (mTempM) {
|
|
vkBn->onReleaseBuffer(mTempM.get(), Backend::DYNAMIC);
|
|
mTempM.reset();
|
|
}
|
|
if (mTempL) {
|
|
vkBn->onReleaseBuffer(mTempL.get(), Backend::DYNAMIC);
|
|
mTempL.reset();
|
|
}
|
|
if (mTempAlpha) {
|
|
vkBn->onReleaseBuffer(mTempAlpha.get(), Backend::DYNAMIC);
|
|
mTempAlpha.reset();
|
|
}
|
|
if (mTempOAcc) {
|
|
vkBn->onReleaseBuffer(mTempOAcc.get(), Backend::DYNAMIC);
|
|
mTempOAcc.reset();
|
|
}
|
|
mPrefillTotalLen = 0;
|
|
|
|
if (mNeedKvCache) {
|
|
const bool useDecodeQ1Subgroup =
|
|
(mQueryLen == 1) && (nullptr != mDecodeQ1SubgroupPipeline) && (nullptr != mDecodeQ1SubgroupSet);
|
|
if (useDecodeQ1Subgroup) {
|
|
const bool useHD128 = (mHeadDim == 128) && (nullptr != mDecodeQ1SubgroupHD128Pipeline) &&
|
|
(nullptr != mDecodeQ1SubgroupHD128Set);
|
|
if (useHD128) {
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_decode_q1_subgroup_hd128_FP16_comp"
|
|
: "glsl_attention_decode_q1_subgroup_hd128_comp",
|
|
mDecodeQ1SubgroupHD128Pipeline, mDecodeQ1SubgroupHD128Set, (uint32_t)mHeadNum, 1, 1);
|
|
} else {
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_decode_q1_subgroup_FP16_comp"
|
|
: "glsl_attention_decode_q1_subgroup_comp",
|
|
mDecodeQ1SubgroupPipeline, mDecodeQ1SubgroupSet, (uint32_t)mHeadNum, 1, 1);
|
|
}
|
|
} else {
|
|
MNN_ASSERT(nullptr != mAttentionPipeline);
|
|
MNN_ASSERT(nullptr != mAttentionSet);
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_fused_packed_FP16_comp" : "glsl_attention_fused_packed_comp",
|
|
mAttentionPipeline, mAttentionSet, UP_DIV(mHeadNum, 8), UP_DIV(mQueryLen, 8), 1);
|
|
}
|
|
} else {
|
|
MNN_ASSERT(nullptr != mAttentionLegacyPipeline);
|
|
MNN_ASSERT(nullptr != mAttentionLegacySet);
|
|
dispatchWithProfile(mUseFP16 ? "glsl_attention_fused_FP16_comp" : "glsl_attention_fused_comp",
|
|
mAttentionLegacyPipeline, mAttentionLegacySet, UP_DIV(mHeadNum, 8), UP_DIV(mQueryLen, 8), 1);
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode VulkanAttention::onBeforeExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
MNN_ASSERT(!inputs.empty());
|
|
MNN_ASSERT(!outputs.empty());
|
|
auto query = inputs[0];
|
|
auto key = inputs[1];
|
|
auto value = inputs[2];
|
|
auto output = outputs[0];
|
|
MNN_ASSERT(nullptr != query && nullptr != key && nullptr != value && nullptr != output);
|
|
MNN_ASSERT(query->length(1) == mQueryLen);
|
|
MNN_ASSERT(key->length(1) == mKeyLen);
|
|
MNN_ASSERT(query->length(2) == mHeadNum);
|
|
MNN_ASSERT(key->length(2) == mKvHeadNum);
|
|
MNN_ASSERT(query->length(3) == mHeadDim);
|
|
MNN_ASSERT(key->length(3) == mHeadDim);
|
|
MNN_ASSERT(value->length(1) == mKeyLen);
|
|
MNN_ASSERT(value->length(2) == mKvHeadNum);
|
|
MNN_ASSERT(value->length(3) == mHeadDim);
|
|
MNN_ASSERT(query->length(0) == 1);
|
|
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
|
|
int pastLenForCompute = 0;
|
|
if (mNeedKvCache) {
|
|
MNN_ASSERT(nullptr != mMeta);
|
|
MNN_ASSERT(mMeta->n_reserve == 0);
|
|
MNN_ASSERT(mMeta->computeReverseSize() == 0);
|
|
const int previous = (int)mMeta->previous;
|
|
const int remove = (int)mMeta->remove;
|
|
const int add = (int)mMeta->add;
|
|
MNN_ASSERT(previous >= 0);
|
|
MNN_ASSERT(remove >= 0);
|
|
MNN_ASSERT(add >= 0);
|
|
MNN_ASSERT(add <= mKeyLen);
|
|
MNN_ASSERT(remove <= previous);
|
|
pastLenForCompute = previous - remove;
|
|
// Ensure capacity for compute window (pastLen + keyLen), because shaders read only from KV cache.
|
|
mKVCache->ensureCapacity(vkBn, pastLenForCompute + mKeyLen, mKvHeadNum, mHeadDim, mUseFP16);
|
|
}
|
|
|
|
const int group = mHeadNum / mKvHeadNum;
|
|
const int totalLenForCompute = pastLenForCompute + mKeyLen;
|
|
|
|
int maskMode = 0;
|
|
int maskQlen = 0;
|
|
int maskKvlen = 0;
|
|
const Tensor* mask = nullptr;
|
|
if (inputs.size() > 3 && nullptr != inputs[3]) {
|
|
mask = inputs[3];
|
|
MNN_ASSERT(mask->getType() == halide_type_of<float>());
|
|
if (mask->shape().empty()) {
|
|
// Match the transformer-fuse attention convention used by the tests:
|
|
// a shape-empty scalar mask is the lower-triangular causal sentinel.
|
|
maskMode = 2;
|
|
} else {
|
|
const int md = mask->dimensions();
|
|
MNN_ASSERT(md >= 2);
|
|
maskQlen = mask->length(md - 2);
|
|
maskKvlen = mask->length(md - 1);
|
|
MNN_ASSERT(maskQlen == mQueryLen);
|
|
MNN_ASSERT(maskKvlen > 0);
|
|
maskMode = 1;
|
|
}
|
|
}
|
|
|
|
auto gpuParam = reinterpret_cast<GpuParam*>(mParam->map());
|
|
gpuParam->s0[0] = mQueryLen;
|
|
gpuParam->s0[1] = mKeyLen;
|
|
gpuParam->s0[2] = mHeadNum;
|
|
gpuParam->s0[3] = mKvHeadNum;
|
|
gpuParam->s1[0] = mHeadDim;
|
|
gpuParam->s1[1] = group;
|
|
gpuParam->s1[2] = pastLenForCompute;
|
|
gpuParam->s1[3] = totalLenForCompute;
|
|
gpuParam->s2[0] = maskQlen;
|
|
gpuParam->s2[1] = maskKvlen;
|
|
gpuParam->s2[2] = maskMode;
|
|
gpuParam->s2[3] = mNeedKvCache ? mKVCache->maxLen : 0;
|
|
gpuParam->f0[0] = _invSqrt((float)mHeadDim);
|
|
gpuParam->f0[1] = 0.0f;
|
|
gpuParam->f0[2] = 0.0f;
|
|
gpuParam->f0[3] = 0.0f;
|
|
mParam->unmap();
|
|
|
|
// Bind buffers (update + attention). Note: when hasMask == 0, bind query buffer as placeholder.
|
|
auto queryBuf = vkBn->getTensorBuffer(query);
|
|
auto keyBuf = vkBn->getTensorBuffer(key);
|
|
auto valueBuf = vkBn->getTensorBuffer(value);
|
|
auto outBuf = vkBn->getTensorBuffer(output);
|
|
const VkDeviceSize queryOffset = queryBuf.second;
|
|
|
|
const VulkanBuffer* cacheKeyBuf = nullptr;
|
|
const VulkanBuffer* cacheValueBuf = nullptr;
|
|
VkDeviceSize cacheKeyOffset = 0;
|
|
VkDeviceSize cacheValueOffset = 0;
|
|
size_t cacheKeySize = 0;
|
|
size_t cacheValueSize = 0;
|
|
|
|
if (mNeedKvCache) {
|
|
cacheKeyBuf = mKVCache->key.get();
|
|
cacheValueBuf = mKVCache->value.get();
|
|
MNN_ASSERT(nullptr != cacheKeyBuf && nullptr != cacheValueBuf);
|
|
cacheKeySize = cacheKeyBuf->size();
|
|
cacheValueSize = cacheValueBuf->size();
|
|
} else {
|
|
// KV cache disabled: alias cache buffers to current K/V (shaders read only from cache bindings).
|
|
cacheKeyBuf = keyBuf.first;
|
|
cacheValueBuf = valueBuf.first;
|
|
cacheKeyOffset = keyBuf.second;
|
|
cacheValueOffset = valueBuf.second;
|
|
cacheKeySize = vkBn->getTensorSize(key);
|
|
cacheValueSize = vkBn->getTensorSize(value);
|
|
}
|
|
|
|
// Update set (only when KV cache is enabled; kv_cache=false uses legacy fused shader directly on input K/V).
|
|
if (mNeedKvCache) {
|
|
mUpdateSet->writeBuffer(keyBuf.first->buffer(), 0, vkBn->getTensorSize(key), keyBuf.second);
|
|
mUpdateSet->writeBuffer(valueBuf.first->buffer(), 1, vkBn->getTensorSize(value), valueBuf.second);
|
|
mUpdateSet->writeBuffer(cacheKeyBuf->buffer(), 2, cacheKeySize, cacheKeyOffset);
|
|
mUpdateSet->writeBuffer(cacheValueBuf->buffer(), 3, cacheValueSize, cacheValueOffset);
|
|
mUpdateSet->writeBuffer(mParam->buffer(), 4, mParam->size());
|
|
}
|
|
|
|
if (mUsePrefill) {
|
|
MNN_ASSERT(totalLenForCompute == mPrefillTotalLen);
|
|
MNN_ASSERT(mQueryLen4 == UP_DIV(mQueryLen, 4) * 4);
|
|
|
|
// Only the bindings that may change between executes are rewritten below.
|
|
|
|
MNN_ASSERT(nullptr != mRearrangeQSet);
|
|
mRearrangeQSet->writeBuffer(queryBuf.first->buffer(), 1, vkBn->getTensorSize(query), queryBuf.second);
|
|
|
|
mQKBlockSet->writeBuffer(cacheKeyBuf->buffer(), 2, cacheKeySize, cacheKeyOffset);
|
|
mQKBlockFullSet->writeBuffer(cacheKeyBuf->buffer(), 2, cacheKeySize, cacheKeyOffset);
|
|
if (maskMode == 1) {
|
|
auto maskBuf = vkBn->getTensorBuffer(mask);
|
|
mQKBlockSet->writeBuffer(maskBuf.first->buffer(), 3, vkBn->getTensorSize(mask), maskBuf.second);
|
|
mQKBlockFullSet->writeBuffer(maskBuf.first->buffer(), 3, vkBn->getTensorSize(mask), maskBuf.second);
|
|
} else {
|
|
mQKBlockSet->writeBuffer(queryBuf.first->buffer(), 3, vkBn->getTensorSize(query), queryBuf.second);
|
|
mQKBlockFullSet->writeBuffer(queryBuf.first->buffer(), 3, vkBn->getTensorSize(query), queryBuf.second);
|
|
}
|
|
|
|
mQKVAccSet->writeBuffer(cacheValueBuf->buffer(), 2, cacheValueSize, cacheValueOffset);
|
|
mQKVAccFullSet->writeBuffer(cacheValueBuf->buffer(), 2, cacheValueSize, cacheValueOffset);
|
|
|
|
mFinalizeSet->writeBuffer(outBuf.first->buffer(), 0, vkBn->getTensorSize(output), outBuf.second);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
// Attention set (fused). Keep packed fused set for fallback even when decode-q1 subgroup is available.
|
|
auto writeAttentionSet = [&](const std::shared_ptr<VulkanLayout::DescriptorSet>& set) {
|
|
MNN_ASSERT(nullptr != set);
|
|
set->writeBuffer(outBuf.first->buffer(), 0, vkBn->getTensorSize(output), outBuf.second);
|
|
set->writeBuffer(queryBuf.first->buffer(), 1, vkBn->getTensorSize(query), queryBuf.second);
|
|
set->writeBuffer(keyBuf.first->buffer(), 2, vkBn->getTensorSize(key), keyBuf.second);
|
|
set->writeBuffer(valueBuf.first->buffer(), 3, vkBn->getTensorSize(value), valueBuf.second);
|
|
set->writeBuffer(cacheKeyBuf->buffer(), 4, cacheKeySize, cacheKeyOffset);
|
|
set->writeBuffer(cacheValueBuf->buffer(), 5, cacheValueSize, cacheValueOffset);
|
|
if (maskMode == 1) {
|
|
auto maskBuf = vkBn->getTensorBuffer(mask);
|
|
set->writeBuffer(maskBuf.first->buffer(), 6, vkBn->getTensorSize(mask), maskBuf.second);
|
|
} else {
|
|
set->writeBuffer(queryBuf.first->buffer(), 6, vkBn->getTensorSize(query), queryBuf.second);
|
|
}
|
|
set->writeBuffer(mParam->buffer(), 7, mParam->size());
|
|
};
|
|
if (mNeedKvCache) {
|
|
MNN_ASSERT(nullptr != mAttentionSet);
|
|
writeAttentionSet(mAttentionSet);
|
|
if (mQueryLen == 1 && nullptr != mDecodeQ1SubgroupSet) {
|
|
writeAttentionSet(mDecodeQ1SubgroupSet);
|
|
}
|
|
if (mQueryLen == 1 && nullptr != mDecodeQ1SubgroupHD128Set) {
|
|
writeAttentionSet(mDecodeQ1SubgroupHD128Set);
|
|
}
|
|
} else {
|
|
MNN_ASSERT(nullptr != mAttentionLegacySet);
|
|
writeAttentionSet(mAttentionLegacySet);
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class VulkanAttentionCreator : public VulkanBackend::Creator {
|
|
public:
|
|
VulkanBasicExecution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op,
|
|
Backend* backend) const override {
|
|
return new VulkanAttention(op, backend);
|
|
}
|
|
};
|
|
|
|
static bool gResistor = []() {
|
|
VulkanBackend::addCreator(OpType_Attention, new VulkanAttentionCreator);
|
|
return true;
|
|
}();
|
|
|
|
} // namespace MNN
|
|
|
|
#endif // MNN_SUPPORT_TRANSFORMER_FUSE
|