wehub-resource-sync
141 lines
6.0 KiB
Plaintext
141 lines
6.0 KiB
Plaintext
// Copyright (c) 2024 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/calc_reduced_attn_kernel.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
|
|
|
|
namespace phi {
|
|
|
|
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
|
|
struct CalcReducedAttnScoresParams : public FlashAttnParamsBase {
|
|
bool return_softmax;
|
|
DenseTensor* softmax;
|
|
|
|
CalcReducedAttnScoresParams(const GPUContext& dev_ctx,
|
|
const int _batch_size,
|
|
const int64_t _max_seqlen_q,
|
|
const int64_t _max_seqlen_k,
|
|
const int _num_heads,
|
|
const int _num_heads_k,
|
|
const int _head_size,
|
|
const float _scale,
|
|
const DataType q_dtype)
|
|
: FlashAttnParamsBase(/*version=*/2,
|
|
/*is_fwd=*/true,
|
|
_batch_size,
|
|
_max_seqlen_q,
|
|
_max_seqlen_k,
|
|
_num_heads,
|
|
_num_heads_k,
|
|
_head_size,
|
|
_scale,
|
|
/*_causal=*/false,
|
|
q_dtype,
|
|
optional<DenseTensor>{},
|
|
optional<DenseTensor>{},
|
|
/*_unpadded_lse=*/false,
|
|
/*_total_q*/ 0) {}
|
|
};
|
|
#endif
|
|
|
|
template <typename T, typename Context>
|
|
void CalcReducedAttnScoresKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& softmax_lse,
|
|
DenseTensor* reduced_scores) {
|
|
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
|
|
PADDLE_ENFORCE_EQ(q.dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"calc_reduced_attention receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
|
|
PADDLE_ENFORCE_EQ(k.dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"calc_reduced_attention receive input with dim "
|
|
"[batch_size, seq_len, num_heads, head_dim]"));
|
|
|
|
if (!reduced_scores->IsInitialized())
|
|
dev_ctx.template Alloc<float>(reduced_scores);
|
|
funcs::SetConstant<Context, float> set_zero;
|
|
set_zero(dev_ctx, reduced_scores, 0.0f);
|
|
// q, k, v [batch_size, seq_len, num_heads, head_dim]
|
|
const int64_t batch_size = q.dims()[0];
|
|
const int64_t seqlen_q = q.dims()[1];
|
|
const int64_t num_heads = q.dims()[2];
|
|
const int64_t head_size = q.dims()[3];
|
|
const int64_t seqlen_k = k.dims()[1];
|
|
const int64_t num_heads_k = k.dims()[2];
|
|
|
|
const float softmax_scale = 1.0f / std::sqrt(head_size);
|
|
const float softmax_unscale = std::sqrt(head_size);
|
|
|
|
using Params = CalcReducedAttnScoresParams;
|
|
|
|
Params params = Params(dev_ctx,
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
softmax_scale,
|
|
q.dtype());
|
|
|
|
cudaStream_t stream = dev_ctx.stream();
|
|
|
|
bool succ = dynload::calc_reduced_attn_scores(q.data(),
|
|
k.data(),
|
|
softmax_lse.data(),
|
|
reduced_scores->data(),
|
|
/*softmax_ptr=*/nullptr,
|
|
params.batch_size,
|
|
params.max_seqlen_q,
|
|
params.max_seqlen_k,
|
|
params.num_heads,
|
|
params.num_heads_k,
|
|
params.head_size,
|
|
params.softmax_scale,
|
|
/*return_softmax=*/false,
|
|
params.is_bf16,
|
|
/*num_splits=*/0,
|
|
stream,
|
|
q.strides()[1],
|
|
k.strides()[1],
|
|
reduced_scores->strides()[1],
|
|
q.strides()[2],
|
|
k.strides()[2],
|
|
reduced_scores->strides()[2],
|
|
q.strides()[0],
|
|
k.strides()[0],
|
|
reduced_scores->strides()[0]);
|
|
CheckFlashAttnStatus(succ);
|
|
#else
|
|
RaiseNotSupportedError();
|
|
#endif
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(calc_reduced_attn_scores,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::CalcReducedAttnScoresKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|