wehub-resource-sync
440 lines
19 KiB
Plaintext
440 lines
19 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 <cuda_fp16.h>
|
|
#include <cub/cub.cuh>
|
|
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/broadcast_function.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
|
|
#include "paddle/phi/kernels/funcs/functors.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
|
#include "paddle/phi/kernels/funcs/transpose_function.cuh"
|
|
#include "paddle/phi/kernels/fusion/gpu/attention_layer.norm.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/attn_gemm.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fmha_ref.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_attention_utils.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_dropout_helper.h"
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
template <typename T, typename Context>
|
|
void FusedAttentionKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const optional<DenseTensor> &ln_scale,
|
|
const optional<DenseTensor> &ln_bias,
|
|
const DenseTensor &qkv_weight,
|
|
const optional<DenseTensor> &qkv_bias,
|
|
const optional<DenseTensor> &cache_kv,
|
|
const optional<DenseTensor> &src_mask,
|
|
const DenseTensor &out_linear_weight,
|
|
const optional<DenseTensor> &out_linear_bias,
|
|
const optional<DenseTensor> &ln_scale_2,
|
|
const optional<DenseTensor> &ln_bias_2,
|
|
int num_heads,
|
|
bool transpose_qkv_wb,
|
|
bool pre_layer_norm,
|
|
float epsilon,
|
|
float attn_dropout_rate,
|
|
bool is_test,
|
|
bool attn_dropout_fix_seed,
|
|
int attn_dropout_seed,
|
|
const std::string &attn_dropout_implementation,
|
|
float dropout_rate,
|
|
bool dropout_fix_seed,
|
|
int dropout_seed,
|
|
const std::string &dropout_implementation,
|
|
float ln_epsilon,
|
|
bool add_residual,
|
|
int ring_id,
|
|
DenseTensor *ln_mean,
|
|
DenseTensor *ln_var,
|
|
DenseTensor *ln_out,
|
|
DenseTensor *qkv_out,
|
|
DenseTensor *qkv_bias_out,
|
|
DenseTensor *transpose_out_2,
|
|
DenseTensor *qk_out,
|
|
DenseTensor *qktv_out,
|
|
DenseTensor *softmax_out,
|
|
DenseTensor *attn_dropout_mask_out,
|
|
DenseTensor *attn_dropout_out,
|
|
DenseTensor *src_mask_out,
|
|
DenseTensor *fmha_out,
|
|
DenseTensor *out_linear_out,
|
|
DenseTensor *dropout_mask_out,
|
|
DenseTensor *ln_mean_2,
|
|
DenseTensor *ln_var_2,
|
|
DenseTensor *bias_dropout_residual_out,
|
|
DenseTensor *cache_kv_out,
|
|
DenseTensor *out) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
if (x.numel() == 0) {
|
|
if (ln_mean) dev_ctx.template Alloc<U>(ln_mean);
|
|
if (ln_var) dev_ctx.template Alloc<U>(ln_var);
|
|
if (ln_out) dev_ctx.template Alloc<T>(ln_out);
|
|
if (qkv_out) dev_ctx.template Alloc<T>(qkv_out);
|
|
if (qkv_bias_out) dev_ctx.template Alloc<T>(qkv_bias_out);
|
|
if (transpose_out_2) dev_ctx.template Alloc<T>(transpose_out_2);
|
|
if (qk_out) dev_ctx.template Alloc<T>(qk_out);
|
|
if (qktv_out) dev_ctx.template Alloc<T>(qktv_out);
|
|
if (softmax_out) dev_ctx.template Alloc<T>(softmax_out);
|
|
if (attn_dropout_mask_out)
|
|
dev_ctx.template Alloc<uint8_t>(attn_dropout_mask_out);
|
|
if (attn_dropout_out) dev_ctx.template Alloc<T>(attn_dropout_out);
|
|
if (src_mask_out) dev_ctx.template Alloc<T>(src_mask_out);
|
|
if (fmha_out) dev_ctx.template Alloc<T>(fmha_out);
|
|
if (out_linear_out) dev_ctx.template Alloc<T>(out_linear_out);
|
|
if (dropout_mask_out) dev_ctx.template Alloc<uint8_t>(dropout_mask_out);
|
|
if (ln_mean_2) dev_ctx.template Alloc<U>(ln_mean_2);
|
|
if (ln_var_2) dev_ctx.template Alloc<U>(ln_var_2);
|
|
if (bias_dropout_residual_out)
|
|
dev_ctx.template Alloc<T>(bias_dropout_residual_out);
|
|
if (cache_kv_out) dev_ctx.template Alloc<T>(cache_kv_out);
|
|
dev_ctx.template Alloc<T>(out);
|
|
return;
|
|
}
|
|
|
|
// x: qkv's input [batch_size, seq_len, dim_embed]
|
|
// if transpose_qkv_wb is False
|
|
// y: qkv's weight: [3, num_head, dim_head, dim_embed]
|
|
// if transpose_qkv_wb is True
|
|
// y: qkv's weight: [dim_embed, 3 * dim_embed]
|
|
|
|
auto *x_p = &x;
|
|
auto *ln_scale_p = ln_scale.get_ptr();
|
|
auto *ln_bias_p = ln_bias.get_ptr();
|
|
|
|
auto *qkv_weight_p = &qkv_weight;
|
|
auto *qkv_bias_p = qkv_bias.get_ptr();
|
|
auto *cache_kv_p = cache_kv.get_ptr();
|
|
|
|
auto *src_mask_p = src_mask.get_ptr();
|
|
auto *out_linear_weight_p = &out_linear_weight;
|
|
|
|
auto *out_linear_bias_p = out_linear_bias.get_ptr();
|
|
|
|
auto *ln_scale_2_p = ln_scale_2.get_ptr();
|
|
auto *ln_bias_2_p = ln_bias_2.get_ptr();
|
|
|
|
const bool has_attn_dropout = (attn_dropout_rate != 0.0f);
|
|
|
|
const bool is_upscale_in_train =
|
|
(dropout_implementation == "upscale_in_train");
|
|
fusion::DropoutParam dropout_param2(dropout_fix_seed,
|
|
0,
|
|
is_test,
|
|
is_upscale_in_train,
|
|
dropout_rate,
|
|
nullptr,
|
|
dropout_seed);
|
|
|
|
const bool has_dropout = (dropout_param2.dropout_prob != 0.0f);
|
|
|
|
bool is_upscale_in_train_1 =
|
|
(attn_dropout_implementation == "upscale_in_train");
|
|
DenseTensor *seed_1 = nullptr;
|
|
|
|
// get data ptr for qkv part.
|
|
const auto input_x_dims = x_p->dims();
|
|
const auto qkv_w_dims = qkv_weight_p->dims();
|
|
|
|
auto *x_data = x_p->data<T>();
|
|
auto *qkv_weight_data = qkv_weight_p->data<T>();
|
|
auto *qkv_bias_data =
|
|
(qkv_bias_p == nullptr) ? nullptr : qkv_bias_p->data<T>();
|
|
auto *qkv_out_data =
|
|
dev_ctx.template Alloc<T>(qkv_out, qkv_out->numel() * sizeof(T));
|
|
auto *qkv_bias_out_data =
|
|
(qkv_bias_p == nullptr)
|
|
? nullptr
|
|
: dev_ctx.template Alloc<T>(qkv_bias_out,
|
|
qkv_bias_out->numel() * sizeof(T));
|
|
|
|
// get data ptr for FMHA.
|
|
auto *transpose_out_2_data = dev_ctx.template Alloc<T>(
|
|
transpose_out_2, transpose_out_2->numel() * sizeof(T));
|
|
auto *cache_kv_out_data =
|
|
(cache_kv_out == nullptr)
|
|
? nullptr
|
|
: dev_ctx.template Alloc<T>(cache_kv_out,
|
|
cache_kv_out->numel() * sizeof(T));
|
|
auto *qk_out_data =
|
|
dev_ctx.template Alloc<T>(qk_out, qk_out->numel() * sizeof(T));
|
|
auto *qktv_out_data =
|
|
dev_ctx.template Alloc<T>(qktv_out, qktv_out->numel() * sizeof(T));
|
|
auto *src_mask_out_data =
|
|
(src_mask_p == nullptr)
|
|
? nullptr
|
|
: dev_ctx.template Alloc<T>(src_mask_out,
|
|
src_mask_out->numel() * sizeof(T));
|
|
auto *softmax_out_data =
|
|
dev_ctx.template Alloc<T>(softmax_out, softmax_out->numel() * sizeof(T));
|
|
auto *attn_dropout_mask_out_data =
|
|
has_attn_dropout ? dev_ctx.template Alloc<uint8_t>(
|
|
attn_dropout_mask_out,
|
|
attn_dropout_mask_out->numel() * sizeof(uint8_t))
|
|
: nullptr;
|
|
auto *attn_dropout_out_data =
|
|
has_attn_dropout
|
|
? dev_ctx.template Alloc<T>(attn_dropout_out,
|
|
attn_dropout_out->numel() * sizeof(T))
|
|
: nullptr;
|
|
auto *fmha_out_data =
|
|
dev_ctx.template Alloc<T>(fmha_out, fmha_out->numel() * sizeof(T));
|
|
|
|
// get data ptr for out_linear.
|
|
auto *out_linear_weight_data = out_linear_weight_p->data<T>();
|
|
auto *out_linear_bias_data =
|
|
(out_linear_bias_p == nullptr) ? nullptr : out_linear_bias_p->data<T>();
|
|
auto *out_linear_out_data = dev_ctx.template Alloc<T>(
|
|
out_linear_out, out_linear_out->numel() * sizeof(T));
|
|
|
|
// get data ptr for bias+dropout+residual+layernorm
|
|
auto *dropout_mask_out_data =
|
|
has_dropout
|
|
? dev_ctx.template Alloc<uint8_t>(
|
|
dropout_mask_out, dropout_mask_out->numel() * sizeof(uint8_t))
|
|
: nullptr;
|
|
auto *final_out_data =
|
|
dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
|
|
|
|
int batch_size = input_x_dims[0];
|
|
int max_seq_len = input_x_dims[1];
|
|
int dim_embed = input_x_dims[2];
|
|
|
|
int num_head;
|
|
int dim_head;
|
|
int nranks = 1;
|
|
// get num_head and dim_head in two different ways
|
|
if (!transpose_qkv_wb) {
|
|
num_head = qkv_w_dims[1];
|
|
dim_head = qkv_w_dims[2];
|
|
} else {
|
|
nranks = (qkv_w_dims[0] * 3) / qkv_w_dims[1];
|
|
num_head = num_heads;
|
|
dim_head = dim_embed / (num_head * nranks);
|
|
}
|
|
|
|
int64_t bsz_seq = static_cast<int64_t>(batch_size) * max_seq_len;
|
|
int64_t hidden_size = static_cast<int64_t>(num_head) * dim_head;
|
|
int64_t output_size = 3 * hidden_size;
|
|
int input_size = dim_embed;
|
|
|
|
auto layer_norm_compute =
|
|
fusion::AttnLayerNorm<T>(dev_ctx, epsilon, bsz_seq, dim_embed);
|
|
|
|
bool compute_bias = true;
|
|
if (qkv_bias_p == nullptr) {
|
|
compute_bias = false;
|
|
}
|
|
// (transA, transB, compute_bias) = (false, true, true)
|
|
bool transB = transpose_qkv_wb ? false : true;
|
|
PADDLE_ENFORCE_LE_INT_MAX(bsz_seq, "bsz_seq");
|
|
PADDLE_ENFORCE_LE_INT_MAX(output_size, "output_size");
|
|
auto qkv_compute = fusion::AttnMatMul<T>(dev_ctx,
|
|
false,
|
|
transB,
|
|
static_cast<int>(bsz_seq),
|
|
static_cast<int>(output_size),
|
|
input_size,
|
|
compute_bias);
|
|
|
|
fusion::AttnDropoutParam attn_dropout_param(is_test,
|
|
attn_dropout_implementation,
|
|
attn_dropout_rate,
|
|
is_upscale_in_train_1,
|
|
attn_dropout_fix_seed,
|
|
attn_dropout_seed,
|
|
seed_1);
|
|
auto fmha_ref_compute = fusion::FMHARef<T>(
|
|
dev_ctx, batch_size, max_seq_len, num_head, dim_head, attn_dropout_param);
|
|
|
|
output_size = hidden_size;
|
|
// (transA, transB, compute_bias) = (false, false, false)
|
|
// NOTE(Yuang Liu): For general input size == output size, change the
|
|
// position won't have effects. For mp, the output size is mp_head * dkey
|
|
// which is actually the input size. While the input size is hidden size,
|
|
// which is actually the output size. So for out linear, switch the
|
|
// input size and output size.
|
|
PADDLE_ENFORCE_LE_INT_MAX(bsz_seq, "bsz_seq");
|
|
PADDLE_ENFORCE_LE_INT_MAX(output_size, "output_size");
|
|
auto out_linear_compute = fusion::AttnMatMul<T>(dev_ctx,
|
|
false,
|
|
false,
|
|
static_cast<int>(bsz_seq),
|
|
input_size,
|
|
static_cast<int>(output_size),
|
|
false);
|
|
fusion::FusedDropoutLayerNormHelper<T, uint8_t>
|
|
fused_dropout_layernorm_helper(
|
|
dev_ctx, bsz_seq, dim_embed, dropout_param2, ln_epsilon);
|
|
|
|
if (pre_layer_norm) {
|
|
auto *ln_scale_data =
|
|
(ln_scale_p == nullptr ? nullptr : ln_scale_p->data<U>());
|
|
auto *ln_bias_data =
|
|
(ln_bias_p == nullptr ? nullptr : ln_bias_p->data<U>());
|
|
auto *ln_mean_data =
|
|
dev_ctx.template Alloc<U>(ln_mean, ln_mean->numel() * sizeof(U));
|
|
auto *ln_var_data =
|
|
dev_ctx.template Alloc<U>(ln_var, ln_var->numel() * sizeof(U));
|
|
auto *ln_out_data =
|
|
dev_ctx.template Alloc<T>(ln_out, ln_out->numel() * sizeof(T));
|
|
|
|
layer_norm_compute.ComputeForward(x_data,
|
|
ln_scale_data,
|
|
ln_bias_data,
|
|
ln_out_data,
|
|
ln_mean_data,
|
|
ln_var_data);
|
|
qkv_compute.ComputeForward(
|
|
qkv_weight_p, ln_out, qkv_bias_p, qkv_out, qkv_bias_out);
|
|
} else {
|
|
qkv_compute.ComputeForward(
|
|
qkv_weight_p, x_p, qkv_bias_p, qkv_out, qkv_bias_out);
|
|
}
|
|
|
|
if (transpose_qkv_wb) {
|
|
// resize the output for fmha compute
|
|
qkv_out->Resize({batch_size, max_seq_len, 3, num_head, dim_head});
|
|
qkv_bias_out->Resize({batch_size, max_seq_len, 3, num_head, dim_head});
|
|
}
|
|
|
|
if (qkv_bias_p == nullptr) {
|
|
fmha_ref_compute.ComputeForward(*qkv_out,
|
|
cache_kv_p,
|
|
src_mask_p,
|
|
transpose_out_2,
|
|
cache_kv_out,
|
|
qk_out,
|
|
src_mask_out,
|
|
softmax_out,
|
|
attn_dropout_mask_out,
|
|
attn_dropout_out,
|
|
qktv_out,
|
|
fmha_out);
|
|
} else {
|
|
fmha_ref_compute.ComputeForward(*qkv_bias_out,
|
|
cache_kv_p,
|
|
src_mask_p,
|
|
transpose_out_2,
|
|
cache_kv_out,
|
|
qk_out,
|
|
src_mask_out,
|
|
softmax_out,
|
|
attn_dropout_mask_out,
|
|
attn_dropout_out,
|
|
qktv_out,
|
|
fmha_out);
|
|
}
|
|
|
|
if (transpose_qkv_wb) {
|
|
// resize the output back to make the shape compatible with infer shape
|
|
qkv_out->Resize({batch_size, max_seq_len, 3 * hidden_size});
|
|
qkv_bias_out->Resize({batch_size, max_seq_len, 3 * hidden_size});
|
|
}
|
|
|
|
// fmha_out: [batch_size, seq_len, num_head, head_dim]
|
|
// weight: [embed_dim, embed_dim]
|
|
// out_linear_out: [batch_size, seq_len, embed_dim]
|
|
out_linear_compute.ComputeForward(
|
|
out_linear_weight_p, fmha_out, nullptr, out_linear_out, nullptr);
|
|
// tensor model parallel
|
|
phi::fusion::AllReduce<T>(*out_linear_out, ring_id, dev_ctx);
|
|
|
|
const T *residual_ptr = add_residual ? x_data : nullptr;
|
|
if (pre_layer_norm) {
|
|
// output = (residual + dropout(input + bias))
|
|
fused_dropout_layernorm_helper.ResidualDropoutBias(dev_ctx,
|
|
out_linear_out_data,
|
|
residual_ptr,
|
|
out_linear_bias_data,
|
|
final_out_data,
|
|
dropout_mask_out_data);
|
|
} else {
|
|
// TODO(Xreki): support post layer_norm case when add_residual is false.
|
|
PADDLE_ENFORCE_EQ(
|
|
add_residual,
|
|
true,
|
|
errors::InvalidArgument("Attribute add_residual is expected to be true "
|
|
"when pre_layer_norm is false."));
|
|
|
|
const U *ln_scale_2_ptr = ln_scale_2_p ? ln_scale_2_p->data<U>() : nullptr;
|
|
const U *ln_bias_2_ptr = ln_bias_2_p ? ln_bias_2_p->data<U>() : nullptr;
|
|
T *bias_dropout_residual_out_ptr = dev_ctx.template Alloc<T>(
|
|
bias_dropout_residual_out,
|
|
bias_dropout_residual_out->numel() * sizeof(T));
|
|
U *ln_mean_2_ptr =
|
|
dev_ctx.template Alloc<U>(ln_mean_2, ln_mean_2->numel() * sizeof(U));
|
|
U *ln_var_2_ptr =
|
|
dev_ctx.template Alloc<U>(ln_var_2, ln_var_2->numel() * sizeof(U));
|
|
|
|
// 0-size
|
|
if (ln_scale_2_p && ln_scale_2_p->numel() == 0) {
|
|
// output = (residual + dropout(input + bias))
|
|
fused_dropout_layernorm_helper.ResidualDropoutBias(dev_ctx,
|
|
out_linear_out_data,
|
|
residual_ptr,
|
|
out_linear_bias_data,
|
|
final_out_data,
|
|
dropout_mask_out_data);
|
|
return;
|
|
}
|
|
|
|
// output = layernorm(residual + dropout(input + bias))
|
|
fused_dropout_layernorm_helper.LayernormResidualDropoutBias(
|
|
dev_ctx,
|
|
out_linear_out_data,
|
|
residual_ptr,
|
|
out_linear_bias_data,
|
|
ln_scale_2_ptr,
|
|
ln_bias_2_ptr,
|
|
bias_dropout_residual_out_ptr,
|
|
dropout_mask_out_data,
|
|
final_out_data,
|
|
ln_mean_2_ptr,
|
|
ln_var_2_ptr);
|
|
}
|
|
}
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(fused_attention,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::fusion::FusedAttentionKernel,
|
|
phi::float16,
|
|
double,
|
|
float) {
|
|
kernel->OutputAt(9).SetDataType(phi::DataType::UINT8);
|
|
kernel->OutputAt(14).SetDataType(phi::DataType::UINT8);
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
|
kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(15).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(16).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
}
|