wehub-resource-sync
833 lines
47 KiB
C++
833 lines
47 KiB
C++
//
|
|
// LinearAttentionBufExecution.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/12.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
|
|
|
#include "LinearAttentionBufExecution.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "core/OpCommonUtils.hpp"
|
|
|
|
namespace MNN {
|
|
namespace OpenCL {
|
|
|
|
LinearAttentionBufExecution::LinearAttentionBufExecution(const MNN::Op *op, Backend *backend)
|
|
: CommonExecution(backend, op) {
|
|
mOpenCLBackend = static_cast<OpenCLBackend *>(backend);
|
|
mMeta = (KVMeta*)(backend->getMetaPtr());
|
|
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 OpenCLStateCache);
|
|
}
|
|
|
|
ErrorCode LinearAttentionBufExecution::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
auto qkv = inputs[0];
|
|
int batch = qkv->length(0);
|
|
int convDim = qkv->length(1);
|
|
int seqLen = qkv->length(2);
|
|
|
|
// ─── Chunked prefill: fully independent branch ───
|
|
mUseChunkedPrefill = (seqLen > 1);
|
|
if (mUseChunkedPrefill) {
|
|
return onResizeChunkedPrefill(inputs, outputs);
|
|
}
|
|
|
|
int H = mNumVHeads;
|
|
int dk = mHeadKDim;
|
|
int dv = mHeadVDim;
|
|
int K_conv = inputs[3]->length(2);
|
|
int convStateSize = K_conv - 1;
|
|
int key_dim = mNumKHeads * dk;
|
|
int val_dim = mNumVHeads * dv;
|
|
int gqa_factor = (mNumVHeads > mNumKHeads) ? (mNumVHeads / mNumKHeads) : 1;
|
|
float qScale = 1.0f / sqrt((float)dk);
|
|
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
|
|
// ─── Persistent state buffers (STATIC): allocate once, shared via onClone ───
|
|
int bytesPerElement = mOpenCLBackend->fpBytes();
|
|
if (mStateCache->mRecurrentState.get() == nullptr) {
|
|
// First time: allocate and zero-initialize
|
|
int rnnSize = batch * H * dk * dv;
|
|
mStateCache->mRecurrentState.reset(Tensor::createDevice<float>({rnnSize}));
|
|
bool success = backend()->onAcquireBuffer(mStateCache->mRecurrentState.get(), Backend::STATIC);
|
|
if (!success) return OUT_OF_MEMORY;
|
|
|
|
{
|
|
cl_int res;
|
|
int bufferBytes = rnnSize * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mRecurrentState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mRecurrentState.get()), mapPtr);
|
|
}
|
|
}
|
|
|
|
if (convStateSize > 0) {
|
|
int convStateTotal = batch * convDim * convStateSize;
|
|
mStateCache->mConvState.reset(Tensor::createDevice<float>({convStateTotal}));
|
|
success &= backend()->onAcquireBuffer(mStateCache->mConvState.get(), Backend::STATIC);
|
|
if (!success) return OUT_OF_MEMORY;
|
|
|
|
cl_int res;
|
|
int bufferBytes = convStateTotal * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mConvState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mConvState.get()), mapPtr);
|
|
}
|
|
}
|
|
} else if (seqLen > 1) {
|
|
// Prefill: reset state for new sequence, UNLESS:
|
|
// 1. Loading from prefix cache (PendingRead), or
|
|
// 2. Reusing KV from previous inference (reuse_kv=true, i.e. previous != remove)
|
|
bool loadingFromDisk = (mMeta != nullptr && mMeta->file_flag == KVMeta::PendingRead && mMeta->file_name.size() > 0);
|
|
bool reusingKV = (mMeta != nullptr && mMeta->previous != mMeta->remove);
|
|
if (!loadingFromDisk && !reusingKV) {
|
|
{
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mRecurrentState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mRecurrentState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mRecurrentState.get()), mapPtr);
|
|
}
|
|
}
|
|
if (mStateCache->mConvState.get() != nullptr) {
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mConvState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mConvState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mConvState.get()), mapPtr);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Decode (seqLen == 1): keep existing state untouched
|
|
|
|
// Allocate temporary conv output buffer
|
|
mConvOut.reset(Tensor::createDevice<float>({batch * convDim * seqLen}));
|
|
mOpenCLBackend->onAcquireBuffer(mConvOut.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mConvOut.get(), Backend::DYNAMIC);
|
|
|
|
// Build kernels
|
|
std::set<std::string> buildOptions;
|
|
int local_size = 16;
|
|
buildOptions.emplace("-DLOCAL_SIZE=" + std::to_string(local_size));
|
|
buildOptions.emplace("-DK_SIZE=" + std::to_string(dv));
|
|
|
|
// Kernel 1: Conv1D + SiLU
|
|
mKernelConvSilu = runtime->buildKernel("linear_attention_buf", "linear_attn_conv_silu", buildOptions, mOpenCLBackend->getPrecision());
|
|
|
|
int totalConvSilu = batch * convDim * seqLen;
|
|
mGWSConvSilu = {(uint32_t)totalConvSilu, 1, 1};
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelConvSilu));
|
|
uint32_t lwsConv = std::min(maxWorkGroupSize, (uint32_t)256);
|
|
lwsConv = std::min(lwsConv, (uint32_t)totalConvSilu);
|
|
mLWSConvSilu = {lwsConv, 1, 1};
|
|
|
|
// Kernel 2: Conv state update
|
|
if (convStateSize > 0) {
|
|
mKernelConvStateUpdate = runtime->buildKernel("linear_attention_buf", "linear_attn_conv_state_update", buildOptions, mOpenCLBackend->getPrecision());
|
|
|
|
int totalConvUpdate = batch * convDim * convStateSize;
|
|
mGWSConvStateUpdate = {(uint32_t)totalConvUpdate, 1, 1};
|
|
maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelConvStateUpdate));
|
|
uint32_t lwsUpdate = std::min(maxWorkGroupSize, (uint32_t)256);
|
|
lwsUpdate = std::min(lwsUpdate, (uint32_t)totalConvUpdate);
|
|
mLWSConvStateUpdate = {lwsUpdate, 1, 1};
|
|
}
|
|
|
|
// Kernel 3: Gated Delta Rule
|
|
auto gateDeltaRuleBuildOptions = buildOptions;
|
|
if(seqLen == 1){
|
|
gateDeltaRuleBuildOptions.emplace("-DDECODE_PHASE");
|
|
}
|
|
mKernelGatedDeltaRule = runtime->buildKernel("linear_attention_buf", "linear_attn_gated_delta_rule", gateDeltaRuleBuildOptions, mOpenCLBackend->getPrecision());
|
|
|
|
// Set kernel arguments
|
|
// Kernel 1: conv_silu
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelConvSilu->get().setArg(idx++, totalConvSilu);
|
|
ret |= mKernelConvSilu->get().setArg(idx++, openCLBuffer(inputs[0])); // qkv
|
|
ret |= mKernelConvSilu->get().setArg(idx++, openCLBuffer(mStateCache->mConvState.get())); // conv_state
|
|
ret |= mKernelConvSilu->get().setArg(idx++, openCLBuffer(inputs[3])); // conv_weight
|
|
ret |= mKernelConvSilu->get().setArg(idx++, openCLBuffer(mConvOut.get())); // conv_out
|
|
ret |= mKernelConvSilu->get().setArg(idx++, batch);
|
|
ret |= mKernelConvSilu->get().setArg(idx++, convDim);
|
|
ret |= mKernelConvSilu->get().setArg(idx++, seqLen);
|
|
ret |= mKernelConvSilu->get().setArg(idx++, K_conv);
|
|
ret |= mKernelConvSilu->get().setArg(idx++, convStateSize);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg linear_attn_conv_silu");
|
|
}
|
|
|
|
// Kernel 2: conv_state_update
|
|
if (convStateSize > 0) {
|
|
int totalConvUpdate = batch * convDim * convStateSize;
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, totalConvUpdate);
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, openCLBuffer(inputs[0])); // qkv
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, openCLBuffer(mStateCache->mConvState.get())); // conv_state
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, batch);
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, convDim);
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, seqLen);
|
|
ret |= mKernelConvStateUpdate->get().setArg(idx++, convStateSize);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg linear_attn_conv_state_update");
|
|
}
|
|
|
|
// Kernel 2.5: l2
|
|
if(mUseQKL2Norm){
|
|
auto l2BuildOptions = buildOptions;
|
|
if(seqLen > 1){
|
|
l2BuildOptions.emplace("-DUSE_VEC");
|
|
}
|
|
mKernell2Norm = runtime->buildKernel("linear_attention_buf", "l2_norm", l2BuildOptions, mOpenCLBackend->getPrecision());
|
|
|
|
mGWSl2Norm = {128, (uint32_t)(mNumKHeads * UP_DIV(seqLen, 4)), (uint32_t)(batch * 2)};
|
|
mLWSl2Norm = {128, 1, 1};
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernell2Norm->get().setArg(idx++, openCLBuffer(mConvOut.get())); // conv_out
|
|
ret |= mKernell2Norm->get().setArg(idx++, openCLBuffer(mConvOut.get())); // conv_out
|
|
ret |= mKernell2Norm->get().setArg(idx++, convDim);
|
|
ret |= mKernell2Norm->get().setArg(idx++, dk);
|
|
ret |= mKernell2Norm->get().setArg(idx++, 1);
|
|
ret |= mKernell2Norm->get().setArg(idx++, key_dim);
|
|
ret |= mKernell2Norm->get().setArg(idx++, seqLen);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg l2 norm");
|
|
}
|
|
|
|
// Kernel 3: gated_delta_rule
|
|
{
|
|
mGWSGatedDeltaRule = {(uint32_t)local_size, (uint32_t)UP_DIV(dv, 4) * H * batch};
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, openCLBuffer(mConvOut.get())); // conv_out
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, openCLBuffer(inputs[1])); // gate
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, openCLBuffer(inputs[2])); // beta
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, openCLBuffer(mStateCache->mRecurrentState.get())); // recurrent_state id = 6
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, openCLBuffer(outputs[0])); // attn_out
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, batch);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, convDim);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, seqLen);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, mNumKHeads);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, mNumVHeads);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, dk);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, dv);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, key_dim);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, val_dim);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelGatedDeltaRule->get().setArg(idx++, qScale);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg linear_attn_gated_delta_rule");
|
|
maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelGatedDeltaRule));
|
|
mLWSGatedDeltaRule = {(uint32_t)local_size, 1};
|
|
}
|
|
|
|
// Round up global work sizes to multiples of local work sizes
|
|
mGWSConvSilu[0] = ROUND_UP(mGWSConvSilu[0], mLWSConvSilu[0]);
|
|
if (convStateSize > 0) {
|
|
mGWSConvStateUpdate[0] = ROUND_UP(mGWSConvStateUpdate[0], mLWSConvStateUpdate[0]);
|
|
}
|
|
mGWSGatedDeltaRule[0] = ROUND_UP(mGWSGatedDeltaRule[0], mLWSGatedDeltaRule[0]);
|
|
|
|
// Record kernels for queue recording optimization
|
|
mOpenCLBackend->startRecord(mRecording);
|
|
mOpenCLBackend->recordKernel3d(mKernelConvSilu, mGWSConvSilu, mLWSConvSilu);
|
|
if (convStateSize > 0) {
|
|
mOpenCLBackend->recordKernel3d(mKernelConvStateUpdate, mGWSConvStateUpdate, mLWSConvStateUpdate);
|
|
}
|
|
if(mUseQKL2Norm){
|
|
mOpenCLBackend->recordKernel3d(mKernell2Norm, mGWSl2Norm, mLWSl2Norm);
|
|
}
|
|
mOpenCLBackend->recordKernel2d(mKernelGatedDeltaRule, mGWSGatedDeltaRule, mLWSGatedDeltaRule);
|
|
mOpenCLBackend->endRecord(mRecording);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode LinearAttentionBufExecution::onResizeChunkedPrefill(
|
|
const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
|
|
auto qkv = inputs[0];
|
|
int batch = qkv->length(0);
|
|
int convDim = qkv->length(1);
|
|
int seqLen = qkv->length(2);
|
|
|
|
int H = mNumVHeads;
|
|
int dk = mHeadKDim;
|
|
int dv = mHeadVDim;
|
|
int K_conv = inputs[3]->length(2);
|
|
int convStateSize = K_conv - 1;
|
|
int key_dim = mNumKHeads * dk;
|
|
int gqa_factor = (mNumVHeads > mNumKHeads) ? (mNumVHeads / mNumKHeads) : 1;
|
|
float qScale = 1.0f / sqrt((float)dk);
|
|
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
int bytesPerElement = mOpenCLBackend->fpBytes();
|
|
|
|
// ─── Persistent state buffers (STATIC): allocate once, shared via onClone ───
|
|
if (mStateCache->mRecurrentState.get() == nullptr) {
|
|
int rnnSize = batch * H * dk * dv;
|
|
mStateCache->mRecurrentState.reset(Tensor::createDevice<float>({rnnSize}));
|
|
mStateCache->mRecurrentStateTune.reset(Tensor::createDevice<float>({rnnSize}));
|
|
bool success = backend()->onAcquireBuffer(mStateCache->mRecurrentState.get(), Backend::STATIC);
|
|
if (!success) return OUT_OF_MEMORY;
|
|
|
|
{
|
|
cl_int res;
|
|
int bufferBytes = rnnSize * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mRecurrentState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mRecurrentState.get()), mapPtr);
|
|
}
|
|
}
|
|
|
|
if (convStateSize > 0) {
|
|
int convStateTotal = batch * convDim * convStateSize;
|
|
mStateCache->mConvState.reset(Tensor::createDevice<float>({convStateTotal}));
|
|
success &= backend()->onAcquireBuffer(mStateCache->mConvState.get(), Backend::STATIC);
|
|
if (!success) return OUT_OF_MEMORY;
|
|
|
|
cl_int res;
|
|
int bufferBytes = convStateTotal * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mConvState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mConvState.get()), mapPtr);
|
|
}
|
|
}
|
|
} else {
|
|
// Prefill (seqLen > 1): reset state for new sequence
|
|
{
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mRecurrentState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mRecurrentState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mRecurrentState.get()), mapPtr);
|
|
}
|
|
}
|
|
if (mStateCache->mConvState.get() != nullptr) {
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mConvState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mConvState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mConvState.get()), mapPtr);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Allocate temporary buffers ───
|
|
// IMPORTANT: All DYNAMIC buffers that are used together during execution must have
|
|
// overlapping lifetimes (acquire all before releasing any) to prevent the memory
|
|
// planner from aliasing them. mConvOutPrefill is read by C2-C5 and C7 concurrently
|
|
// with chunk buffers, so they must all be alive simultaneously.
|
|
mConvOutPrefill.reset(Tensor::createDevice<float>({batch * convDim * seqLen}));
|
|
mOpenCLBackend->onAcquireBuffer(mStateCache->mRecurrentStateTune.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mStateCache->mRecurrentStateTune.get(), Backend::DYNAMIC);
|
|
|
|
int chunkSize = mChunkSize;
|
|
int numChunks = UP_DIV(seqLen, chunkSize);
|
|
mNumChunks = numChunks;
|
|
|
|
// Allocate intermediate buffers
|
|
// Critical buffers use float32 for precision (allocate 2x elements in half mode to get 4N bytes = N floats)
|
|
// Non-critical buffers use FLOAT (matches onAcquireBuffer precision)
|
|
int fpBytes = mOpenCLBackend->fpBytes();
|
|
auto f32Elems = [fpBytes](int n) { return (n * 4 + fpBytes - 1) / fpBytes; };
|
|
mGCumsumBuf.reset(Tensor::createDevice<float>({f32Elems(batch * H * numChunks * chunkSize)}));
|
|
mAttnMatrixBuf.reset(Tensor::createDevice<float>({f32Elems(batch * H * numChunks * chunkSize * chunkSize)}));
|
|
mVCorrectedBuf.reset(Tensor::createDevice<float>({f32Elems(batch * H * numChunks * chunkSize * dv)}));
|
|
mKCumdecayBuf.reset(Tensor::createDevice<float>({f32Elems(batch * H * numChunks * chunkSize * dk)}));
|
|
mVNewBuf.reset(Tensor::createDevice<float>({f32Elems(batch * H * chunkSize * dv)}));
|
|
// Acquire all buffers used concurrently during execution BEFORE releasing any
|
|
mOpenCLBackend->onAcquireBuffer(mConvOutPrefill.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onAcquireBuffer(mGCumsumBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onAcquireBuffer(mAttnMatrixBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onAcquireBuffer(mVCorrectedBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onAcquireBuffer(mKCumdecayBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onAcquireBuffer(mVNewBuf.get(), Backend::DYNAMIC);
|
|
// Release all together — planner now sees overlapping lifetimes, no aliasing
|
|
mOpenCLBackend->onReleaseBuffer(mConvOutPrefill.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mGCumsumBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mAttnMatrixBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mVCorrectedBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mKCumdecayBuf.get(), Backend::DYNAMIC);
|
|
mOpenCLBackend->onReleaseBuffer(mVNewBuf.get(), Backend::DYNAMIC);
|
|
|
|
// ─── Build common kernels for prefill ───
|
|
std::set<std::string> buildOptions;
|
|
|
|
int local_size = 16;
|
|
buildOptions.emplace("-DLOCAL_SIZE=" + std::to_string(local_size));
|
|
buildOptions.emplace("-DK_SIZE=" + std::to_string(dv));
|
|
// Conv1D + SiLU
|
|
mKernelConvSiluPrefill = runtime->buildKernel("linear_attention_buf", "linear_attn_conv_silu", buildOptions, mOpenCLBackend->getPrecision());
|
|
int totalConvSilu = batch * convDim * seqLen;
|
|
mGWSConvSiluPrefill = {(uint32_t)totalConvSilu, 1, 1};
|
|
{
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelConvSiluPrefill));
|
|
uint32_t lwsConv = std::min(maxWorkGroupSize, (uint32_t)256);
|
|
lwsConv = std::min(lwsConv, (uint32_t)totalConvSilu);
|
|
mLWSConvSiluPrefill = {lwsConv, 1, 1};
|
|
}
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, totalConvSilu);
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, openCLBuffer(inputs[0]));
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, openCLBuffer(mStateCache->mConvState.get()));
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, openCLBuffer(inputs[3]));
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, batch);
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, convDim);
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, seqLen);
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, K_conv);
|
|
ret |= mKernelConvSiluPrefill->get().setArg(idx++, convStateSize);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg linear_attn_conv_silu (prefill)");
|
|
}
|
|
mGWSConvSiluPrefill[0] = ROUND_UP(mGWSConvSiluPrefill[0], mLWSConvSiluPrefill[0]);
|
|
|
|
// Conv state update
|
|
if (convStateSize > 0) {
|
|
mKernelConvStateUpdatePrefill = runtime->buildKernel("linear_attention_buf", "linear_attn_conv_state_update", buildOptions, mOpenCLBackend->getPrecision());
|
|
int totalConvUpdate = batch * convDim * convStateSize;
|
|
mGWSConvStateUpdatePrefill = {(uint32_t)totalConvUpdate, 1, 1};
|
|
{
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelConvStateUpdatePrefill));
|
|
uint32_t lwsUpdate = std::min(maxWorkGroupSize, (uint32_t)256);
|
|
lwsUpdate = std::min(lwsUpdate, (uint32_t)totalConvUpdate);
|
|
mLWSConvStateUpdatePrefill = {lwsUpdate, 1, 1};
|
|
}
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, totalConvUpdate);
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, openCLBuffer(inputs[0]));
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, openCLBuffer(mStateCache->mConvState.get()));
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, batch);
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, convDim);
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, seqLen);
|
|
ret |= mKernelConvStateUpdatePrefill->get().setArg(idx++, convStateSize);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg linear_attn_conv_state_update (prefill)");
|
|
}
|
|
mGWSConvStateUpdatePrefill[0] = ROUND_UP(mGWSConvStateUpdatePrefill[0], mLWSConvStateUpdatePrefill[0]);
|
|
}
|
|
|
|
// L2 norm
|
|
if (mUseQKL2Norm) {
|
|
auto l2BuildOptions = buildOptions;
|
|
l2BuildOptions.emplace("-DUSE_VEC");
|
|
mKernell2NormPrefill = runtime->buildKernel("linear_attention_buf", "l2_norm", l2BuildOptions, mOpenCLBackend->getPrecision());
|
|
mGWSl2NormPrefill = {128, (uint32_t)(mNumKHeads * UP_DIV(seqLen, 4)), (uint32_t)(batch * 2)};
|
|
mLWSl2NormPrefill = {128, 1, 1};
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, convDim);
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, dk);
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, 1);
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, key_dim);
|
|
ret |= mKernell2NormPrefill->get().setArg(idx++, seqLen);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg l2 norm (prefill)");
|
|
}
|
|
|
|
// ─── Build chunked prefill kernels ───
|
|
std::set<std::string> chunkOpts = buildOptions;
|
|
chunkOpts.emplace("-DCHUNK_PREFILL");
|
|
chunkOpts.emplace("-DCHUNK_SIZE=" + std::to_string(chunkSize));
|
|
|
|
// C1: chunk_g_cumsum
|
|
mKernelChunkGCumsum = runtime->buildKernel("linear_attention_buf", "chunk_g_cumsum", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkGCumsum = {(uint32_t)H, (uint32_t)numChunks, (uint32_t)batch};
|
|
mLWSChunkGCumsum = {1, 1, 1};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkGCumsum->get().setArg(idx++, openCLBuffer(inputs[1])); // gate
|
|
ret |= mKernelChunkGCumsum->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get())); // g_cumsum
|
|
ret |= mKernelChunkGCumsum->get().setArg(idx++, H);
|
|
ret |= mKernelChunkGCumsum->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkGCumsum->get().setArg(idx++, numChunks);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_g_cumsum");
|
|
}
|
|
|
|
{
|
|
{
|
|
// C2: chunk_build_neumann_attn
|
|
mKernelChunkNeumannAttn0 = runtime->buildKernel("linear_attention_buf", "chunk_build_neumann_attn_step0", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkNeumannAttn0 = {(uint32_t)chunkSize * chunkSize, (uint32_t)(H * numChunks), (uint32_t)batch};
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, openCLBuffer(inputs[2])); // beta
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get()));
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get()));
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, batch);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, H);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkNeumannAttn0->get().setArg(idx++, numChunks);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_build_neumann_attn_step0");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkNeumannAttn0));
|
|
mLWSChunkNeumannAttn0 = localWS3DDefault(mGWSChunkNeumannAttn0, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_build_neumann_attn_step0", mKernelChunkNeumannAttn0, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
}
|
|
{
|
|
// C2: chunk_build_neumann_attn
|
|
mKernelChunkNeumannAttn1 = runtime->buildKernel("linear_attention_buf", "chunk_build_neumann_attn_step1", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkNeumannAttn1 = {(uint32_t)chunkSize, (uint32_t)(H * numChunks), (uint32_t)batch};
|
|
mLWSChunkNeumannAttn1 = {(uint32_t)chunkSize, 1, 1};
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get()));
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, batch);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, H);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkNeumannAttn1->get().setArg(idx++, numChunks);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_build_neumann_attn_step1");
|
|
}
|
|
}
|
|
|
|
// C3: chunk_correct_v
|
|
mKernelChunkCorrectV = runtime->buildKernel("linear_attention_buf", "chunk_correct_v", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkCorrectV = {(uint32_t)UP_DIV(dv, 4), (uint32_t)(chunkSize * numChunks), (uint32_t)(batch * H)};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get()));
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(inputs[2])); // beta
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get()));
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(mVCorrectedBuf.get()));
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, openCLBuffer(mKCumdecayBuf.get()));
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, mGWSChunkCorrectV[0]);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, mGWSChunkCorrectV[1]);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, mGWSChunkCorrectV[2]);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, H);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, dv);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkCorrectV->get().setArg(idx++, numChunks);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_correct_v");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkCorrectV));
|
|
mLWSChunkCorrectV = localWS3DDefault(mGWSChunkCorrectV, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_correct_v", mKernelChunkCorrectV, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
}
|
|
|
|
// C5: chunk_qk_attn (reuses attn_matrix buffer)
|
|
mKernelChunkQKAttn = runtime->buildKernel("linear_attention_buf", "chunk_qk_attn", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkQKAttn = {(uint32_t)chunkSize, (uint32_t)(chunkSize * numChunks), (uint32_t)(batch * H)};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get()));
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get())); // overwrite
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, mGWSChunkQKAttn[0]);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, mGWSChunkQKAttn[1]);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, mGWSChunkQKAttn[2]);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, H);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, numChunks);
|
|
ret |= mKernelChunkQKAttn->get().setArg(idx++, qScale);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_qk_attn");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkQKAttn));
|
|
mLWSChunkQKAttn = localWS3DDefault(mGWSChunkQKAttn, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_qk_attn", mKernelChunkQKAttn, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
}
|
|
|
|
// C6: chunk_compute_vnew (per-chunk, chunk_idx=11 set dynamically)
|
|
mKernelChunkVnew = runtime->buildKernel("linear_attention_buf", "chunk_compute_vnew", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkVnew = {(uint32_t)UP_DIV(dv, 4), (uint32_t)chunkSize, (uint32_t)(batch * H)};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, openCLBuffer(mVCorrectedBuf.get()));
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, openCLBuffer(mKCumdecayBuf.get()));
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, openCLBuffer(mStateCache->mRecurrentStateTune.get())); // arg 2: state (tune first)
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, openCLBuffer(mVNewBuf.get()));
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, mGWSChunkVnew[0]);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, mGWSChunkVnew[1]);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, mGWSChunkVnew[2]);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, dv);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, H);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, numChunks);
|
|
ret |= mKernelChunkVnew->get().setArg(idx++, 0); // chunk_idx placeholder
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_compute_vnew");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkVnew));
|
|
mLWSChunkVnew = localWS3DDefault(mGWSChunkVnew, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_compute_vnew", mKernelChunkVnew, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
// Swap to real state buffer after tuning
|
|
ret |= mKernelChunkVnew->get().setArg(2, openCLBuffer(mStateCache->mRecurrentState.get()));
|
|
}
|
|
|
|
// C6.5: chunk_output (per-chunk, chunk_idx=17 set dynamically)
|
|
mKernelChunkOutput = runtime->buildKernel("linear_attention_buf", "chunk_output", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkOutput = {(uint32_t)UP_DIV(dv, 4) * chunkSize, (uint32_t)H, (uint32_t)batch};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get())); // qk_attn after C5
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(mVNewBuf.get()));
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get()));
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(mStateCache->mRecurrentState.get())); // arg 4: state (tune first)
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, openCLBuffer(outputs[0]));
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, mGWSChunkOutput[0]);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, mGWSChunkOutput[1]);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, mGWSChunkOutput[2]);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, H);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, dv);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, numChunks);
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, 0); // chunk_idx placeholder
|
|
ret |= mKernelChunkOutput->get().setArg(idx++, qScale);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_output");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkOutput));
|
|
mLWSChunkOutput = localWS3DDefault(mGWSChunkOutput, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_output", mKernelChunkOutput, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
}
|
|
|
|
// C7: chunk_output_state_update (per-chunk, chunk_idx=17 set dynamically)
|
|
mKernelChunkOutputUpdate = runtime->buildKernel("linear_attention_buf", "chunk_output_state_update", chunkOpts, mOpenCLBackend->getPrecision());
|
|
mGWSChunkOutputUpdate = {(uint32_t)UP_DIV(dv, 4) * dk, (uint32_t)H, (uint32_t)batch};
|
|
{
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(mConvOutPrefill.get()));
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(mAttnMatrixBuf.get())); // qk_attn after C5
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(mVNewBuf.get()));
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(mGCumsumBuf.get()));
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(mStateCache->mRecurrentStateTune.get())); // arg 4: state (tune first)
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, openCLBuffer(outputs[0]));
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, mGWSChunkOutputUpdate[0]);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, mGWSChunkOutputUpdate[1]);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, mGWSChunkOutputUpdate[2]);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, convDim);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, seqLen);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, H);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, dk);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, dv);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, key_dim);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, gqa_factor);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, numChunks);
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, 0); // chunk_idx placeholder
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(idx++, qScale);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg chunk_output_state_update");
|
|
auto maxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(mKernelChunkOutputUpdate));
|
|
mLWSChunkOutputUpdate = localWS3DDefault(mGWSChunkOutputUpdate, maxWorkGroupSize, mOpenCLBackend->getOpenCLRuntime(), "chunk_output_state_update", mKernelChunkOutputUpdate, mOpenCLBackend->getCLTuneLevel(), "linear_attention_buf").first;
|
|
// Swap to real state buffer after tuning
|
|
ret |= mKernelChunkOutputUpdate->get().setArg(4, openCLBuffer(mStateCache->mRecurrentState.get()));
|
|
}
|
|
|
|
// Round up chunked GWS
|
|
for (auto& gws_lws : std::vector<std::pair<std::vector<uint32_t>*, std::vector<uint32_t>*>>{
|
|
{&mGWSChunkCorrectV, &mLWSChunkCorrectV}, {&mGWSChunkNeumannAttn0, &mLWSChunkNeumannAttn0}, {&mGWSChunkNeumannAttn1, &mLWSChunkNeumannAttn1},
|
|
{&mGWSChunkQKAttn, &mLWSChunkQKAttn}, {&mGWSChunkVnew, &mLWSChunkVnew},{&mGWSChunkOutput, &mLWSChunkOutput},
|
|
{&mGWSChunkOutputUpdate, &mLWSChunkOutputUpdate}}) {
|
|
for (int d = 0; d < 3; ++d) {
|
|
(*gws_lws.first)[d] = ROUND_UP((*gws_lws.first)[d], std::max((uint32_t)1, (*gws_lws.second)[d]));
|
|
}
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode LinearAttentionBufExecution::onExecuteChunkedPrefill(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
int convStateSize = inputs[3]->length(2) - 1;
|
|
|
|
#ifdef ENABLE_OPENCL_TIME_PROFILER
|
|
{ cl::Event event; run3DKernelDefault(mKernelConvSiluPrefill, mGWSConvSiluPrefill, mLWSConvSiluPrefill, runtime, &event); runtime->pushEvent({"linear_attn_conv_silu", event}); }
|
|
if (convStateSize > 0) {
|
|
cl::Event event; run3DKernelDefault(mKernelConvStateUpdatePrefill, mGWSConvStateUpdatePrefill, mLWSConvStateUpdatePrefill, runtime, &event); runtime->pushEvent({"linear_attn_conv_state_update", event});
|
|
}
|
|
if (mUseQKL2Norm) {
|
|
cl::Event event; run3DKernelDefault(mKernell2NormPrefill, mGWSl2NormPrefill, mLWSl2NormPrefill, runtime, &event); runtime->pushEvent({"l2_norm", event});
|
|
}
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkGCumsum, mGWSChunkGCumsum, mLWSChunkGCumsum, runtime, &e); runtime->pushEvent({"chunk_g_cumsum", e}); }
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkNeumannAttn0, mGWSChunkNeumannAttn0, mLWSChunkNeumannAttn0, runtime, &e); runtime->pushEvent({"chunk_build_neumann_attn0", e}); }
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkNeumannAttn1, mGWSChunkNeumannAttn1, mLWSChunkNeumannAttn1, runtime, &e); runtime->pushEvent({"chunk_build_neumann_attn1", e}); }
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkCorrectV, mGWSChunkCorrectV, mLWSChunkCorrectV, runtime, &e); runtime->pushEvent({"chunk_correct_v", e}); }
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkQKAttn, mGWSChunkQKAttn, mLWSChunkQKAttn, runtime, &e); runtime->pushEvent({"chunk_qk_attn", e}); }
|
|
for (int c = 0; c < mNumChunks; ++c) {
|
|
mKernelChunkVnew->get().setArg(11, c);
|
|
mKernelChunkOutput->get().setArg(17, c);
|
|
mKernelChunkOutputUpdate->get().setArg(17, c);
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkVnew, mGWSChunkVnew, mLWSChunkVnew, runtime, &e); runtime->pushEvent({"chunk_vnew_" + std::to_string(c), e}); }
|
|
{ cl::Event e; runKernel2D(mKernelChunkOutput, mGWSChunkOutput, mLWSChunkOutput, runtime, &e); runtime->pushEvent({"chunk_output_" + std::to_string(c), e}); }
|
|
{ cl::Event e; run3DKernelDefault(mKernelChunkOutputUpdate, mGWSChunkOutputUpdate, mLWSChunkOutputUpdate, runtime, &e); runtime->pushEvent({"chunk_update" + std::to_string(c), e}); }
|
|
}
|
|
#else
|
|
// Common kernels
|
|
run3DKernelDefault(mKernelConvSiluPrefill, mGWSConvSiluPrefill, mLWSConvSiluPrefill, runtime);
|
|
if (convStateSize > 0) {
|
|
run3DKernelDefault(mKernelConvStateUpdatePrefill, mGWSConvStateUpdatePrefill, mLWSConvStateUpdatePrefill, runtime);
|
|
}
|
|
if (mUseQKL2Norm) {
|
|
run3DKernelDefault(mKernell2NormPrefill, mGWSl2NormPrefill, mLWSl2NormPrefill, runtime);
|
|
}
|
|
// Chunked prefill: C1 → C2 → C3, C4 → C5 → loop(C6, C7)
|
|
run3DKernelDefault(mKernelChunkGCumsum, mGWSChunkGCumsum, mLWSChunkGCumsum, runtime);
|
|
run3DKernelDefault(mKernelChunkNeumannAttn0, mGWSChunkNeumannAttn0, mLWSChunkNeumannAttn0, runtime);
|
|
run3DKernelDefault(mKernelChunkNeumannAttn1, mGWSChunkNeumannAttn1, mLWSChunkNeumannAttn1, runtime);
|
|
run3DKernelDefault(mKernelChunkCorrectV, mGWSChunkCorrectV, mLWSChunkCorrectV, runtime);
|
|
run3DKernelDefault(mKernelChunkQKAttn, mGWSChunkQKAttn, mLWSChunkQKAttn, runtime);
|
|
for (int c = 0; c < mNumChunks; ++c) {
|
|
mKernelChunkVnew->get().setArg(11, c); // chunk_idx at position 11
|
|
mKernelChunkOutput->get().setArg(17, c); // chunk_idx at position 17
|
|
mKernelChunkOutputUpdate->get().setArg(17, c); // chunk_idx at position 17
|
|
run3DKernelDefault(mKernelChunkVnew, mGWSChunkVnew, mLWSChunkVnew, runtime);
|
|
runKernel2D(mKernelChunkOutput, mGWSChunkOutput, mLWSChunkOutput, runtime);
|
|
run3DKernelDefault(mKernelChunkOutputUpdate, mGWSChunkOutputUpdate, mLWSChunkOutputUpdate, runtime);
|
|
}
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode LinearAttentionBufExecution::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) {
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
int bytesPerElement = mOpenCLBackend->fpBytes();
|
|
if (mStateCache->mConvState.get() != nullptr) {
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mConvState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mConvState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mConvState.get()), mapPtr);
|
|
}
|
|
}
|
|
{
|
|
cl_int res;
|
|
int bufferBytes = mStateCache->mRecurrentState->elementSize() * bytesPerElement;
|
|
void* mapPtr = runtime->commandQueue().enqueueMapBuffer(
|
|
openCLBuffer(mStateCache->mRecurrentState.get()), true, CL_MAP_WRITE, 0, bufferBytes, nullptr, nullptr, &res);
|
|
if (mapPtr != nullptr && res == CL_SUCCESS) {
|
|
::memset(mapPtr, 0, bufferBytes);
|
|
runtime->commandQueue().enqueueUnmapMemObject(openCLBuffer(mStateCache->mRecurrentState.get()), mapPtr);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (mUseChunkedPrefill) {
|
|
return onExecuteChunkedPrefill(inputs, outputs);
|
|
}
|
|
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
int convStateSize = inputs[3]->length(2) - 1;
|
|
|
|
#ifdef ENABLE_OPENCL_TIME_PROFILER
|
|
{
|
|
cl::Event event;
|
|
run3DKernelDefault(mKernelConvSilu, mGWSConvSilu, mLWSConvSilu, runtime, &event);
|
|
runtime->pushEvent({"linear_attn_conv_silu", event});
|
|
}
|
|
if (convStateSize > 0) {
|
|
cl::Event event;
|
|
run3DKernelDefault(mKernelConvStateUpdate, mGWSConvStateUpdate, mLWSConvStateUpdate, runtime, &event);
|
|
runtime->pushEvent({"linear_attn_conv_state_update", event});
|
|
}
|
|
if(mUseQKL2Norm){
|
|
cl::Event event;
|
|
run3DKernelDefault(mKernell2Norm, mGWSl2Norm, mLWSl2Norm, runtime, &event);
|
|
runtime->pushEvent({"l2_norm", event});
|
|
}
|
|
{
|
|
cl::Event event;
|
|
runKernel2D(mKernelGatedDeltaRule, mGWSGatedDeltaRule, mLWSGatedDeltaRule, runtime, &event);
|
|
runtime->pushEvent({"linear_attn_gated_delta_rule", event});
|
|
}
|
|
#else
|
|
if(mOpenCLBackend->isUseRecordQueue()){
|
|
mOpenCLBackend->addRecord(mRecording, mOpRecordUpdateInfo);
|
|
return NO_ERROR;
|
|
}
|
|
run3DKernelDefault(mKernelConvSilu, mGWSConvSilu, mLWSConvSilu, runtime);
|
|
if (convStateSize > 0) {
|
|
run3DKernelDefault(mKernelConvStateUpdate, mGWSConvStateUpdate, mLWSConvStateUpdate, runtime);
|
|
}
|
|
if(mUseQKL2Norm){
|
|
run3DKernelDefault(mKernell2Norm, mGWSl2Norm, mLWSl2Norm, runtime);
|
|
}
|
|
runKernel2D(mKernelGatedDeltaRule, mGWSGatedDeltaRule, mLWSGatedDeltaRule, runtime);
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
bool LinearAttentionBufExecution::onClone(Backend* bn, const Op* op, Execution** dst) {
|
|
if (nullptr == dst) {
|
|
return true;
|
|
}
|
|
auto exe = new LinearAttentionBufExecution(op, bn);
|
|
// Share persistent state buffers between prefill and decode Executions
|
|
exe->mStateCache = mStateCache;
|
|
*dst = exe;
|
|
return true;
|
|
}
|
|
|
|
class LinearAttentionBufCreator : public OpenCLBackend::Creator {
|
|
public:
|
|
virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs,
|
|
const MNN::Op *op, Backend *backend) const override {
|
|
for (int i = 0; i < inputs.size(); ++i) {
|
|
TensorUtils::setTensorSupportPack(inputs[i], false);
|
|
}
|
|
for (int i = 0; i < outputs.size(); ++i) {
|
|
TensorUtils::setTensorSupportPack(outputs[i], false);
|
|
}
|
|
OPENCL_CREATOR_CHECK(new LinearAttentionBufExecution(op, backend));
|
|
}
|
|
};
|
|
REGISTER_OPENCL_OP_CREATOR_TRANSFORMER(LinearAttentionBufCreator, OpType_LinearAttention, BUFFER);
|
|
|
|
} // namespace OpenCL
|
|
} // namespace MNN
|
|
#endif /* MNN_SUPPORT_TRANSFORMER_FUSE */ |