wehub-resource-sync
225 lines
8.4 KiB
Plaintext
225 lines
8.4 KiB
Plaintext
//
|
|
// FmhaCommonExecution.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2024/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
#include "FmhaCommonExecution.hpp"
|
|
|
|
namespace MNN {
|
|
namespace CUDA {
|
|
|
|
template <typename T>
|
|
__global__ void SPLIT_FusedQKV(const size_t count, const T* fused_qkv, T* ptr_q,
|
|
T* ptr_k, T* ptr_v,
|
|
int head_size
|
|
) {
|
|
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (count); i += blockDim.x * gridDim.x) {
|
|
//[B, S, H, 3, D] -> [B, S, H, D]
|
|
const int bsh = i / head_size;
|
|
const int d = i % head_size;
|
|
|
|
ptr_q[i] = fused_qkv[(bsh * 3 + 0) * head_size + d];
|
|
ptr_k[i] = fused_qkv[(bsh * 3 + 1) * head_size + d];
|
|
ptr_v[i] = fused_qkv[(bsh * 3 + 2) * head_size + d];
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SPLIT_FusedKV(const size_t count, const T* fused_kv,
|
|
T* ptr_k, T* ptr_v,
|
|
int head_size
|
|
) {
|
|
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (count); i += blockDim.x * gridDim.x) {
|
|
//[B, S, H, 2, D] -> [B, S, H, D]
|
|
const int bsh = i / head_size;
|
|
const int d = i % head_size;
|
|
|
|
ptr_k[i] = fused_kv[(bsh * 2 + 0) * head_size + d];
|
|
ptr_v[i] = fused_kv[(bsh * 2 + 1) * head_size + d];
|
|
}
|
|
}
|
|
|
|
template <
|
|
int kQueriesPerBlock,
|
|
int kKeysPerBlock,
|
|
int kMaxK
|
|
>
|
|
int FmhaCommonExecution::run_attention(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
using Attention = AttentionKernel<
|
|
cutlass::half_t, // scalar_t
|
|
cutlass::arch::Sm80, // ArchTag
|
|
true, // Memory is aligned
|
|
kQueriesPerBlock,
|
|
kKeysPerBlock,
|
|
kMaxK,
|
|
false, // Supports dropout
|
|
false // Supports bias
|
|
>;
|
|
|
|
typename Attention::Params p;
|
|
// set parameters
|
|
{
|
|
// TODO : Split fused qkv [B, S, H, 3, D] --> [B, S, H, D]
|
|
p.query_ptr = (cutlass::half_t *)mQ_Buffer;//inputs[0]->deviceId();
|
|
p.key_ptr = (cutlass::half_t *)mK_Buffer;//inputs[0]->deviceId();
|
|
p.value_ptr = (cutlass::half_t *)mV_Buffer;//inputs[0]->deviceId();
|
|
p.logsumexp_ptr = nullptr; // Only needed for bw
|
|
p.output_accum_ptr = nullptr;
|
|
if (Attention::kNeedsOutputAccumulatorBuffer) {
|
|
p.output_accum_ptr = (float *)mAcc_Buffer;
|
|
}
|
|
p.output_ptr = (cutlass::half_t *)outputs[0]->deviceId();
|
|
|
|
// TODO: support arbitrary seq lengths
|
|
// if (cu_seqlens_q.has_value()) {
|
|
// p.cu_seqlens_q_ptr = (int32_t*)cu_seqlens_q->data_ptr();
|
|
// p.cu_seqlens_k_ptr = (int32_t*)cu_seqlens_k->data_ptr();
|
|
// }
|
|
|
|
p.scale = 1.0f / sqrtf(mHeadSize);
|
|
|
|
p.num_heads = mNumHeads;
|
|
p.num_batches = mBatchSize;
|
|
p.head_dim = mHeadSize;
|
|
p.head_dim_value = mHeadSizeV;
|
|
p.num_queries = mSeqLen;
|
|
p.num_keys = mSeqLenKV;
|
|
if (false/*options.causal*/) {
|
|
p.custom_mask_type = Attention::CausalFromTopLeft;
|
|
}
|
|
|
|
// All tensors are in BMHK shapes
|
|
p.q_strideH = mHeadSize;
|
|
p.k_strideH = mHeadSize;
|
|
p.v_strideH = mHeadSizeV;
|
|
p.q_strideM = p.q_strideH * mNumHeads;
|
|
p.k_strideM = p.k_strideH * mNumHeads;
|
|
p.v_strideM = p.v_strideH * mNumHeads;
|
|
p.q_strideB = p.q_strideM * mSeqLen;
|
|
p.k_strideB = p.k_strideM * mSeqLenKV;
|
|
p.v_strideB = p.v_strideM * mSeqLenKV;
|
|
p.o_strideM = mHeadSizeV * mNumHeads;
|
|
}
|
|
// launch kernel :)
|
|
constexpr auto kernel_fn = attention_kernel_batched_impl<Attention>;
|
|
int smem_bytes = sizeof(typename Attention::SharedStorage);
|
|
if (smem_bytes > 0xc000) {
|
|
cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
|
|
}
|
|
if (!Attention::check_supported(p)) {
|
|
MNN_ERROR("Attention Kernel does not support these inputs\n");
|
|
return -1;
|
|
}
|
|
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes>>>(p);
|
|
return 0;
|
|
}
|
|
|
|
FmhaCommonExecution::FmhaCommonExecution(const MNN::Op* op, Backend* backend) : Execution(backend) {
|
|
if(op->type() == OpType_FmhaV2) {
|
|
auto fmha_v2_param = op->main_as_FmhaV2Param();
|
|
mNumHeads = fmha_v2_param->heads();
|
|
mType = 0;
|
|
} else if(op->type() == OpType_Fmhca) {
|
|
auto fmhca_param = op->main_as_FmhcaParam();
|
|
mNumHeads = fmhca_param->heads();
|
|
mType = 1;
|
|
}
|
|
}
|
|
ErrorCode FmhaCommonExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto pool = static_cast<CUDABackend*>(backend())->getBufferPool();
|
|
auto runtime = static_cast<CUDABackend*>(backend())->getCUDARuntime();
|
|
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
MNN_ASSERT(output->dimensions() == 3);
|
|
|
|
mBatchSize = output->length(0);
|
|
mSeqLen = output->length(1);
|
|
mHeadSizeV = outputs[0]->length(2)/mNumHeads;
|
|
mHeadSize = mHeadSizeV;
|
|
mSeqLenKV = mSeqLen;
|
|
if(mType == 1) {
|
|
mSeqLenKV = inputs[1]->length(1);
|
|
mHeadSize = inputs[0]->length(2)/mNumHeads;
|
|
}
|
|
|
|
mSM = runtime->compute_capability();
|
|
MemChunk buffer_q;
|
|
if(mType == 0) {
|
|
buffer_q = pool->alloc(mBatchSize * mSeqLen * mHeadSize * mNumHeads * sizeof(half));
|
|
mQ_Buffer = (void*)((uint8_t*)buffer_q.first + buffer_q.second);
|
|
}
|
|
auto buffer_k = pool->alloc(mBatchSize * mSeqLenKV * mHeadSize * mNumHeads * sizeof(half));
|
|
mK_Buffer = (void*)((uint8_t*)buffer_k.first + buffer_k.second);
|
|
auto buffer_v = pool->alloc(mBatchSize * mSeqLenKV * mHeadSizeV * mNumHeads * sizeof(half));
|
|
mV_Buffer = (void*)((uint8_t*)buffer_v.first + buffer_v.second);
|
|
// output size
|
|
auto buffer_acc = pool->alloc(mBatchSize * mSeqLen * mHeadSizeV * mNumHeads * sizeof(float));
|
|
mAcc_Buffer = (void*)((uint8_t*)buffer_acc.first + buffer_acc.second);
|
|
|
|
if(mType == 0) {
|
|
pool->free(buffer_q);
|
|
}
|
|
pool->free(buffer_k);
|
|
pool->free(buffer_v);
|
|
pool->free(buffer_acc);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode FmhaCommonExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("start FmhaCommonExecution onExecute...");
|
|
#endif
|
|
|
|
auto runtime = static_cast<CUDABackend*>(backend())->getCUDARuntime();
|
|
|
|
size_t count = mBatchSize * mSeqLenKV * mHeadSizeV * mNumHeads;
|
|
int block_num = runtime->blocks_num(count);
|
|
int thread_num = runtime->threads_num();
|
|
//printf("type:%d, %p %p %p %p, %d %d %d %d %d\n", mType, inputs[0]->deviceId(), mQ_Buffer, mK_Buffer, mV_Buffer, mHeadSizeV, mNumHeads, mSeqLen, mSeqLenKV, mBatchSize);
|
|
if(mType == 0) {
|
|
SPLIT_FusedQKV<<<block_num, thread_num>>>(count, (const half*)inputs[0]->deviceId(), (half *)mQ_Buffer, (half *)mK_Buffer, (half *)mV_Buffer, mHeadSizeV);
|
|
checkKernelErrors;
|
|
}
|
|
if(mType == 1) {
|
|
mQ_Buffer = (void *)inputs[0]->deviceId();
|
|
SPLIT_FusedKV<<<block_num, thread_num>>>(count, (const half*)inputs[1]->deviceId(), (half *)mK_Buffer, (half *)mV_Buffer, mHeadSizeV);
|
|
checkKernelErrors;
|
|
}
|
|
// Determine kernel configuration based on head size.
|
|
// If head size is less than or equal to 64, each block operates over 64 queries and
|
|
// 64 keys, and partial results can be stored in the register file.
|
|
// If head size is greater than 64, each block operates over 32 queries and 128 keys,
|
|
// and partial results are stored in shared memory.
|
|
int ret = 0;
|
|
if (mHeadSize > 64) {
|
|
static int const kQueriesPerBlock = 32;
|
|
static int const kKeysPerBlock = 128;
|
|
if (mHeadSize <= 128) {
|
|
ret = run_attention<kQueriesPerBlock, kKeysPerBlock, 128>(inputs, outputs);
|
|
} else {
|
|
ret = run_attention<kQueriesPerBlock, kKeysPerBlock, 65536>(inputs, outputs);
|
|
}
|
|
} else {
|
|
static constexpr int kMaxK = 64; // <- Decrease to 32/16 if your problem is smaller
|
|
static int const kQueriesPerBlock = 64;
|
|
static int const kKeysPerBlock = 64;
|
|
ret = run_attention<kQueriesPerBlock, kKeysPerBlock, kMaxK>(inputs, outputs);
|
|
}
|
|
|
|
// printf("fmha shape b:%d s:%d %d h_num:%d h_size:%d, %d\n", mBatchSize, mSeqLen, mSeqLenKV, mNumHeads, mHeadSize, mHeadSizeV);
|
|
checkKernelErrors;
|
|
if(ret != 0) {
|
|
MNN_ERROR("FmhaCommonExecution error\n");
|
|
}
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("end FmhaCommonExecution onExecute...");
|
|
#endif
|
|
return NO_ERROR;
|
|
}
|
|
|
|
} // namespace CUDA
|
|
} // namespace MNN
|