Skip to content

GitLab

"src/targets/vscode:/vscode.git/clone" did not exist on "36eaf9e532828fa014e7ed2eebb40911ed73324b"
Unverified Commit 452822a4 authored by q.yao's avatar q.yao Committed by GitHub
Browse files

Add flashattention2 (#196) 



* first

* fix causal mask

* disable flash attention2 on sm70

* fix 2

* update readme

* clang-format

* disable ft2 on windows

* fix lint

* fix build

* fix build

* fix long kv seq

* fix lint

* sync copy output

---------
Co-authored-by: default avatargrimoire <yaoqian@pjlab.org.cn>
Co-authored-by: default avatarirexyc <irexyc@gmail.com>
parent d4d609bd
Hide whitespace changes
Inline Side-by-side
Showing with 222 additions and 133 deletions
src/turbomind/models/llama/fused_multi_head_attention/llama_flash_attention_kernel.cu
View file @ 452822a4
#include "src/turbomind/models/llama/llama_kernels.h" #include "src/turbomind/models/llama/llama_kernels.h"
#include "src/turbomind/utils/cuda_utils.h" #include "src/turbomind/utils/cuda_utils.h"
#include "42_fused_multi_head_attention/kernel_forward.h" #include "41_fused_multi_head_attention/kernel_forward.h"
#include "mma_accum_lambda_iterator.h"
#include "tile_smem_loader.h"
#include <cuda_fp16.h> #include <cuda_fp16.h>
#include <cutlass/arch/arch.h> #include <cutlass/arch/arch.h>
#include <cutlass/gemm/gemm.h> #include <cutlass/gemm/gemm.h>
...@@ -77,20 +75,10 @@ struct LlamaAttentionKernel: ...@@ -77,20 +75,10 @@ struct LlamaAttentionKernel:
int v_batch_seqs_offset = 0; int v_batch_seqs_offset = 0;
int o_batch_seqs_offset = 0; int o_batch_seqs_offset = 0;
int32_t o_strideM_custom = 0;
int32_t group_size = 1; int32_t group_size = 1;
float scale; float scale;
CUTLASS_HOST_DEVICE int32_t o_strideM() const
{
if (o_strideM_custom == 0)
return BaseParams::head_dim_value;
else
return o_strideM_custom;
}
template<typename ptr_t> template<typename ptr_t>
CUTLASS_DEVICE void CUTLASS_DEVICE void
update_batched_ptr(ptr_t& data_ptr, ptr_t* batch_seq_ptr, int batch_seq_offset, int batch_id, int strideB) update_batched_ptr(ptr_t& data_ptr, ptr_t* batch_seq_ptr, int batch_seq_offset, int batch_id, int strideB)
...@@ -107,8 +95,8 @@ struct LlamaAttentionKernel: ...@@ -107,8 +95,8 @@ struct LlamaAttentionKernel:
auto& query_ptr = BaseParams::query_ptr; auto& query_ptr = BaseParams::query_ptr;
auto& key_ptr = BaseParams::key_ptr; auto& key_ptr = BaseParams::key_ptr;
auto& value_ptr = BaseParams::value_ptr; auto& value_ptr = BaseParams::value_ptr;
auto& cu_seqlens_q_ptr = BaseParams::cu_seqlens_q_ptr; auto& cu_seqlens_q_ptr = BaseParams::seqstart_q_ptr;
auto& cu_seqlens_k_ptr = BaseParams::cu_seqlens_k_ptr; auto& cu_seqlens_k_ptr = BaseParams::seqstart_k_ptr;
auto& output_ptr = BaseParams::output_ptr; auto& output_ptr = BaseParams::output_ptr;
auto& output_accum_ptr = BaseParams::output_accum_ptr; auto& output_accum_ptr = BaseParams::output_accum_ptr;
...@@ -119,22 +107,19 @@ struct LlamaAttentionKernel: ...@@ -119,22 +107,19 @@ struct LlamaAttentionKernel:
auto& num_queries = BaseParams::num_queries; auto& num_queries = BaseParams::num_queries;
auto& num_keys = BaseParams::num_keys; auto& num_keys = BaseParams::num_keys;
auto& causal = BaseParams::causal;
auto& q_strideM = BaseParams::q_strideM; auto& q_strideM = BaseParams::q_strideM;
auto& k_strideM = BaseParams::k_strideM; auto& k_strideM = BaseParams::k_strideM;
auto& v_strideM = BaseParams::v_strideM; auto& v_strideM = BaseParams::v_strideM;
auto& o_strideM = BaseParams::o_strideM;
// Everything below is only used in `advance_to_block` // Everything below is only used in `advance_to_block`
// and shouldn't use registers // and shouldn't use registers
auto& q_strideH = BaseParams::q_strideH; auto& q_strideH = BaseParams::q_strideH;
auto& k_strideH = BaseParams::k_strideH; auto& k_strideH = BaseParams::k_strideH;
auto& v_strideH = BaseParams::v_strideH; auto& v_strideH = BaseParams::v_strideH;
auto& o_strideH = BaseParams::o_strideH;
auto& q_strideB = BaseParams::q_strideB; auto& q_strideB = BaseParams::q_strideB;
auto& k_strideB = BaseParams::k_strideB; auto& k_strideB = BaseParams::k_strideB;
auto& v_strideB = BaseParams::v_strideB; auto& v_strideB = BaseParams::v_strideB;
auto& o_strideB = BaseParams::o_strideB;
auto& num_batches = BaseParams::num_batches; auto& num_batches = BaseParams::num_batches;
auto& num_heads = BaseParams::num_heads; auto& num_heads = BaseParams::num_heads;
...@@ -142,6 +127,8 @@ struct LlamaAttentionKernel: ...@@ -142,6 +127,8 @@ struct LlamaAttentionKernel:
auto head_id = blockIdx.y; auto head_id = blockIdx.y;
auto query_start = blockIdx.x * kQueriesPerBlock; auto query_start = blockIdx.x * kQueriesPerBlock;
auto o_strideB = o_strideM * num_queries;
auto lse_dim = ceil_div((int32_t)num_queries, kAlignLSE) * kAlignLSE; auto lse_dim = ceil_div((int32_t)num_queries, kAlignLSE) * kAlignLSE;
int64_t q_start, k_start; int64_t q_start, k_start;
...@@ -203,10 +190,10 @@ struct LlamaAttentionKernel: ...@@ -203,10 +190,10 @@ struct LlamaAttentionKernel:
query_ptr += (qq_start + query_start) * q_strideM + head_id * q_strideH; query_ptr += (qq_start + query_start) * q_strideM + head_id * q_strideH;
key_ptr += k_start * k_strideM + int64_t(head_id / group_size) * k_strideH; key_ptr += k_start * k_strideM + int64_t(head_id / group_size) * k_strideH;
value_ptr += k_start * v_strideM + int64_t(head_id / group_size) * v_strideH; value_ptr += k_start * v_strideM + int64_t(head_id / group_size) * v_strideH;
output_ptr += int64_t(qo_start + query_start) * o_strideM() + head_id * o_strideH; output_ptr += int64_t(qo_start + query_start) * o_strideM + head_id * head_dim_value;
if (output_accum_ptr != nullptr) { if (output_accum_ptr != nullptr) {
output_accum_ptr += int64_t(query_start) * o_strideM() + head_id * o_strideH; output_accum_ptr += int64_t(query_start) * o_strideM + head_id * head_dim_value;
} }
else { else {
// Accumulate directly in the destination buffer (eg for f32) // Accumulate directly in the destination buffer (eg for f32)
...@@ -218,9 +205,6 @@ struct LlamaAttentionKernel: ...@@ -218,9 +205,6 @@ struct LlamaAttentionKernel:
} }
num_queries -= query_start; num_queries -= query_start;
if (causal) {
num_keys = cutlass::fast_min(int32_t(query_start + kQueriesPerBlock), num_keys);
}
num_batches = 0; // no longer used after num_batches = 0; // no longer used after
// Make sure the compiler knows these variables are the same on all // Make sure the compiler knows these variables are the same on all
...@@ -328,7 +312,7 @@ struct LlamaAttentionKernel: ...@@ -328,7 +312,7 @@ struct LlamaAttentionKernel:
auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator {
using OutputTileIterator = typename MM1::OutputTileIterator; using OutputTileIterator = typename MM1::OutputTileIterator;
return OutputTileIterator(typename OutputTileIterator::Params{(int32_t)p.o_strideM()}, return OutputTileIterator(typename OutputTileIterator::Params{(int32_t)p.o_strideM},
p.output_ptr, p.output_ptr,
typename OutputTileIterator::TensorCoord{p.num_queries, p.head_dim_value}, typename OutputTileIterator::TensorCoord{p.num_queries, p.head_dim_value},
thread_id(), thread_id(),
...@@ -338,7 +322,7 @@ struct LlamaAttentionKernel: ...@@ -338,7 +322,7 @@ struct LlamaAttentionKernel:
auto createOutputAccumIter = [&](int col) -> typename MM1::OutputTileIteratorAccum { auto createOutputAccumIter = [&](int col) -> typename MM1::OutputTileIteratorAccum {
using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum; using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum;
return OutputTileIteratorAccum( return OutputTileIteratorAccum(
typename OutputTileIteratorAccum::Params{(int32_t)p.o_strideM()}, typename OutputTileIteratorAccum::Params{(int32_t)p.o_strideM},
p.output_accum_ptr, p.output_accum_ptr,
typename OutputTileIteratorAccum::TensorCoord{p.num_queries, p.head_dim_value}, typename OutputTileIteratorAccum::TensorCoord{p.num_queries, p.head_dim_value},
thread_id(), thread_id(),
...@@ -453,28 +437,27 @@ struct LlamaAttentionKernel: ...@@ -453,28 +437,27 @@ struct LlamaAttentionKernel:
[&](int accum_m) {}); [&](int accum_m) {});
} }
DISPATCH_BOOL( DISPATCH_BOOL(iter_key_start == 0, kIsFirst, ([&] {
iter_key_start == 0, kIsFirst, ([&] { DISPATCH_BOOL(
DISPATCH_BOOL( p.num_keys - iter_key_start >= kKeysPerBlock, kFullColumns, ([&] {
p.num_keys - iter_key_start >= kKeysPerBlock, kFullColumns, ([&] { // Update `mi` from accum stored in registers
// Update `mi` from accum stored in registers // Also updates `accum` with accum[i] <-
// Also updates `accum` with accum[i] <- // exp(accum[i] * scale
// exp(accum[i] * scale // - mi)
// - mi) Base::iterative_softmax<MM0::Mma::Operator::IteratorC, kFullColumns, kIsFirst>(
MM0::ScalingCoefsUpdater::update<kQueriesPerBlock, kFullColumns, kIsFirst, kKeepOutputInRF>( accum_o,
accum_o, accum,
accum, mi,
mi, m_prime,
m_prime, s_prime,
s_prime, lane_id(),
lane_id(), thread_id(),
thread_id(), warp_id(),
warp_id(), p.num_keys - iter_key_start,
p.num_keys - iter_key_start, iteratorC_tile_offset,
iteratorC_tile_offset, kSupportsBias ? 1.0f : p.scale);
kSupportsBias ? 1.0f : p.scale); }));
})); }));
}));
// Output results to shared-memory // Output results to shared-memory
int warp_idx_mn_0 = my_warp_id % (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN); int warp_idx_mn_0 = my_warp_id % (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN);
...@@ -651,13 +634,13 @@ struct LlamaAttentionKernel: ...@@ -651,13 +634,13 @@ struct LlamaAttentionKernel:
}; };
template<typename T, typename Attention> template<typename T, typename Attention>
void invokeFlashAttention_impl(int batch_size, void invokeFlashAttention_impl(int batch_size,
int head_num, int head_num,
int key_len, int key_len,
int seq_len, int seq_len,
int size_per_head, int size_per_head,
typename FlashAttentionOp<T>::Params& attention_params, typename FlashAttentionOpImpl<T, 1>::Params& attention_params,
cudaStream_t st) cudaStream_t st)
{ {
T* out_ptr = attention_params.attn_out; T* out_ptr = attention_params.attn_out;
T* query_ptr = attention_params.query; T* query_ptr = attention_params.query;
...@@ -685,11 +668,11 @@ void invokeFlashAttention_impl(int batch_size, ...@@ -685,11 +668,11 @@ void invokeFlashAttention_impl(int batch_size,
// fill param // fill param
typename Attention::Params params{}; typename Attention::Params params{};
{ {
params.query_ptr = (scalar_t*)(query_ptr); params.query_ptr = (scalar_t*)(query_ptr);
params.key_ptr = (scalar_t*)(key_ptr); params.key_ptr = (scalar_t*)(key_ptr);
params.value_ptr = (scalar_t*)(value_ptr); params.value_ptr = (scalar_t*)(value_ptr);
params.attn_bias_ptr = (scalar_t*)(mask_ptr); params.attn_bias_ptr = (scalar_t*)(mask_ptr);
params.cu_seqlens_q_ptr = cu_seqlens_q_ptr; params.seqstart_q_ptr = cu_seqlens_q_ptr;
params.output_ptr = (scalar_t*)(out_ptr); params.output_ptr = (scalar_t*)(out_ptr);
params.output_accum_ptr = kNeedsOutputAccumulatorBuffer ? output_accum_ptr : nullptr; params.output_accum_ptr = kNeedsOutputAccumulatorBuffer ? output_accum_ptr : nullptr;
...@@ -725,9 +708,7 @@ void invokeFlashAttention_impl(int batch_size, ...@@ -725,9 +708,7 @@ void invokeFlashAttention_impl(int batch_size,
params.v_batch_seqs_ptr = (scalar_t**)layout_v.batch_seqs; params.v_batch_seqs_ptr = (scalar_t**)layout_v.batch_seqs;
params.v_batch_seqs_offset = layout_v.batch_seqs_offset; params.v_batch_seqs_offset = layout_v.batch_seqs_offset;
params.o_strideH = layout_o.stride_head; params.o_strideM = layout_o.stride_seq;
params.o_strideM_custom = layout_o.stride_seq;
params.o_strideB = layout_o.stride_batch;
params.o_use_seqlens = layout_o.use_seqlens; params.o_use_seqlens = layout_o.use_seqlens;
params.o_batch_seqs_ptr = (scalar_t**)layout_o.batch_seqs; params.o_batch_seqs_ptr = (scalar_t**)layout_o.batch_seqs;
params.o_batch_seqs_offset = layout_o.batch_seqs_offset; params.o_batch_seqs_offset = layout_o.batch_seqs_offset;
...@@ -784,14 +765,14 @@ bool get_needs_accum_buffer() ...@@ -784,14 +765,14 @@ bool get_needs_accum_buffer()
} }
template<typename T, int kQueriesPerBlock, int kKeysPerBlock> template<typename T, int kQueriesPerBlock, int kKeysPerBlock>
void invoke_attention_impl(bool single_val_iteration, void invoke_attention_impl(bool single_val_iteration,
int batch_size, int batch_size,
int head_num, int head_num,
int key_len, int key_len,
int seq_len, int seq_len,
int size_per_head, int size_per_head,
typename FlashAttentionOp<T>::Params& params, typename FlashAttentionOpImpl<T, 1>::Params& params,
cudaStream_t st) cudaStream_t st)
{ {
using scalar_t = using scalar_t =
typename std::conditional_t<std::is_same<half, typename std::decay<T>::type>::value, cutlass::half_t, T>; typename std::conditional_t<std::is_same<half, typename std::decay<T>::type>::value, cutlass::half_t, T>;
...@@ -830,18 +811,34 @@ void invoke_attention_impl(bool single_val_iter ...@@ -830,18 +811,34 @@ void invoke_attention_impl(bool single_val_iter
} }
template<typename T> template<typename T>
class FlashAttentionOp<T>::impl { class FlashAttentionOpImpl<T, 1> {
public:
using AttentionLayout = BaseAttentionLayout<T>;
using Params = BaseAttentionParams<T>;
public:
FlashAttentionOpImpl(int batch_size, int head_num, int key_len, int seq_len, int size_per_head);
~FlashAttentionOpImpl();
int get_workspace_size() const;
void operator()(Params& params, cudaStream_t st) const;
private:
class impl;
std::unique_ptr<impl> pimpl;
};
template<typename T>
class FlashAttentionOpImpl<T, 1>::impl {
private: private:
static constexpr int kQueriesPerBlock = 32; static constexpr int kQueriesPerBlock = 32;
static constexpr int kKeysPerBlock = 128; static constexpr int kKeysPerBlock = 128;
using ArchTag = cutlass::arch::Sm80;
using scalar_t = using scalar_t =
typename std::conditional_t<std::is_same<half, typename std::decay<T>::type>::value, cutlass::half_t, T>; typename std::conditional_t<std::is_same<half, typename std::decay<T>::type>::value, cutlass::half_t, T>;
using SingleValueAttention = LlamaAttentionKernel<scalar_t, ArchTag, kQueriesPerBlock, kKeysPerBlock, true>; using Params = typename FlashAttentionOpImpl<T, 1>::Params;
using MultiValueAttention = LlamaAttentionKernel<scalar_t, ArchTag, kQueriesPerBlock, kKeysPerBlock, false>;
using AttentionLayout = typename FlashAttentionOp<T>::AttentionLayout;
using Params = typename FlashAttentionOp<T>::Params;
int batch_size_; int batch_size_;
int head_num_; int head_num_;
...@@ -887,29 +884,30 @@ public: ...@@ -887,29 +884,30 @@ public:
}; };
template<typename T> template<typename T>
FlashAttentionOp<T>::FlashAttentionOp(int batch_size, int head_num, int key_len, int seq_len, int size_per_head): FlashAttentionOpImpl<T, 1>::FlashAttentionOpImpl(
pimpl{std::make_unique<FlashAttentionOp<T>::impl>(batch_size, head_num, key_len, seq_len, size_per_head)} int batch_size, int head_num, int key_len, int seq_len, int size_per_head):
pimpl{std::make_unique<FlashAttentionOpImpl<T, 1>::impl>(batch_size, head_num, key_len, seq_len, size_per_head)}
{ {
} }
template<typename T> template<typename T>
FlashAttentionOp<T>::~FlashAttentionOp() FlashAttentionOpImpl<T, 1>::~FlashAttentionOpImpl()
{ {
} }
template<typename T> template<typename T>
int FlashAttentionOp<T>::get_workspace_size() const int FlashAttentionOpImpl<T, 1>::get_workspace_size() const
{ {
return pimpl->get_workspace_size(); return pimpl->get_workspace_size();
} }
template<typename T> template<typename T>
void FlashAttentionOp<T>::operator()(Params& params, cudaStream_t st) const void FlashAttentionOpImpl<T, 1>::operator()(Params& params, cudaStream_t st) const
{ {
pimpl->operator()(params, st); pimpl->operator()(params, st);
} }
template class FlashAttentionOp<float>; template class FlashAttentionOpImpl<float, 1>;
template class FlashAttentionOp<half>; template class FlashAttentionOpImpl<half, 1>;
} // namespace turbomind } // namespace turbomind
src/turbomind/models/llama/llama_kernels.cu
View file @ 452822a4
...@@ -329,11 +329,11 @@ static inline __device__ uint32_t float4_to_char4(float x, float y, float z, flo ...@@ -329,11 +329,11 @@ static inline __device__ uint32_t float4_to_char4(float x, float y, float z, flo
asm volatile("cvt.pack.sat.s8.s32.b32 %0, %1, %2, %0;\n" : "+r"(dst) : "r"(b), "r"(a)); asm volatile("cvt.pack.sat.s8.s32.b32 %0, %1, %2, %0;\n" : "+r"(dst) : "r"(b), "r"(a));
#else #else
char4 tmp; char4 tmp;
tmp.x = x; tmp.x = x;
tmp.y = y; tmp.y = y;
tmp.z = z; tmp.z = z;
tmp.w = w; tmp.w = w;
dst = reinterpret_cast<const uint32_t&>(tmp); dst = reinterpret_cast<const uint32_t&>(tmp);
#endif #endif
return dst; return dst;
} }
...@@ -724,4 +724,61 @@ void invokeGatherOutput(int* output_ids, ...@@ -724,4 +724,61 @@ void invokeGatherOutput(int* output_ids,
output_ids, ids, context_length, max_context_len, max_gen_step, max_output_len, batch_size); output_ids, ids, context_length, max_context_len, max_gen_step, max_output_len, batch_size);
} }
#define VERSION_SWITCH(VERSION, CONST_NAME, ...) \
[&] { \
if (VERSION == 2) { \
constexpr static int CONST_NAME = 2; \
return __VA_ARGS__(); \
} \
else { \
constexpr static int CONST_NAME = 1; \
return __VA_ARGS__(); \
} \
}()
template<typename T>
FlashAttentionOp<T>::FlashAttentionOp(int batch_size, int head_num, int key_len, int seq_len, int size_per_head):
batch_size_(batch_size), head_num_(head_num), key_len_(key_len), seq_len_(seq_len), size_per_head_(size_per_head)
{
#ifdef _MSC_VER
op_version_ = 1;
#else
op_version_ = std::is_same<half, typename std::decay<T>::type>::value ? 2 : 1;
if (op_version_ == 2 && getSMVersion() < 80) {
op_version_ = 1;
}
#endif
}
template<typename T>
int FlashAttentionOp<T>::get_workspace_size() const
{
#ifdef _MSC_VER
FlashAttentionOpImpl<T, 1> attention_op(batch_size_, head_num_, key_len_, seq_len_, size_per_head_);
return attention_op.get_workspace_size();
#else
return VERSION_SWITCH(op_version_, OP_VERSION, [&]() {
FlashAttentionOpImpl<T, OP_VERSION> attention_op(batch_size_, head_num_, key_len_, seq_len_, size_per_head_);
return attention_op.get_workspace_size();
});
#endif
}
template<typename T>
void FlashAttentionOp<T>::operator()(Params& params, cudaStream_t st) const
{
#ifdef _MSC_VER
FlashAttentionOpImpl<T, 1> attention_op(batch_size_, head_num_, key_len_, seq_len_, size_per_head_);
return attention_op(params, st);
#else
return VERSION_SWITCH(op_version_, OP_VERSION, [&]() {
FlashAttentionOpImpl<T, OP_VERSION> attention_op(batch_size_, head_num_, key_len_, seq_len_, size_per_head_);
return attention_op(params, st);
});
#endif
}
template class FlashAttentionOp<float>;
template class FlashAttentionOp<half>;
} // namespace turbomind } // namespace turbomind
src/turbomind/models/llama/llama_kernels.h
View file @ 452822a4
...@@ -79,36 +79,41 @@ void invokeGatherOutput(int* output_ids, ...@@ -79,36 +79,41 @@ void invokeGatherOutput(int* output_ids,
void invokeMyCopyInt(int* dst, const int* src, size_t count, cudaStream_t st); void invokeMyCopyInt(int* dst, const int* src, size_t count, cudaStream_t st);
template<typename T> template<typename T>
class FlashAttentionOp { struct BaseAttentionLayout {
int stride_batch;
int stride_seq;
int stride_head;
bool use_seqlens = false;
int batch_seqs_offset = 0;
T** batch_seqs = nullptr;
};
template<typename T>
struct BaseAttentionParams {
T* attn_out;
T* query;
T* key;
T* val;
T* mask;
float* out_accum = nullptr;
int* cu_seqlens_q = nullptr;
int* cu_seqlens_k = nullptr;
size_t group_size = 1;
BaseAttentionLayout<T> layout_q;
BaseAttentionLayout<T> layout_k;
BaseAttentionLayout<T> layout_v;
BaseAttentionLayout<T> layout_o;
};
template<typename T, int version>
class FlashAttentionOpImpl {
public: public:
struct AttentionLayout { using AttentionLayout = BaseAttentionLayout<T>;
int stride_batch; using Params = BaseAttentionParams<T>;
int stride_seq;
int stride_head;
bool use_seqlens = false;
int batch_seqs_offset = 0;
T** batch_seqs = nullptr;
};
struct Params {
T* attn_out;
T* query;
T* key;
T* val;
T* mask;
float* out_accum = nullptr;
int* cu_seqlens_q = nullptr;
int* cu_seqlens_k = nullptr;
size_t group_size = 1;
AttentionLayout layout_q;
AttentionLayout layout_k;
AttentionLayout layout_v;
AttentionLayout layout_o;
};
public: public:
FlashAttentionOp(int batch_size, int head_num, int key_len, int seq_len, int size_per_head); FlashAttentionOpImpl(int batch_size, int head_num, int key_len, int seq_len, int size_per_head);
~FlashAttentionOp(); ~FlashAttentionOpImpl();
int get_workspace_size() const; int get_workspace_size() const;
...@@ -119,6 +124,28 @@ private: ...@@ -119,6 +124,28 @@ private:
std::unique_ptr<impl> pimpl; std::unique_ptr<impl> pimpl;
}; };
template<typename T>
class FlashAttentionOp {
public:
using AttentionLayout = BaseAttentionLayout<T>;
using Params = BaseAttentionParams<T>;
public:
FlashAttentionOp(int batch_size, int head_num, int key_len, int seq_len, int size_per_head);
int get_workspace_size() const;
void operator()(Params& params, cudaStream_t st) const;
private:
int batch_size_;
int head_num_;
int key_len_;
int seq_len_;
int size_per_head_;
int op_version_;
};
template<typename T> template<typename T>
inline void dump(const T* x, int size, cudaStream_t st, const char* msg, bool full = false) inline void dump(const T* x, int size, cudaStream_t st, const char* msg, bool full = false)
{ {
......
tests/csrc/unittests/test_context_attention_layer.cu
View file @ 452822a4
...@@ -10,6 +10,8 @@ ...@@ -10,6 +10,8 @@
#include <thrust/host_vector.h> #include <thrust/host_vector.h>
#include <thrust/transform.h> #include <thrust/transform.h>
#undef TORCH_CUDA
#include "src/turbomind/kernels/bert_preprocess_kernels.h" #include "src/turbomind/kernels/bert_preprocess_kernels.h"
#include "src/turbomind/kernels/unfused_attention_kernels.h" #include "src/turbomind/kernels/unfused_attention_kernels.h"
#include "src/turbomind/models/llama/llama_kernels.h" #include "src/turbomind/models/llama/llama_kernels.h"
...@@ -202,15 +204,6 @@ void naive_mha(scalar_t* out_ptr, ...@@ -202,15 +204,6 @@ void naive_mha(scalar_t* out_ptr,
pad_out(out_ptr, cu_seqlens, batch_size, head_num * seq_len * size_per_head, head_num * size_per_head, stream); pad_out(out_ptr, cu_seqlens, batch_size, head_num * seq_len * size_per_head, head_num * size_per_head, stream);
} }
template<typename scalar_t>
struct UpdateMask {
UpdateMask() {}
__host__ __device__ scalar_t operator()(const scalar_t& x) const
{
return x > scalar_t(0.0f) ? scalar_t(1.0f) : scalar_t(0.0f);
}
};
static const char* usage = "Usage: %s <batch-size> <num-heads> <key-len> <query-len> <size-per-head>\n" static const char* usage = "Usage: %s <batch-size> <num-heads> <key-len> <query-len> <size-per-head>\n"
"Example: $test_context_attention_layer 2, 8, 1024, 512, 128\n"; "Example: $test_context_attention_layer 2, 8, 1024, 512, 128\n";
...@@ -290,13 +283,14 @@ int main(int argc, const char* argv[]) ...@@ -290,13 +283,14 @@ int main(int argc, const char* argv[])
cudaRandomUniform<scalar_t>(key_ptr, batch_size * num_heads * key_len * size_per_head); cudaRandomUniform<scalar_t>(key_ptr, batch_size * num_heads * key_len * size_per_head);
cudaRandomUniform<scalar_t>(val_ptr, batch_size * num_heads * key_len * size_per_head); cudaRandomUniform<scalar_t>(val_ptr, batch_size * num_heads * key_len * size_per_head);
cudaRandomUniform<scalar_t>(mask_ptr, batch_size * seq_len * key_len); cudaRandomUniform<scalar_t>(mask_ptr, batch_size * seq_len * key_len);
thrust::transform(
thrust::device, mask_ptr, mask_ptr + batch_size * seq_len * key_len, mask_ptr, UpdateMask<scalar_t>());
// create random length for batch // create random length for batch
std::uniform_int_distribution<int> dist{seq_len / 2, seq_len}; std::uniform_int_distribution<int> dist{seq_len / 2, seq_len};
auto gen = [&dist, &mersenne_engine]() { return dist(mersenne_engine); }; auto gen = [&dist, &mersenne_engine]() { return dist(mersenne_engine); };
std::generate(begin(input_lengths_host), end(input_lengths_host), gen); std::generate(begin(input_lengths_host), end(input_lengths_host), gen);
// for(int batch_id=0;batch_id<batch_size;++batch_id){
// input_lengths_host[batch_id] = seq_len;
// }
thrust::copy(input_lengths_host.begin(), input_lengths_host.end(), input_lengths.begin()); thrust::copy(input_lengths_host.begin(), input_lengths_host.end(), input_lengths.begin());
size_t h_token_num = 0; size_t h_token_num = 0;
size_t* h_pinned_token_num; size_t* h_pinned_token_num;
...@@ -312,6 +306,12 @@ int main(int argc, const char* argv[]) ...@@ -312,6 +306,12 @@ int main(int argc, const char* argv[])
stream); stream);
cudaFreeHost((void*)h_pinned_token_num); cudaFreeHost((void*)h_pinned_token_num);
int* k_lens = (int*)allocator.malloc(batch_size * sizeof(int));
deviceFill(k_lens, batch_size, key_len, stream);
invokeCreateCausalMasks(mask_ptr, input_lengths_ptr, k_lens, seq_len, key_len, batch_size, stream);
// deviceFill(mask_ptr, batch_size*key_len*seq_len, scalar_t(1), stream);
// compute gt // compute gt
naive_mha<scalar_t>(expect_out_ptr, naive_mha<scalar_t>(expect_out_ptr,
query_ptr, query_ptr,
...@@ -334,7 +334,12 @@ int main(int argc, const char* argv[]) ...@@ -334,7 +334,12 @@ int main(int argc, const char* argv[])
&cublas_wrapper); &cublas_wrapper);
// compute actual // compute actual
using AttentionOp = FlashAttentionOp<scalar_t>; #ifdef _MSC_VER
static constexpr int FMHA_VERSION = 1;
#else
static constexpr int FMHA_VERSION = 2;
#endif
using AttentionOp = FlashAttentionOpImpl<scalar_t, FMHA_VERSION>;
using Layout = typename AttentionOp::AttentionLayout; using Layout = typename AttentionOp::AttentionLayout;
Layout layout_q{num_heads * seq_len * size_per_head, size_per_head, seq_len * size_per_head}; Layout layout_q{num_heads * seq_len * size_per_head, size_per_head, seq_len * size_per_head};
Layout layout_k{num_heads * key_len * size_per_head, size_per_head, key_len * size_per_head}; Layout layout_k{num_heads * key_len * size_per_head, size_per_head, key_len * size_per_head};
...@@ -359,11 +364,13 @@ int main(int argc, const char* argv[]) ...@@ -359,11 +364,13 @@ int main(int argc, const char* argv[])
flash_attention(attn_params, stream); flash_attention(attn_params, stream);
sync_check_cuda_error(); sync_check_cuda_error();
// int num_rows = 8; int num_rows = 8;
// printf("expect:\n"); // printf("query:\n");
// printMatrix(expect_out_ptr, num_rows, size_per_head, size_per_head, true); // printMatrix(query_ptr, num_rows, 8, size_per_head, true);
// printf("actual:\n"); printf("expect:\n");
// printMatrix(actual_out_ptr, num_rows, size_per_head, size_per_head, true); printMatrix(expect_out_ptr, num_rows, 8, size_per_head, true);
printf("actual:\n");
printMatrix(actual_out_ptr, num_rows, 8, size_per_head, true);
checkResult( checkResult(
"all close:", actual_out_ptr, expect_out_ptr, batch_size * num_heads * seq_len * size_per_head, true, true); "all close:", actual_out_ptr, expect_out_ptr, batch_size * num_heads * seq_len * size_per_head, true, true);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment