Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: 开始实现 attention #89

Draft
wants to merge 17 commits into
base: dev
Choose a base branch
from
Draft

feat: 开始实现 attention #89

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
cbe7c78
feat: 开始实现 attention
YdrMaster Feb 2, 2024
6d4465c
build(kernel): try to use local cccl instead of system
YdrMaster Feb 2, 2024
e54c7db
refactor(kernel): 实现一种不依赖模板参数的 BlockReduce 并用于 softmax
YdrMaster Feb 2, 2024
8e1564a
feat(kernel): 确认供 attention 调用的 softmax 接口
YdrMaster Feb 2, 2024
ccf293f
feat(kernel): 开始实现 softmax
YdrMaster Feb 2, 2024
20a34ac
feat(kernel): 完成不带 kv cache 的简单 attention
YdrMaster Feb 2, 2024
7286459
fix: 整理和改正 attention 构造问题
YdrMaster Feb 2, 2024
d7bbd3b
fix(kernel): 设置 sharedMemory
YdrMaster Feb 2, 2024
9f736e9
test(llm): 添加 attention 前端单测
YdrMaster Feb 2, 2024
4d47e5d
feat(kernel): 封装 attLen 计算
YdrMaster Feb 2, 2024
6211b92
test(kernel): 添加 attention 单测
YdrMaster Feb 2, 2024
1705126
fix(kernel): 解决 attention 访存错误的问题
YdrMaster Feb 2, 2024
f5d9756
fix(kernel): 仍然使用 cub::BlockReduce 并改正 Attention
YdrMaster Feb 2, 2024
e16c679
feat(kernel): 支持 kv cache attention
YdrMaster Feb 2, 2024
bfa8e9f
fix (kernel): 修复attention算子中的concat
PanZezhong1725 Feb 21, 2024
a686d2f
temp
YdrMaster Feb 21, 2024
0d1fce4
temp
YdrMaster Feb 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
feat(kernel): 封装 attLen 计算 
Signed-off-by: YdrMaster <ydrml@hotmail.com>
  • Loading branch information
@YdrMaster
YdrMaster committed Feb 19, 2024
commit 4d47e5d7ae4b2811101490084ac7904cbaa4841b
3 changes: 3 additions & 0 deletions src/04kernel/include/kernel/attributes/attention_info.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@ namespace refactor::kernel {
DataType dataType;
dim_t batch, nHead, nKVHead, seqLen, headDim, cacheLen;
bool concatCache, resetCache;

dim_t attLen(dim_t pastSeqLen) const noexcept;
size_t attSize(dim_t pastSeqLen) const noexcept;
};

}// namespace refactor::kernel
Expand Down
13 changes: 13 additions & 0 deletions src/04kernel/src/attributes/attention_info.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
#include "kernel/attributes/attention_info.h"

namespace refactor::kernel {

dim_t AttentionInfo::attLen(dim_t pastSeqLen) const noexcept {
return pastSeqLen + seqLen;
}

size_t AttentionInfo::attSize(dim_t pastSeqLen) const noexcept {
return batch * nHead * seqLen * attLen(pastSeqLen) * dataType.size();
}

}// namespace refactor::kernel
17 changes: 9 additions & 8 deletions src/04kernel/src/kernels/attention/cuda_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ namespace refactor::kernel {
MatMulDescriptor mul;
MatrixDescriptor q, k, v, att;
cublasLtMatmulAlgo_t algoQK, algoAV;
size_t attSize, workspaceSizeQK, workspaceSizeAV;
size_t workspaceSizeQK, workspaceSizeAV;

Descriptors(CublasLtContext const &context,
AttentionInfo info)
Expand Down Expand Up @@ -112,8 +112,7 @@ namespace refactor::kernel {
.order = ROW_MAJOR,
.batchCount = static_cast<int32_t>(info.batch * info.nHead),
.batchStride = static_cast<int64_t>(info.seqLen * info.seqLen),
}),
attSize(info.batch * info.nHead * info.seqLen * info.seqLen * info.dataType.size()) {
}) {
auto [algoQK_, workspaceSizeQK_] = tune(context.handle, mul, q, k, att);
auto [algoAV_, workspaceSizeAV_] = tune(context.handle, mul, att, v, q);
algoQK = algoQK_;
Expand All @@ -125,7 +124,7 @@ namespace refactor::kernel {

auto const &context = *res.fetchOrStore<CublasLtContext>();
auto d = std::make_shared<Descriptors>(context, info);
auto workspaceSize = d->attSize;
auto workspaceSize = info.attSize(0);
workspaceSize = hardware::alignBytes(workspaceSize, 256);
workspaceSize += d->workspaceSizeQK;
workspaceSize += d->workspaceSizeAV;
Expand All @@ -139,7 +138,7 @@ namespace refactor::kernel {
auto v = inputs[2];
auto o = outputs[0];
auto att = reinterpret_cast<half *>(workspace);
auto workspaceQK = reinterpret_cast<uint8_t *>(workspace) + hardware::alignBytes(d->attSize, 256);
auto workspaceQK = reinterpret_cast<uint8_t *>(workspace) + hardware::alignBytes(info.attSize(0), 256);
auto workspaceAV = workspaceQK + hardware::alignBytes(d->workspaceSizeQK, 256);
{
half alpha = rsqrtf(info.headDim), beta = 0;
Expand All @@ -155,10 +154,12 @@ namespace refactor::kernel {
workspaceQK, d->workspaceSizeQK,
cudaStreamLegacy);
}
auto attLen = info.attLen(0);
auto bufLen = attLen;
softmax<<<dim3(info.batch * info.nHead, info.seqLen),
info.seqLen,
info.seqLen * sizeof(float)>>>(
att, causualMask, info.seqLen, info.seqLen);
std::min(1024u, attLen),
attLen * sizeof(float)>>>(
att, causualMask, attLen, bufLen);
{
half alpha = 1, beta = 0;
cublasLtMatmul(
Expand Down