wehub-resource-sync
829 lines
30 KiB
Plaintext
829 lines
30 KiB
Plaintext
// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/phi/kernels/flash_attn_kernel.h"
|
|
|
|
#include <cstddef>
|
|
#include "glog/logging.h" // For VLOG()
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/platform/device_context.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/core/utils/data_type.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
|
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
|
|
#include "paddle/phi/kernels/slice_kernel.h"
|
|
#include "paddle/utils/none.h"
|
|
|
|
#ifdef PADDLE_WITH_FLASHATTN_V3
|
|
#include "paddle/phi/kernels/gpu/flash_attn_v3_kernel.h"
|
|
#endif
|
|
|
|
COMMON_DECLARE_int32(flash_attn_version);
|
|
COMMON_DECLARE_bool(cudnn_deterministic);
|
|
|
|
namespace phi {
|
|
template <typename OutT>
|
|
struct ZeroFunctor {
|
|
__device__ __forceinline__ OutT operator()() const {
|
|
return static_cast<OutT>(0);
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnUnpaddedBaseKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& cu_seqlens_q,
|
|
const DenseTensor& cu_seqlens_k,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const Scalar& max_seqlen_q_,
|
|
const Scalar& max_seqlen_k_,
|
|
float scale,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset,
|
|
bool varlen_padded) {
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
if (!out->IsInitialized()) dev_ctx.template Alloc<T>(out);
|
|
if (varlen_padded) {
|
|
std::vector<const DenseTensor*> inputs{};
|
|
std::vector<DenseTensor*> outputs{out};
|
|
|
|
funcs::ElementwiseKernel<T>(dev_ctx, inputs, &outputs, ZeroFunctor<T>());
|
|
}
|
|
#ifdef PADDLE_WITH_HIP
|
|
hipStream_t stream = dev_ctx.stream();
|
|
#else
|
|
cudaStream_t stream = dev_ctx.stream();
|
|
#endif
|
|
|
|
// q, k, v [total_q/k/v, num_heads, head_dim]
|
|
auto dims = q.dims();
|
|
PADDLE_ENFORCE_EQ(
|
|
dims.size(),
|
|
3,
|
|
common::errors::InvalidArgument("flash_attn_raw receive input with dim "
|
|
"[total_seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(
|
|
k.dims().size(),
|
|
3,
|
|
common::errors::InvalidArgument("flash_attn_raw receive input with dim "
|
|
"[total_seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(
|
|
v.dims().size(),
|
|
3,
|
|
common::errors::InvalidArgument("flash_attn_raw receive input with dim "
|
|
"[total_seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(
|
|
out->dims().size(),
|
|
3,
|
|
common::errors::InvalidArgument("flash_attn_raw receive input with dim "
|
|
"[total_seq_len, num_heads, head_dim]"));
|
|
|
|
const int64_t batch_size = cu_seqlens_q.numel() - 1;
|
|
const int64_t num_heads = dims[1];
|
|
const int64_t head_size = dims[2];
|
|
const int64_t num_heads_k = k.dims()[1];
|
|
const int64_t total_q = dims[0];
|
|
|
|
// TODO(umiswing): add shape check
|
|
|
|
int64_t max_seqlen_q = max_seqlen_q_.to<int64_t>();
|
|
int64_t max_seqlen_k = max_seqlen_k_.to<int64_t>();
|
|
|
|
FlashAttnFwdParamsV2<T> params =
|
|
FlashAttnFwdParamsV2<T>(dev_ctx,
|
|
/*version=*/2,
|
|
batch_size,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
dropout,
|
|
scale,
|
|
causal,
|
|
return_softmax,
|
|
q.dtype(),
|
|
is_test,
|
|
rng_name,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
nullptr, // startend_row_indices
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset,
|
|
/*_unpadded_lse*/ true,
|
|
/*_total_q*/ total_q);
|
|
|
|
VLOG(10) << "FlashAttn fwd seed: " << params.seed
|
|
<< ", offset: " << params.offset;
|
|
bool succ = dynload::flash_attn_varlen_fwd(
|
|
q.data(),
|
|
k.data(),
|
|
v.data(),
|
|
cu_seqlens_q.data<int32_t>(),
|
|
cu_seqlens_k.data<int32_t>(),
|
|
params.rng_state.data(),
|
|
out->data(),
|
|
params.return_softmax ? softmax->data() : nullptr,
|
|
softmax_lse->data(),
|
|
params.batch_size,
|
|
params.max_seqlen_q,
|
|
params.max_seqlen_k,
|
|
params.seqlen_q_rounded,
|
|
params.seqlen_k_rounded,
|
|
params.num_heads,
|
|
params.num_heads_k,
|
|
params.head_size,
|
|
params.head_size_rounded,
|
|
params.dropout,
|
|
params.softmax_scale,
|
|
1.0f / params.softmax_scale,
|
|
params.causal,
|
|
params.return_softmax,
|
|
params.is_bf16,
|
|
stream,
|
|
params.seed,
|
|
params.offset,
|
|
params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
|
|
params.attn_mask_tensor ? params.mask_dims.data() : nullptr,
|
|
q.strides()[0],
|
|
k.strides()[0],
|
|
v.strides()[0],
|
|
q.strides()[1],
|
|
k.strides()[1],
|
|
v.strides()[1],
|
|
out->strides()[0],
|
|
out->strides()[1],
|
|
max_seqlen_q * q.strides()[0],
|
|
max_seqlen_k * k.strides()[0],
|
|
max_seqlen_k * v.strides()[0],
|
|
max_seqlen_q * out->strides()[0],
|
|
#ifdef PADDLE_WITH_CUDA
|
|
varlen_padded,
|
|
params.total_q
|
|
#else
|
|
varlen_padded
|
|
#endif
|
|
);
|
|
CheckFlashAttnStatus(succ);
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnUnpaddedKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& cu_seqlens_q,
|
|
const DenseTensor& cu_seqlens_k,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const Scalar& max_seqlen_q,
|
|
const Scalar& max_seqlen_k,
|
|
float scale,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
FlashAttnUnpaddedBaseKernel<T>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
is_test,
|
|
rng_name,
|
|
out,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset,
|
|
false /*varlen_padded*/);
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
|
|
static void sliceFlattenView(const DenseTensor& in,
|
|
DenseTensor* out,
|
|
int axis,
|
|
int64_t offset,
|
|
int64_t sliceLength) {
|
|
PADDLE_ENFORCE_LT(
|
|
axis,
|
|
in.dims().size(),
|
|
common::errors::InvalidArgument("sliceView receive axis out of bound"));
|
|
std::array<int64_t, DDim::kMaxRank> dimArr;
|
|
std::array<int64_t, DDim::kMaxRank> strideArr;
|
|
auto id = dimArr.begin(), is = strideArr.begin();
|
|
for (int i = 0; i < in.dims().size(); i++) {
|
|
if (i == axis) continue;
|
|
if (i == axis + 1)
|
|
*id = in.dims()[i] * sliceLength;
|
|
else
|
|
*id = in.dims()[i];
|
|
*is = in.strides()[i];
|
|
id++;
|
|
is++;
|
|
}
|
|
*out = DenseTensor{
|
|
in.Holder(),
|
|
DenseTensorMeta{in.dtype(),
|
|
DDim{dimArr.data(), in.dims().size() - 1},
|
|
DDim(strideArr.data(), in.dims().size() - 1)}};
|
|
out->set_offset(in.offset() +
|
|
offset * in.strides()[axis] * SizeOf(out->dtype()));
|
|
}
|
|
template <typename T, typename Context>
|
|
void FlashAttnVarlenQKVPackedKernel(
|
|
const Context& dev_ctx,
|
|
const DenseTensor& qkv,
|
|
const DenseTensor& cu_seqlens_q,
|
|
const DenseTensor& cu_seqlens_k,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const Scalar& max_seqlen_q,
|
|
const Scalar& max_seqlen_k,
|
|
float scale,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
bool varlen_padded,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
const auto head_groupnum = qkv.dims()[1]; // nheads/nheads_k + 1 + 1
|
|
DenseTensor q, k, v;
|
|
sliceFlattenView(qkv, &q, 1, 0, head_groupnum - 2);
|
|
sliceFlattenView(qkv, &k, 1, head_groupnum - 2, 1);
|
|
sliceFlattenView(qkv, &v, 1, head_groupnum - 1, 1);
|
|
FlashAttnUnpaddedBaseKernel<T>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
is_test,
|
|
rng_name,
|
|
out,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset,
|
|
varlen_padded);
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnBaseKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const optional<DenseTensor>& startend_row_indices,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
// q, k, v [batch_size, seq_len, num_heads, head_dim]
|
|
const auto& dims = q.dims();
|
|
PADDLE_ENFORCE_EQ(dims.size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(k.dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(v.dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
PADDLE_ENFORCE_EQ(out->dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
const int64_t batch_size = dims[0];
|
|
const int64_t seqlen_q = dims[1];
|
|
const int64_t num_heads = dims[2];
|
|
const int64_t head_size = dims[3];
|
|
const int64_t seqlen_k = k.dims()[1];
|
|
const int64_t num_heads_k = k.dims()[2];
|
|
const int64_t head_size_v = v.dims()[3];
|
|
|
|
PADDLE_ENFORCE_EQ(head_size,
|
|
head_size_v,
|
|
common::errors::InvalidArgument(
|
|
"This kernel does not support headdim != headdim_v, "
|
|
"but got headdim = %d and headdim_v = %d",
|
|
head_size,
|
|
head_size_v));
|
|
|
|
// TODO(umiswing): Add check shape
|
|
|
|
// TODO(GuoxiaWang): implement use_gqa_packing logic
|
|
bool use_gqa_packing = false;
|
|
const float softmax_scale = 1.0f / std::sqrt(head_size);
|
|
const float softmax_unscale = std::sqrt(head_size);
|
|
|
|
int version = FLAGS_flash_attn_version == 3 && FLAGS_cudnn_deterministic &&
|
|
head_size > 128
|
|
? 2
|
|
: FLAGS_flash_attn_version;
|
|
FlashAttnFwdParamsV2<T> params =
|
|
FlashAttnFwdParamsV2<T>(dev_ctx,
|
|
version,
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
dropout,
|
|
softmax_scale,
|
|
causal,
|
|
return_softmax,
|
|
q.dtype(),
|
|
is_test,
|
|
rng_name,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
startend_row_indices,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset,
|
|
/*_unpadded_lse*/ false,
|
|
/*_total_q*/ 0);
|
|
|
|
VLOG(10) << "[FlashAttn Forward" << version << "] q.shape=[" << q.dims()
|
|
<< "], k.shape=[" << k.dims() << "], v.shape=[" << v.dims() << "]";
|
|
VLOG(10) << "[FlashAttn Forward" << version << "] dropout=" << dropout
|
|
<< ", seed=" << params.seed << ", offset=" << params.offset;
|
|
VLOG(10) << "[FlashAttn Forward" << version
|
|
<< "] softmax_scale=" << softmax_scale
|
|
<< ", softmax_unscale=" << softmax_unscale;
|
|
if (attn_mask.get_ptr()) {
|
|
VLOG(10) << "[FlashAttn Forward" << version << "] attn_mask.shape=["
|
|
<< (attn_mask.get_ptr())->dims() << "]";
|
|
}
|
|
if (!out->IsInitialized()) dev_ctx.template Alloc<T>(out);
|
|
|
|
DenseTensor flashmask_maxmin, downstart_row_indices, upend_row_indices,
|
|
downend_row_indices, upstart_row_indices;
|
|
void *downstart_row_indices_data = nullptr, *upend_row_indices_data = nullptr,
|
|
*downend_row_indices_data = nullptr, *upstart_row_indices_data = nullptr;
|
|
bool is_flashmask = params.startend_row_indices != nullptr;
|
|
if (is_flashmask) {
|
|
PADDLE_ENFORCE_EQ(
|
|
startend_row_indices->dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flashmask_attention receive startend_row_indices with dim "
|
|
"[batch_size, num_heads,seq_len, mask_bounds]"));
|
|
PADDLE_ENFORCE_EQ(startend_row_indices->dims()[3] == 1 ||
|
|
startend_row_indices->dims()[3] == 2 ||
|
|
startend_row_indices->dims()[3] == 4,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"flashmask_attention startend_row_indices "
|
|
"mask_bounds must in [1,2,4]"));
|
|
auto flashmask_maxmin_shape = params.startend_row_indices->dims();
|
|
flashmask_maxmin_shape[2] = (flashmask_maxmin_shape[2] + 31) / 32 * 8;
|
|
flashmask_maxmin.set_type(DataType::INT32);
|
|
flashmask_maxmin.Resize(flashmask_maxmin_shape);
|
|
dev_ctx.template Alloc<T>(&flashmask_maxmin);
|
|
|
|
downstart_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {0}, {1});
|
|
downstart_row_indices_data = downstart_row_indices.data();
|
|
if (startend_row_indices->dims()[3] == 2) {
|
|
if (!causal) {
|
|
upend_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
upend_row_indices_data = upend_row_indices.data();
|
|
} else {
|
|
downend_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
downend_row_indices_data = downend_row_indices.data();
|
|
}
|
|
} else if (startend_row_indices->dims()[3] == 4) {
|
|
upend_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {3}, {4});
|
|
upend_row_indices_data = upend_row_indices.data();
|
|
downend_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
downend_row_indices_data = downend_row_indices.data();
|
|
upstart_row_indices =
|
|
Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {2}, {3});
|
|
upstart_row_indices_data = upstart_row_indices.data();
|
|
}
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
hipStream_t stream = dev_ctx.stream();
|
|
#else
|
|
cudaStream_t stream = dev_ctx.stream();
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
bool succ = dynload::flash_attn_fwd(
|
|
q.data(),
|
|
k.data(),
|
|
v.data(),
|
|
params.rng_state.data(),
|
|
out->data(),
|
|
params.return_softmax ? params.softmax->data() : nullptr,
|
|
params.softmax_lse->data(),
|
|
params.batch_size,
|
|
params.max_seqlen_q,
|
|
params.max_seqlen_k,
|
|
params.seqlen_q_rounded,
|
|
params.seqlen_k_rounded,
|
|
params.num_heads,
|
|
params.num_heads_k,
|
|
params.head_size,
|
|
params.head_size_rounded,
|
|
params.dropout,
|
|
params.softmax_scale,
|
|
softmax_unscale, // for unscale
|
|
params.causal,
|
|
params.return_softmax,
|
|
params.is_bf16,
|
|
stream,
|
|
params.seed,
|
|
params.offset,
|
|
params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
|
|
params.mask_dims.data(),
|
|
is_flashmask ? downstart_row_indices_data : nullptr,
|
|
is_flashmask ? params.startend_row_indices_dims.data() : nullptr,
|
|
is_flashmask ? upend_row_indices_data : nullptr,
|
|
is_flashmask ? downend_row_indices_data : nullptr,
|
|
is_flashmask ? upstart_row_indices_data : nullptr,
|
|
is_flashmask ? flashmask_maxmin.data() : nullptr,
|
|
q.strides()[1],
|
|
k.strides()[1],
|
|
v.strides()[1],
|
|
q.strides()[2],
|
|
k.strides()[2],
|
|
v.strides()[2],
|
|
out->strides()[1],
|
|
out->strides()[2],
|
|
q.strides()[0],
|
|
k.strides()[0],
|
|
v.strides()[0],
|
|
out->strides()[0]);
|
|
#else
|
|
bool succ;
|
|
int arch =
|
|
backends::gpu::GetGPUComputeCapability(dev_ctx.GetPlace().GetDeviceId());
|
|
|
|
if (arch == 80 && version == 3) {
|
|
RaiseNotSupportedError(3);
|
|
}
|
|
|
|
if (arch == 90 && version == 3) {
|
|
#ifdef PADDLE_WITH_FLASHATTN_V3
|
|
if (is_flashmask || params.attn_mask_tensor) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"FlashMask or Dense Mask is unsupported in FlashAttention V3"));
|
|
}
|
|
|
|
FlashAttnV3Kernel<T, Context>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
paddle::none, // q_v_
|
|
paddle::none, // q_descale_
|
|
paddle::none, // k_descale_
|
|
paddle::none, // v_descale_
|
|
params.softmax_scale,
|
|
params.causal,
|
|
-1, // window_size_left
|
|
-1, // window_size_right
|
|
0.f, // softcap
|
|
1, // num_splits
|
|
false, // manual_set_pack_gqa
|
|
false, // pack_gqa_
|
|
0, // sm_margin
|
|
out,
|
|
softmax_lse);
|
|
#else
|
|
RaiseNotSupportedError(3);
|
|
#endif
|
|
} else {
|
|
succ = dynload::flash_attn_fwd(
|
|
q.data(),
|
|
k.data(),
|
|
v.data(),
|
|
params.rng_state.data(),
|
|
out->data(),
|
|
params.return_softmax ? params.softmax->data() : nullptr,
|
|
params.softmax_lse->data(),
|
|
params.batch_size,
|
|
params.max_seqlen_q,
|
|
params.max_seqlen_k,
|
|
params.seqlen_q_rounded,
|
|
params.seqlen_k_rounded,
|
|
params.num_heads,
|
|
params.num_heads_k,
|
|
params.head_size,
|
|
params.head_size_rounded,
|
|
params.dropout,
|
|
params.softmax_scale,
|
|
softmax_unscale,
|
|
params.causal,
|
|
params.return_softmax,
|
|
params.is_bf16,
|
|
stream,
|
|
params.seed,
|
|
params.offset,
|
|
params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
|
|
params.mask_dims.data(),
|
|
is_flashmask ? downstart_row_indices_data : nullptr,
|
|
is_flashmask ? params.startend_row_indices_dims.data() : nullptr,
|
|
is_flashmask ? upend_row_indices_data : nullptr,
|
|
is_flashmask ? downend_row_indices_data : nullptr,
|
|
is_flashmask ? upstart_row_indices_data : nullptr,
|
|
is_flashmask ? flashmask_maxmin.data() : nullptr,
|
|
q.strides()[1],
|
|
k.strides()[1],
|
|
v.strides()[1],
|
|
q.strides()[2],
|
|
k.strides()[2],
|
|
v.strides()[2],
|
|
out->strides()[1],
|
|
out->strides()[2],
|
|
q.strides()[0],
|
|
k.strides()[0],
|
|
v.strides()[0],
|
|
out->strides()[0]);
|
|
}
|
|
#endif
|
|
if (version != 3) {
|
|
CheckFlashAttnStatus(succ); // umiswing: no return status in fa3
|
|
}
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
if (q.numel() == 0 || k.numel() == 0 || v.numel() == 0) {
|
|
if (out) {
|
|
Full<T, Context>(dev_ctx, out->dims(), 0, out);
|
|
}
|
|
if (softmax) {
|
|
Full<T, Context>(dev_ctx, softmax->dims(), 0, softmax);
|
|
}
|
|
if (softmax_lse) {
|
|
Full<T, Context>(dev_ctx, softmax_lse->dims(), 0, softmax_lse);
|
|
}
|
|
if (seed_offset) {
|
|
Full<T, Context>(dev_ctx, seed_offset->dims(), 0, seed_offset);
|
|
}
|
|
return;
|
|
}
|
|
FlashAttnBaseKernel<T, Context>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
paddle::none,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
is_test,
|
|
rng_name,
|
|
out,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnQKVPackedKernel(const Context& dev_ctx,
|
|
const DenseTensor& qkv,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
const auto head_groupnum = qkv.dims()[2]; // nheads/nheads_k + 1 + 1
|
|
DenseTensor q, k, v;
|
|
sliceFlattenView(qkv, &q, 2, 0, head_groupnum - 2);
|
|
sliceFlattenView(qkv, &k, 2, head_groupnum - 2, 1);
|
|
sliceFlattenView(qkv, &v, 2, head_groupnum - 1, 1);
|
|
FlashAttnBaseKernel<T, Context>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
fixed_seed_offset,
|
|
attn_mask,
|
|
paddle::none,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
is_test,
|
|
rng_name,
|
|
out,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset);
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashMaskKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& startend_row_indices,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
float dropout,
|
|
bool causal,
|
|
bool return_softmax,
|
|
bool is_test,
|
|
const std::string& rng_name,
|
|
DenseTensor* out,
|
|
DenseTensor* softmax,
|
|
DenseTensor* softmax_lse,
|
|
DenseTensor* seed_offset) {
|
|
// Handle 0-size tensors: return zeros without calling CUDA kernel
|
|
// to avoid invalid memory access
|
|
if (q.numel() == 0 || k.numel() == 0 || v.numel() == 0) {
|
|
if (out) {
|
|
Full<T, Context>(dev_ctx, out->dims(), 0, out);
|
|
}
|
|
if (softmax) {
|
|
Full<T, Context>(dev_ctx, softmax->dims(), 0, softmax);
|
|
}
|
|
if (softmax_lse) {
|
|
Full<T, Context>(dev_ctx, softmax_lse->dims(), 0, softmax_lse);
|
|
}
|
|
if (seed_offset) {
|
|
Full<T, Context>(dev_ctx, seed_offset->dims(), 0, seed_offset);
|
|
}
|
|
return;
|
|
}
|
|
|
|
FlashAttnBaseKernel<T, Context>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
fixed_seed_offset,
|
|
paddle::none,
|
|
startend_row_indices,
|
|
dropout,
|
|
causal,
|
|
return_softmax,
|
|
is_test,
|
|
rng_name,
|
|
out,
|
|
softmax,
|
|
softmax_lse,
|
|
seed_offset);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(flash_attn_unpadded,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnUnpaddedKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(5).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flash_attn_varlen_qkvpacked,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnVarlenQKVPackedKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(3).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flash_attn,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(3).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flash_attn_qkvpacked,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnQKVPackedKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(1).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flashmask_attention,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashMaskKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(4).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|