wehub-resource-sync
1376 lines
51 KiB
Plaintext
1376 lines
51 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/fusion/gpu/masked_multihead_attention_kernel.h"
|
||
#include "paddle/common/enforce.h"
|
||
#include "paddle/phi/core/kernel_registry.h"
|
||
#include "paddle/phi/kernels/funcs/aligned_vector.h"
|
||
#include "paddle/phi/kernels/fusion/gpu/mmha_util.cu.h"
|
||
|
||
namespace phi {
|
||
namespace fusion {
|
||
|
||
#ifndef PADDLE_WITH_HIP
|
||
|
||
constexpr unsigned int str2int(const char *str, int h = 0) {
|
||
return !str[h] ? 5381 : (str2int(str, h + 1) * 33) ^ str[h];
|
||
}
|
||
|
||
template <typename T>
|
||
struct Masked_multihead_attention_params {
|
||
float *qk_sum_max_split_seq;
|
||
float *split_out;
|
||
// qkv_out, [B, 1(seq_len), 3, num_head * dim_head]
|
||
const T *qkv;
|
||
// bias, [3, num_head, dim_head]
|
||
T *qkv_bias;
|
||
// [2, B, num_head, max_seq_len(valid cache_seq_len), dim_head]
|
||
// k [B, num_head, dim_head/x, max_seq_len, x], that is `seq_len` first
|
||
// v [B, num_head, max_seq_len, dim_head]
|
||
T *cache_kv;
|
||
// [B, max_seq_len]
|
||
const int *beam_cache_offset = nullptr;
|
||
|
||
const int *sequence_lengths{nullptr};
|
||
|
||
// The RoPE embedding, [2, B, rotary_seq_len, 1, dim_head]
|
||
// rotary_emb_dims = 1 if pos_ids_extra is null else 2
|
||
const float *rotary_emb;
|
||
|
||
// TODO(wangxi): optimize with input_lengths and max_input_len?
|
||
// [bsz, 1, 1, time_step(cache_seq_length)+1]
|
||
const T *attn_mask;
|
||
|
||
int rotary_emb_dims;
|
||
int batch_size; // batch * beam
|
||
int beam_width;
|
||
int cache_batch_size;
|
||
int num_head;
|
||
// k_num_head and v_num_head must be equal, we unify them.
|
||
// kv_num_head = k_num_head && kv_num_head == v_num_head
|
||
int kv_num_head;
|
||
int timestep; // cache_seq_length
|
||
int max_seq_length;
|
||
|
||
// 1.f / sqrt(Dh)
|
||
float inv_sqrt_dh;
|
||
|
||
int steps_per_block;
|
||
int split_seq = 1;
|
||
|
||
bool add_qkv_bias;
|
||
bool neox_rotary_style;
|
||
// whether to broadcast num_heads(2nd) dimension for attn_mask
|
||
// in MMHA, if false, attn_mask shape should be
|
||
// [bsz, num_heads, 1, time_step(cache_seq_length)+1]
|
||
bool mask_broadcast_num_heads;
|
||
};
|
||
|
||
template <typename T,
|
||
int Dh,
|
||
int Dh_MAX,
|
||
int THREADS_PER_KEY,
|
||
int THREADS_PER_VALUE,
|
||
int THREADS_PER_BLOCK,
|
||
typename LoadFunc,
|
||
typename StoreFunc,
|
||
bool SPLIT>
|
||
__global__ void masked_multihead_attention_kernel(
|
||
Masked_multihead_attention_params<T> params,
|
||
LoadFunc load_func,
|
||
StoreFunc store_func) {
|
||
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
|
||
const int bi = blockIdx.z;
|
||
// params.sequence_lengths[bi] means how many k and v we have cached in
|
||
// cache_kv.
|
||
if (params.sequence_lengths && params.sequence_lengths[bi] < 0) {
|
||
return;
|
||
}
|
||
|
||
typedef PDDataTypeTraits<T> traits_;
|
||
typedef typename traits_::DataType DataType_;
|
||
|
||
static_assert(Dh_MAX % THREADS_PER_KEY == 0, "");
|
||
static_assert(Dh_MAX % THREADS_PER_VALUE == 0, "");
|
||
|
||
constexpr int WARP_SIZE = 32;
|
||
constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE;
|
||
|
||
extern __shared__ char smem_[];
|
||
|
||
float *qk_smem = reinterpret_cast<float *>(smem_);
|
||
|
||
char *logits_smem_ = smem_;
|
||
// fp32 accum for logits
|
||
float *logits_smem = reinterpret_cast<float *>(logits_smem_);
|
||
|
||
T *out_smem = reinterpret_cast<T *>(smem_);
|
||
|
||
__shared__ float red_smem[WARPS_PER_BLOCK * 2];
|
||
using Qk_vec = typename Qk_vec_<T, Dh_MAX>::Type;
|
||
using Qk_vec_RoPE = typename Qk_vec_RoPE_<T, float, Dh_MAX>::Type;
|
||
__shared__ __align__(sizeof(Qk_vec)) T q_smem[Dh_MAX];
|
||
|
||
// beam id
|
||
const int beami = bi % params.beam_width;
|
||
// real batch id
|
||
const int bbi = bi / params.beam_width;
|
||
const int hi = blockIdx.y;
|
||
const int bhi = bi * params.num_head + hi;
|
||
|
||
const int kv_num_head = params.kv_num_head;
|
||
const int num_head_per_group = params.num_head / kv_num_head;
|
||
// hi means the head index in query processed by this cuda thread.
|
||
// kv_bhi means the merged batch and head index in key and value processed by
|
||
// this cuda thread.
|
||
const int kv_bhi = bi * kv_num_head + hi / num_head_per_group;
|
||
|
||
const int bbhi = bbi * params.beam_width * params.num_head + hi;
|
||
const int tid = threadIdx.x;
|
||
|
||
const int bi_seq_len_offset = bi * params.max_seq_length;
|
||
|
||
float qk_max = -FLT_MAX;
|
||
float qk = 0;
|
||
|
||
int act_time_step = params.sequence_lengths == nullptr
|
||
? params.timestep
|
||
: params.sequence_lengths[bi];
|
||
|
||
// with SPLIT, The last single q*k*v is computed by the last threadBlock of
|
||
// split_index
|
||
const int split_index = blockIdx.x;
|
||
int start_seq = 0;
|
||
int end_seq = act_time_step;
|
||
bool is_last_block = (SPLIT == false);
|
||
int real_split_each_batch = (act_time_step - 1) / params.steps_per_block + 1;
|
||
if constexpr (SPLIT) {
|
||
if (split_index >= real_split_each_batch) return;
|
||
|
||
start_seq = split_index * params.steps_per_block;
|
||
end_seq = start_seq + params.steps_per_block;
|
||
if (split_index == real_split_each_batch - 1) {
|
||
is_last_block = true;
|
||
end_seq = act_time_step;
|
||
}
|
||
}
|
||
int curr_seq_section = end_seq - start_seq;
|
||
|
||
// qkv [B, S=1, num_head + 2 * kv_num_head, head_dim]
|
||
// this hi means the head index in query!
|
||
int qkv_base_offset = bi * (params.num_head + 2 * kv_num_head) * Dh + hi * Dh;
|
||
|
||
// QK_VEC_SIZE is only used for compute q dot k .
|
||
constexpr int QK_VEC_SIZE = sizeof(Qk_vec) / sizeof(T);
|
||
static_assert(Dh_MAX % QK_VEC_SIZE == 0, "");
|
||
// Use block reduction if needed
|
||
// static_assert(Dh_MAX / QK_VEC_SIZE <= WARP_SIZE, "");
|
||
constexpr int QK_VECS_PER_WARP = Dh_MAX / QK_VEC_SIZE;
|
||
|
||
// cache_k, [B, num_head, head_dim / x, max_seq_len, x]
|
||
// x == 4/8 for FP32/FP16, 128bit, 16Byte
|
||
constexpr int QK_ELTS_IN_16B = 16 / sizeof(T);
|
||
constexpr int QK_VECS_IN_16B = 16 / sizeof(Qk_vec);
|
||
|
||
// const T *q_base = params.qkv;
|
||
// const T *k_base = params.qkv + params.num_head * Dh;
|
||
T *q_bias_base = nullptr;
|
||
T *k_bias_base = nullptr;
|
||
|
||
if (params.add_qkv_bias) {
|
||
q_bias_base = params.qkv_bias;
|
||
k_bias_base = params.qkv_bias + params.num_head * Dh;
|
||
}
|
||
|
||
// q and k have only head_dim scalar elements.
|
||
// below only compute q dot k = 1 element.
|
||
// q has QK_VECS_PER_WARP elements, [Qk_vec, Qk_vec, ..., Qk_vec]
|
||
// k has QK_VECS_PER_WARP elements: [Qk_vec, Qk_vec, ..., Qk_vec]
|
||
// per cuda thread read a Qk_vec of q and k and compute q dot k.
|
||
if (tid < QK_VECS_PER_WARP) {
|
||
int qk_offset = qkv_base_offset + tid * QK_VEC_SIZE;
|
||
int q_bias_offset = hi * Dh + tid * QK_VEC_SIZE;
|
||
int k_bias_offset = hi / num_head_per_group * Dh + tid * QK_VEC_SIZE;
|
||
|
||
Qk_vec q;
|
||
zero(q);
|
||
// q = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
// ? *reinterpret_cast<const Qk_vec *>(&q_base[qk_offset])
|
||
// : q;
|
||
if (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh) {
|
||
load_func.template load<Qk_vec>(q, qk_offset);
|
||
}
|
||
|
||
Qk_vec k;
|
||
zero(k);
|
||
// k = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
// ? *reinterpret_cast<const Qk_vec *>(&k_base[qk_offset])
|
||
// : k;
|
||
if ((Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh) && is_last_block) {
|
||
load_func.template load<Qk_vec>(k,
|
||
params.num_head * Dh + qk_offset -
|
||
hi * Dh +
|
||
hi / num_head_per_group * Dh);
|
||
}
|
||
|
||
if (params.add_qkv_bias) {
|
||
Qk_vec q_bias;
|
||
zero(q_bias);
|
||
Qk_vec k_bias;
|
||
zero(k_bias);
|
||
|
||
q_bias =
|
||
(Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec *>(&q_bias_base[q_bias_offset])
|
||
: q_bias;
|
||
k_bias =
|
||
(Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec *>(&k_bias_base[k_bias_offset])
|
||
: k_bias;
|
||
|
||
q = add(q, q_bias);
|
||
// TODO(wangxi): See this https://github.com/microsoft/unilm/issues/510
|
||
// we may not require k_bias.
|
||
k = add(k, k_bias);
|
||
}
|
||
|
||
if (!params.neox_rotary_style) {
|
||
if (params.rotary_emb_dims != 0) {
|
||
int rotary_offset = bi * Dh + tid * QK_VEC_SIZE;
|
||
const float *cos_base = params.rotary_emb;
|
||
const float *sin_base = params.rotary_emb + params.batch_size * Dh;
|
||
Qk_vec_RoPE cos_emb, sin_emb;
|
||
zero(cos_emb);
|
||
zero(sin_emb);
|
||
cos_emb = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec_RoPE *>(
|
||
&cos_base[rotary_offset])
|
||
: cos_emb;
|
||
sin_emb = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec_RoPE *>(
|
||
&sin_base[rotary_offset])
|
||
: sin_emb;
|
||
apply_rotary_embedding(q, k, cos_emb, sin_emb);
|
||
}
|
||
} else {
|
||
/* old rotary pos emb */
|
||
if (params.rotary_emb_dims != 0) {
|
||
int last_dim = Dh / params.rotary_emb_dims;
|
||
int half_lastdim = last_dim / 2;
|
||
int rotary_offset = bi * Dh + tid * QK_VEC_SIZE;
|
||
const float *cos_base = params.rotary_emb;
|
||
const float *sin_base = params.rotary_emb + params.batch_size * Dh;
|
||
int stride = half_lastdim / QK_VEC_SIZE;
|
||
int stride_all_lastdim = 2 * stride;
|
||
int right_id = tid / stride_all_lastdim * stride_all_lastdim +
|
||
(tid + stride) % (stride_all_lastdim);
|
||
int qk_right_offset = qkv_base_offset + right_id * QK_VEC_SIZE;
|
||
int q_right_bias_offset = hi * Dh + right_id * QK_VEC_SIZE;
|
||
int k_right_bias_offset =
|
||
hi / num_head_per_group * Dh + right_id * QK_VEC_SIZE;
|
||
Qk_vec q_right;
|
||
zero(q_right);
|
||
// q_right =
|
||
// (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh)
|
||
// ? *reinterpret_cast<const Qk_vec *>(&q_base[qk_right_offset])
|
||
// : q_right;
|
||
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
||
load_func.template load<Qk_vec>(q_right, qk_right_offset);
|
||
}
|
||
Qk_vec k_right;
|
||
zero(k_right);
|
||
// k_right =
|
||
// (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh)
|
||
// ? *reinterpret_cast<const Qk_vec *>(&k_base[qk_right_offset])
|
||
// : k_right;
|
||
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
||
load_func.template load<Qk_vec>(k_right,
|
||
params.num_head * Dh +
|
||
qk_right_offset - hi * Dh +
|
||
hi / num_head_per_group * Dh);
|
||
}
|
||
|
||
if (params.add_qkv_bias) {
|
||
Qk_vec q_right_bias;
|
||
zero(q_right_bias);
|
||
q_right_bias = (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec *>(
|
||
&q_bias_base[q_right_bias_offset])
|
||
: q_right_bias;
|
||
Qk_vec k_right_bias;
|
||
zero(k_right_bias);
|
||
k_right_bias = (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec *>(
|
||
&k_bias_base[k_right_bias_offset])
|
||
: k_right_bias;
|
||
|
||
q_right = add(q_right, q_right_bias);
|
||
k_right = add(k_right, k_right_bias);
|
||
}
|
||
|
||
Qk_vec_RoPE cos_emb;
|
||
zero(cos_emb);
|
||
cos_emb = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec_RoPE *>(
|
||
&cos_base[rotary_offset])
|
||
: cos_emb;
|
||
|
||
Qk_vec_RoPE sin_emb;
|
||
zero(sin_emb);
|
||
sin_emb = (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh)
|
||
? *reinterpret_cast<const Qk_vec_RoPE *>(
|
||
&sin_base[rotary_offset])
|
||
: sin_emb;
|
||
float alpha = (tid % stride_all_lastdim) < stride
|
||
? static_cast<float>(-1)
|
||
: static_cast<float>(1);
|
||
q = apply_rotary_emb<Qk_vec, Qk_vec_RoPE>(
|
||
q, q_right, cos_emb, sin_emb, alpha);
|
||
k = apply_rotary_emb<Qk_vec, Qk_vec_RoPE>(
|
||
k, k_right, cos_emb, sin_emb, alpha);
|
||
}
|
||
}
|
||
|
||
*reinterpret_cast<Qk_vec *>(&q_smem[tid * QK_VEC_SIZE]) = q;
|
||
|
||
if (is_last_block) {
|
||
int co = tid / QK_VECS_IN_16B;
|
||
int ci = (tid % QK_VECS_IN_16B) * QK_VEC_SIZE;
|
||
int offset = kv_bhi * params.max_seq_length * Dh +
|
||
co * params.max_seq_length * QK_ELTS_IN_16B +
|
||
act_time_step * QK_ELTS_IN_16B + ci;
|
||
if (Dh == Dh_MAX || co < Dh / QK_ELTS_IN_16B) {
|
||
*reinterpret_cast<Qk_vec *>(¶ms.cache_kv[offset]) = k;
|
||
}
|
||
|
||
qk = dot<Qk_vec, Qk_vec>(q, k);
|
||
// QK_VECS_PER_WARP is <= WARP_SIZE, reduce it within a warp!
|
||
if (QK_VECS_PER_WARP <= WARP_SIZE) {
|
||
#pragma unroll
|
||
for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) {
|
||
qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// when QK_VECS_PER_WARP > WARP_SIZE, we need to reduce the qk in smem!
|
||
if (QK_VECS_PER_WARP > WARP_SIZE && is_last_block) {
|
||
constexpr int WARPS_PER_RED =
|
||
(QK_VECS_PER_WARP + WARP_SIZE - 1) / WARP_SIZE;
|
||
qk = block_sum<WARPS_PER_RED>(&red_smem[WARPS_PER_RED], qk);
|
||
}
|
||
|
||
// Let only the last cuda ThreadBlock compute the final q*k.
|
||
if (tid == 0 && is_last_block) {
|
||
// NOTE(wangxi): mask must be 0.0
|
||
// T mask = params.attn_mask[
|
||
// bi * (params.timestep + 1) + params.timestep];
|
||
// qk += static_cast<float>(mask);
|
||
qk *= params.inv_sqrt_dh;
|
||
if (params.attn_mask) {
|
||
auto mask_bhi = params.mask_broadcast_num_heads ? bi : bhi;
|
||
T mask =
|
||
params.attn_mask[mask_bhi * (params.timestep + 1) + act_time_step];
|
||
qk += static_cast<float>(mask);
|
||
}
|
||
qk_max = qk;
|
||
qk_smem[act_time_step - start_seq] = qk;
|
||
}
|
||
__syncthreads();
|
||
|
||
using K_vec = typename K_vec_<T, THREADS_PER_KEY>::Type;
|
||
constexpr int K_VEC_SIZE = sizeof(K_vec) / sizeof(T);
|
||
static_assert(Dh_MAX % K_VEC_SIZE == 0, "");
|
||
constexpr int K_ELTS_PER_THREAD = Dh_MAX / THREADS_PER_KEY;
|
||
constexpr int K_VECS_PER_THREAD = K_ELTS_PER_THREAD / K_VEC_SIZE;
|
||
|
||
int ko = tid / THREADS_PER_KEY + start_seq;
|
||
int ki = (tid % THREADS_PER_KEY) * K_VEC_SIZE;
|
||
|
||
static_assert(Dh_MAX == THREADS_PER_KEY * K_VEC_SIZE * K_VECS_PER_THREAD, "");
|
||
|
||
K_vec q[K_VECS_PER_THREAD];
|
||
#pragma unroll
|
||
for (int i = 0; i < K_VECS_PER_THREAD; ++i) {
|
||
q[i] = *reinterpret_cast<const K_vec *>(
|
||
&q_smem[ki + i * THREADS_PER_KEY * K_VEC_SIZE]);
|
||
}
|
||
|
||
constexpr int K_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_KEY;
|
||
constexpr int K_PER_WARP = WARP_SIZE / THREADS_PER_KEY;
|
||
|
||
T *k_cache = ¶ms.cache_kv[kv_bhi * params.max_seq_length * Dh + ki];
|
||
T *k_cache_batch = ¶ms.cache_kv[bbhi * params.max_seq_length * Dh + ki];
|
||
int ti_end = div_up(curr_seq_section, K_PER_WARP) * K_PER_WARP + start_seq;
|
||
|
||
const int *beam_offsets = params.beam_cache_offset
|
||
? ¶ms.beam_cache_offset[bi_seq_len_offset]
|
||
: nullptr;
|
||
|
||
#pragma unroll
|
||
for (int ti = ko; ti < ti_end; ti += K_PER_ITER) {
|
||
const int beam_offset = beam_offsets ? beam_offsets[ti] * params.num_head *
|
||
params.max_seq_length * Dh
|
||
: 0;
|
||
K_vec k[K_VECS_PER_THREAD];
|
||
K_vec k_vec_zero;
|
||
zero(k_vec_zero);
|
||
#pragma unroll
|
||
for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) {
|
||
int jj = ii * params.max_seq_length + ti;
|
||
if (ti < end_seq) {
|
||
if (beam_offset) {
|
||
k[ii] =
|
||
(Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.max_seq_length)
|
||
? *reinterpret_cast<const K_vec *>(
|
||
&k_cache_batch[beam_offset + jj * QK_ELTS_IN_16B])
|
||
: k_vec_zero;
|
||
} else {
|
||
k[ii] =
|
||
(Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.max_seq_length)
|
||
? *reinterpret_cast<const K_vec *>(
|
||
&k_cache[jj * QK_ELTS_IN_16B])
|
||
: k_vec_zero;
|
||
}
|
||
}
|
||
}
|
||
|
||
// NOTE(liyurui): We should multiple q with inv_sqrt_dh first, for dot(q, k)
|
||
// may overflow with FP16 in large model.
|
||
float qk = Qk_dot<T, THREADS_PER_KEY>::dot(q, k, params.inv_sqrt_dh);
|
||
|
||
// bool is_mask = false;
|
||
if (ti < end_seq && tid % THREADS_PER_KEY == 0) {
|
||
// qk_max = is_mask ? qk_max : fmaxf(qk_max, qk);
|
||
auto mask_bhi = params.mask_broadcast_num_heads ? bi : bhi;
|
||
// T mask = params.attn_mask[mask_bhi * (params.timestep + 1) + ti];
|
||
if (params.attn_mask) {
|
||
T mask = params.attn_mask[mask_bhi * (params.timestep + 1) + ti];
|
||
qk += static_cast<float>(mask);
|
||
}
|
||
qk_max = fmaxf(qk_max, qk);
|
||
|
||
qk_smem[ti - start_seq] = qk;
|
||
}
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int mask = WARP_SIZE / 2; mask >= THREADS_PER_KEY; mask /= 2) {
|
||
qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask));
|
||
}
|
||
|
||
const int warp = tid / WARP_SIZE;
|
||
const int lane = tid % WARP_SIZE;
|
||
|
||
if (lane == 0) {
|
||
red_smem[warp] = qk_max;
|
||
}
|
||
|
||
__syncthreads();
|
||
|
||
qk_max = lane < WARPS_PER_BLOCK ? red_smem[lane] : -FLT_MAX;
|
||
#pragma unroll
|
||
for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) {
|
||
qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask));
|
||
}
|
||
|
||
qk_max = __shfl_sync(uint32_t(-1), qk_max, 0);
|
||
|
||
int useful_smem_index =
|
||
is_last_block ? curr_seq_section : curr_seq_section - 1;
|
||
float sum = 0.f;
|
||
for (int ti = tid; ti <= useful_smem_index; ti += THREADS_PER_BLOCK) {
|
||
// bool is_mask = false;
|
||
// float logit = is_mask ? 0.f : __expf(qk_smem[ti] - qk_max);
|
||
float logit = __expf(qk_smem[ti] - qk_max);
|
||
sum += logit;
|
||
qk_smem[ti] = logit;
|
||
}
|
||
|
||
sum = block_sum<WARPS_PER_BLOCK>(&red_smem[WARPS_PER_BLOCK], sum);
|
||
|
||
int bhsi = bhi * params.split_seq;
|
||
if (SPLIT && tid == 0) {
|
||
float2 sum_max = {sum, qk_max};
|
||
*reinterpret_cast<float2 *>(
|
||
¶ms.qk_sum_max_split_seq[(bhsi + split_index) * 2]) = sum_max;
|
||
}
|
||
|
||
// FIXME(wangxi): need add 1.e-6f?
|
||
float inv_sum = __fdividef(1.f, sum + 1.e-6f);
|
||
|
||
for (int ti = tid; ti <= useful_smem_index; ti += THREADS_PER_BLOCK) {
|
||
convert_from_float(logits_smem[ti], qk_smem[ti] * inv_sum);
|
||
}
|
||
|
||
__syncthreads();
|
||
|
||
constexpr int V_VEC_SIZE = Dh_MAX / THREADS_PER_VALUE;
|
||
using V_vec = typename V_vec_<T, V_VEC_SIZE>::Type;
|
||
|
||
// now we have got [1, seq] ，distributed in logits_smem.
|
||
// next we compute [1, seq] * [seq, head_dim] = [1, head_dim]
|
||
// THREADS_PER_VALUE means num of threads per value's head_dim.
|
||
// we split the seq dimension for more cuda threads to compute.
|
||
// vo means the first seq index processed by this cuda thread in the value.
|
||
// vi means the head_dim index processed by this cuda thread in the value.
|
||
// so this cuda thread compute [1, k] * [k, vi:vi+V_VEC_SIZE] and k starts
|
||
// from vo and increases by a step V_PER_ITER.
|
||
int vo = tid / THREADS_PER_VALUE + start_seq;
|
||
int vi = (tid % THREADS_PER_VALUE) * V_VEC_SIZE;
|
||
|
||
T *v_cache = ¶ms.cache_kv[params.cache_batch_size * kv_num_head *
|
||
params.max_seq_length * Dh +
|
||
kv_bhi * params.max_seq_length * Dh + vi];
|
||
T *v_cache_batch = ¶ms.cache_kv[params.batch_size * params.num_head *
|
||
params.max_seq_length * Dh +
|
||
bbhi * params.max_seq_length * Dh + vi];
|
||
|
||
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
||
using V_vec_acum = typename V_vec_acum_fp32_<V_vec>::Type;
|
||
#else
|
||
using V_vec_acum = V_vec;
|
||
#endif
|
||
|
||
V_vec_acum out;
|
||
zero(out);
|
||
// V_PER_ITER is used to strip-mined the seq dimension.
|
||
constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE;
|
||
if (Dh == Dh_MAX || vi < Dh) {
|
||
#pragma unroll
|
||
for (int ti = vo; ti < end_seq; ti += V_PER_ITER) {
|
||
const int beam_offset =
|
||
beam_offsets
|
||
? beam_offsets[ti] * params.num_head * params.max_seq_length * Dh
|
||
: 0;
|
||
V_vec v;
|
||
if (beam_offset) {
|
||
v = *reinterpret_cast<const V_vec *>(
|
||
&v_cache_batch[beam_offset + ti * Dh]);
|
||
} else {
|
||
v = *reinterpret_cast<const V_vec *>(&v_cache[ti * Dh]);
|
||
}
|
||
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
||
float logit = logits_smem[ti - start_seq];
|
||
out = fma(logit, cast_to_float(v), out);
|
||
#else
|
||
DataType_ logit = static_cast<DataType_>(logits_smem[ti - start_seq]);
|
||
// Update the partial sums.
|
||
out = fma(logit, v, out);
|
||
#endif
|
||
}
|
||
}
|
||
|
||
V_vec v_bias;
|
||
zero(v_bias);
|
||
// now we process the last v.
|
||
if (vo == (act_time_step % V_PER_ITER + start_seq) &&
|
||
(Dh == Dh_MAX || vi < Dh) && is_last_block) {
|
||
// V_vec v = *reinterpret_cast<const V_vec *>(
|
||
// ¶ms.qkv[2 * params.num_head * Dh + qkv_base_offset + vi]);
|
||
V_vec v;
|
||
load_func.template load<V_vec>(v,
|
||
qkv_base_offset + vi - hi * Dh +
|
||
params.num_head * Dh + kv_num_head * Dh +
|
||
hi / num_head_per_group * Dh);
|
||
if (params.add_qkv_bias) {
|
||
v_bias = *reinterpret_cast<const V_vec *>(
|
||
¶ms.qkv_bias[(kv_num_head + params.num_head) * Dh +
|
||
hi / num_head_per_group * Dh + vi]);
|
||
v = add(v, v_bias);
|
||
}
|
||
|
||
*reinterpret_cast<V_vec *>(&v_cache[act_time_step * Dh]) = v;
|
||
|
||
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
||
out = fma(logits_smem[act_time_step - start_seq], cast_to_float(v), out);
|
||
#else
|
||
out = fma(logits_smem[act_time_step - start_seq], v, out);
|
||
#endif
|
||
}
|
||
|
||
__syncthreads();
|
||
|
||
// now we do the reduction in the seq dimension to get [1, head_dim].
|
||
if (Dh == Dh_MAX || vi < Dh) {
|
||
int vo_blk = vo - start_seq; // vo id of current block
|
||
#pragma unroll
|
||
for (int active_groups = V_PER_ITER; active_groups >= 2;
|
||
active_groups /= 2) {
|
||
int midpoint = active_groups / 2;
|
||
if (vo_blk >= midpoint && vo_blk < active_groups &&
|
||
(Dh == Dh_MAX || vi < Dh)) {
|
||
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
||
convert_from_float(*reinterpret_cast<V_vec *>(
|
||
&out_smem[(vo_blk - midpoint) * Dh + vi]),
|
||
out);
|
||
#else
|
||
*reinterpret_cast<V_vec *>(&out_smem[(vo_blk - midpoint) * Dh + vi]) =
|
||
out;
|
||
#endif
|
||
}
|
||
__syncthreads();
|
||
if (vo_blk < midpoint && (Dh == Dh_MAX || vi < Dh)) {
|
||
out = add(*reinterpret_cast<const V_vec *>(&out_smem[vo_blk * Dh + vi]),
|
||
out);
|
||
}
|
||
__syncthreads();
|
||
}
|
||
}
|
||
|
||
if (vo == start_seq && (Dh == Dh_MAX || vi < Dh)) {
|
||
if (SPLIT && real_split_each_batch > 1) {
|
||
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
||
*(reinterpret_cast<V_vec_acum *>(
|
||
¶ms.split_out[(bhsi + split_index) * Dh + vi])) = out;
|
||
#else
|
||
*(reinterpret_cast<V_vec_acum_fp32_<V_vec>::Type *>(
|
||
¶ms.split_out[(bhsi + split_index) * Dh + vi])) =
|
||
cast_to_float(out);
|
||
#endif
|
||
} else {
|
||
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
||
V_vec tmp_out;
|
||
convert_from_float(tmp_out, out);
|
||
store_func.template store<V_vec>(tmp_out, bhi * Dh + vi);
|
||
#else
|
||
store_func.template store<V_vec>(out, bhi * Dh + vi);
|
||
#endif
|
||
}
|
||
}
|
||
|
||
#else
|
||
assert(false);
|
||
#endif
|
||
}
|
||
|
||
template <typename T, int Dh, int Dh_MAX, typename StoreFunc>
|
||
__global__ void post_process_kernel(Masked_multihead_attention_params<T> params,
|
||
StoreFunc store_func) {
|
||
const int bi = blockIdx.y;
|
||
int act_time_step = params.sequence_lengths == nullptr
|
||
? params.timestep
|
||
: params.sequence_lengths[bi];
|
||
int real_split_each_batch = (act_time_step - 1) / params.steps_per_block + 1;
|
||
if (real_split_each_batch <= 1) {
|
||
return;
|
||
}
|
||
|
||
const int tid = threadIdx.x;
|
||
const int hi = blockIdx.x;
|
||
const int bhi = (bi * params.num_head + hi);
|
||
const int bhsi = (bi * params.num_head + hi) * params.split_seq;
|
||
extern __shared__ float2 qk_sum_max_smem[];
|
||
|
||
for (int i = tid; i < real_split_each_batch; i += blockDim.x) {
|
||
qk_sum_max_smem[i] = *reinterpret_cast<float2 *>(
|
||
¶ms.qk_sum_max_split_seq[(bhsi + i) * 2]);
|
||
}
|
||
__syncthreads();
|
||
|
||
float max = -FLT_MAX;
|
||
float sum = 0;
|
||
float v = 0;
|
||
if (tid < Dh) {
|
||
#pragma unroll
|
||
for (int i = 0; i < real_split_each_batch; ++i) {
|
||
float2 sum_max = qk_sum_max_smem[i];
|
||
float tmp_max = sum_max.y;
|
||
max = tmp_max > max ? tmp_max : max;
|
||
}
|
||
#pragma unroll
|
||
for (int i = 0; i < real_split_each_batch; ++i) {
|
||
float2 sum_max = qk_sum_max_smem[i];
|
||
// split_out:[bsz , num_head, split_seq, dim_head]
|
||
float this_v = params.split_out[(bhsi + i) * Dh + tid];
|
||
|
||
float real_this_sum = sum_max.x * __expf(sum_max.y - max);
|
||
v += real_this_sum * this_v;
|
||
sum += real_this_sum;
|
||
}
|
||
|
||
v /= sum;
|
||
T tmp_v = (T)v;
|
||
store_func.template store<T>(tmp_v, bhi * Dh + tid);
|
||
// params.out[bhi * Dh + tid] = (T)(v);
|
||
}
|
||
}
|
||
|
||
template <typename T, bool SPLIT>
|
||
inline size_t smem_size_in_bytes(
|
||
const Masked_multihead_attention_params<T> ¶ms,
|
||
int dim_head,
|
||
int threads_per_value,
|
||
int threads_per_block) {
|
||
// for qk_smem and logits_smem(both float)
|
||
size_t qk_sz = div_up(params.timestep, 4) * 16;
|
||
if (SPLIT) {
|
||
qk_sz = div_up(params.steps_per_block, 4) * 16;
|
||
}
|
||
// for reduce (logits dot V) result
|
||
int rows_per_red = threads_per_block / threads_per_value;
|
||
size_t red_sz = rows_per_red * dim_head * sizeof(T) / 2;
|
||
return max(qk_sz, red_sz);
|
||
}
|
||
|
||
#define MMHA_LAUNCH_KERNEL(T, \
|
||
Dh, \
|
||
Dh_MAX, \
|
||
THDS_PER_KEY, \
|
||
THDS_PER_VALUE, \
|
||
THDS_PER_BLOCK, \
|
||
stream, \
|
||
load_func, \
|
||
store_func) \
|
||
size_t smem_sz_size = smem_size_in_bytes<T, SPLIT>( \
|
||
params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
|
||
PADDLE_ENFORCE_LE_INT_MAX(smem_sz_size, \
|
||
"masked_multihead_attention shared memory size"); \
|
||
int smem_sz = static_cast<int>(smem_sz_size); \
|
||
constexpr auto kernel_fn = \
|
||
masked_multihead_attention_kernel<T, \
|
||
Dh, \
|
||
Dh_MAX, \
|
||
THDS_PER_KEY, \
|
||
THDS_PER_VALUE, \
|
||
THDS_PER_BLOCK, \
|
||
decltype(load_func), \
|
||
decltype(store_func), \
|
||
SPLIT>; \
|
||
if (smem_sz > 0xc000) { \
|
||
cudaFuncSetAttribute( \
|
||
kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz); \
|
||
} \
|
||
dim3 grid(params.split_seq, params.num_head, params.batch_size); \
|
||
kernel_fn<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>( \
|
||
params, load_func, store_func)
|
||
|
||
template <typename T,
|
||
int Dh,
|
||
int Dh_MAX,
|
||
typename LoadFunc,
|
||
typename StoreFunc,
|
||
bool SPLIT>
|
||
void fmha_launch_kernel(const Masked_multihead_attention_params<T> ¶ms,
|
||
const cudaStream_t &stream,
|
||
LoadFunc load_func,
|
||
StoreFunc store_func) {
|
||
constexpr int THREADS_PER_VALUE = Dh_MAX * sizeof(T) / 16;
|
||
// If try adjusting the hyperparam, THDS_PER_KEY can try [1, 2, 4]
|
||
// for LLM: multiBatch(8)/longSeq(>2048) case, reduce THDS_PER_KEY may work
|
||
// for super longSeq(>3072) case, larger steps_per_block（256） may work
|
||
if constexpr (SPLIT) {
|
||
MMHA_LAUNCH_KERNEL(T,
|
||
Dh,
|
||
Dh_MAX,
|
||
4,
|
||
THREADS_PER_VALUE,
|
||
128,
|
||
stream,
|
||
load_func,
|
||
store_func);
|
||
} else {
|
||
if (params.timestep < 32) {
|
||
MMHA_LAUNCH_KERNEL(T,
|
||
Dh,
|
||
Dh_MAX,
|
||
4,
|
||
THREADS_PER_VALUE,
|
||
64,
|
||
stream,
|
||
load_func,
|
||
store_func);
|
||
} else if (params.timestep < 2048) {
|
||
#if defined(MMHA_USE_HMMA_FOR_REDUCTION) && defined(__CUDA_ARCH__) && \
|
||
__CUDA_ARCH__ >= 750
|
||
MMHA_LAUNCH_KERNEL(T,
|
||
Dh,
|
||
Dh_MAX,
|
||
4,
|
||
THREADS_PER_VALUE,
|
||
256,
|
||
stream,
|
||
load_func,
|
||
store_func);
|
||
#else
|
||
MMHA_LAUNCH_KERNEL(T,
|
||
Dh,
|
||
Dh_MAX,
|
||
2,
|
||
THREADS_PER_VALUE,
|
||
128,
|
||
stream,
|
||
load_func,
|
||
store_func);
|
||
#endif
|
||
} else {
|
||
MMHA_LAUNCH_KERNEL(T,
|
||
Dh,
|
||
Dh_MAX,
|
||
1,
|
||
THREADS_PER_VALUE,
|
||
256,
|
||
stream,
|
||
load_func,
|
||
store_func);
|
||
}
|
||
}
|
||
}
|
||
|
||
#define FMHA_LAUNCH_KERNEL(dim_head_, dim_head_max_, stream) \
|
||
case dim_head_: \
|
||
fmha_launch_kernel<T, \
|
||
dim_head_, \
|
||
dim_head_max_, \
|
||
decltype(load_func), \
|
||
decltype(store_func), \
|
||
SPLIT>(params, stream, load_func, store_func); \
|
||
if (SPLIT) { \
|
||
post_process_kernel<T, dim_head_, dim_head_max_> \
|
||
<<<grid, dim_head_max_, smem_sz, stream>>>(params, store_func); \
|
||
} \
|
||
break;
|
||
|
||
template <typename T, typename LoadFunc, typename StoreFunc, bool SPLIT>
|
||
void fmha_impl(const GPUContext &dev_ctx,
|
||
const Masked_multihead_attention_params<T> ¶ms,
|
||
int dim_head,
|
||
LoadFunc load_func,
|
||
StoreFunc store_func) {
|
||
dim3 grid(params.num_head, params.batch_size);
|
||
int smem_sz = params.split_seq * sizeof(float2);
|
||
auto stream = dev_ctx.stream();
|
||
switch (dim_head) {
|
||
FMHA_LAUNCH_KERNEL(16, 32, stream)
|
||
FMHA_LAUNCH_KERNEL(32, 32, stream)
|
||
FMHA_LAUNCH_KERNEL(64, 64, stream)
|
||
FMHA_LAUNCH_KERNEL(80, 128, stream)
|
||
FMHA_LAUNCH_KERNEL(96, 128, stream)
|
||
FMHA_LAUNCH_KERNEL(128, 128, stream)
|
||
FMHA_LAUNCH_KERNEL(192, 256, stream)
|
||
default:
|
||
PADDLE_THROW(common::errors::Unimplemented(
|
||
"Dim_head = %d is unsupported!", dim_head));
|
||
}
|
||
}
|
||
|
||
template <typename T, bool SPLIT = false>
|
||
void DispatchFMHA(const GPUContext &dev_ctx,
|
||
const DenseTensor &qkv_tensor,
|
||
const Masked_multihead_attention_params<T> ¶ms,
|
||
int num_head,
|
||
int dim_head,
|
||
DenseTensor *out_tensor,
|
||
const DenseTensor *dequant_qkv_scales = nullptr,
|
||
const float quant_fmha_out_scale = -1,
|
||
const int quant_round_type = 1,
|
||
const float quant_max_bound = 127.0f,
|
||
const float quant_min_bound = -127.0f) {
|
||
if (dequant_qkv_scales != nullptr && quant_fmha_out_scale > 0) {
|
||
MMHALoad<T, int32_t> load_func(qkv_tensor.data<int32_t>(),
|
||
dequant_qkv_scales->data<float>(),
|
||
3 * num_head * dim_head);
|
||
MMHAStore<T, int8_t> store_func(out_tensor->data<int8_t>(),
|
||
quant_round_type,
|
||
quant_fmha_out_scale,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else if (dequant_qkv_scales == nullptr && quant_fmha_out_scale > 0) {
|
||
MMHALoad<T> load_func(qkv_tensor.data<T>());
|
||
MMHAStore<T, int8_t> store_func(out_tensor->data<int8_t>(),
|
||
quant_round_type,
|
||
quant_fmha_out_scale,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else if (dequant_qkv_scales != nullptr && quant_fmha_out_scale <= 0) {
|
||
MMHALoad<T, int32_t> load_func(qkv_tensor.data<int32_t>(),
|
||
dequant_qkv_scales->data<float>(),
|
||
3 * num_head * dim_head);
|
||
MMHAStore<T> store_func(out_tensor->data<T>());
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else {
|
||
MMHALoad<T> load_func(qkv_tensor.data<T>());
|
||
MMHAStore<T> store_func(out_tensor->data<T>());
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
}
|
||
}
|
||
|
||
template <typename T, bool SPLIT = false>
|
||
void DispatchFMHA(const GPUContext &dev_ctx,
|
||
const DenseTensor &qkv_tensor,
|
||
const DenseTensor &shift,
|
||
const DenseTensor &smooth,
|
||
const Masked_multihead_attention_params<T> ¶ms,
|
||
int num_head,
|
||
int dim_head,
|
||
DenseTensor *out_tensor,
|
||
const DenseTensor *dequant_qkv_scales = nullptr,
|
||
const float quant_fmha_out_scale = -1,
|
||
const int quant_round_type = 1,
|
||
const float quant_max_bound = 127.0f,
|
||
const float quant_min_bound = -127.0f) {
|
||
if (dequant_qkv_scales != nullptr && quant_fmha_out_scale > 0) {
|
||
MMHALoad<T, int32_t> load_func(qkv_tensor.data<int32_t>(),
|
||
dequant_qkv_scales->data<float>(),
|
||
3 * num_head * dim_head);
|
||
MMHAStore<T, int8_t, true> store_func(out_tensor->data<int8_t>(),
|
||
shift.data<T>(),
|
||
smooth.data<T>(),
|
||
num_head * dim_head,
|
||
quant_round_type,
|
||
quant_fmha_out_scale,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else if (dequant_qkv_scales == nullptr && quant_fmha_out_scale > 0) {
|
||
MMHALoad<T> load_func(qkv_tensor.data<T>());
|
||
MMHAStore<T, int8_t, true> store_func(out_tensor->data<int8_t>(),
|
||
shift.data<T>(),
|
||
smooth.data<T>(),
|
||
num_head * dim_head,
|
||
quant_round_type,
|
||
quant_fmha_out_scale,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else if (dequant_qkv_scales != nullptr && quant_fmha_out_scale <= 0) {
|
||
MMHALoad<T, int32_t> load_func(qkv_tensor.data<int32_t>(),
|
||
dequant_qkv_scales->data<float>(),
|
||
3 * num_head * dim_head);
|
||
MMHAStore<T, T, true> store_func(out_tensor->data<T>(),
|
||
shift.data<T>(),
|
||
smooth.data<T>(),
|
||
num_head * dim_head);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
} else {
|
||
MMHALoad<T> load_func(qkv_tensor.data<T>());
|
||
MMHAStore<T, T, true> store_func(out_tensor->data<T>(),
|
||
shift.data<T>(),
|
||
smooth.data<T>(),
|
||
num_head * dim_head);
|
||
fmha_impl<T, decltype(load_func), decltype(store_func), SPLIT>(
|
||
dev_ctx, params, dim_head, load_func, store_func);
|
||
}
|
||
}
|
||
|
||
struct NormalVersion {};
|
||
struct UnusedVersion {};
|
||
|
||
template <typename T>
|
||
struct DispatchDtypeTrait {
|
||
using FuncVersion = NormalVersion;
|
||
};
|
||
|
||
template <>
|
||
struct DispatchDtypeTrait<int32_t> {
|
||
using FuncVersion = UnusedVersion;
|
||
};
|
||
|
||
template <typename T, typename Context>
|
||
void DispatchWithDtype(const Context &dev_ctx,
|
||
const DenseTensor &x,
|
||
const DenseTensor &cache_kv,
|
||
const optional<DenseTensor> &bias,
|
||
const optional<DenseTensor> &src_mask,
|
||
const optional<DenseTensor> &cum_offsets,
|
||
const optional<DenseTensor> &sequence_lengths,
|
||
const optional<DenseTensor> &rotary_tensor,
|
||
const optional<DenseTensor> &beam_cache_offset,
|
||
const optional<DenseTensor> &qkv_out_scale,
|
||
const optional<DenseTensor> &out_shift,
|
||
const optional<DenseTensor> &out_smooth,
|
||
int seq_len,
|
||
int rotary_emb_dims,
|
||
const bool use_neox_rotary_style,
|
||
const float out_scale,
|
||
const int quant_round_type,
|
||
const float quant_max_bound,
|
||
const float quant_min_bound,
|
||
DenseTensor *out,
|
||
DenseTensor *cache_kv_out,
|
||
DenseTensor *beam_cache_offset_out,
|
||
NormalVersion) {
|
||
const auto &x_dims = x.dims();
|
||
int bsz = x_dims[0];
|
||
int64_t cache_bsz = cache_kv.dims()[1];
|
||
// TODO(large-tensor): downstream functors may still use int
|
||
|
||
int64_t max_seq_len = cache_kv.dims()[3];
|
||
// TODO(large-tensor): downstream functors may still use int
|
||
|
||
int64_t dim_head = cache_kv.dims()[4];
|
||
// TODO(large-tensor): downstream functors may still use int
|
||
|
||
int timestep = max_seq_len;
|
||
float inv_sqrt_dh = 1. / sqrt(dim_head);
|
||
|
||
int64_t k_num_head = cache_kv.dims()[2];
|
||
// TODO(large-tensor): downstream functors may still use int
|
||
|
||
int v_num_head = k_num_head;
|
||
// this num_head means query's head
|
||
int num_head =
|
||
x.dims()[x.dims().size() - 1] / dim_head - k_num_head - v_num_head;
|
||
|
||
Masked_multihead_attention_params<T> params;
|
||
|
||
bool mask_broadcast_num_heads = true;
|
||
|
||
params.add_qkv_bias = false;
|
||
if (bias) {
|
||
params.add_qkv_bias = true;
|
||
params.qkv_bias = const_cast<T *>(bias->data<T>());
|
||
}
|
||
|
||
if (src_mask) {
|
||
if (src_mask->dims()[1] == 1) {
|
||
mask_broadcast_num_heads = true;
|
||
} else if (src_mask->dims()[1] == num_head) {
|
||
mask_broadcast_num_heads = false;
|
||
} else {
|
||
PADDLE_THROW(errors::InvalidArgument(
|
||
"Unknown dimension for attn_mask, the num_head(2nd) "
|
||
"dimension is invalid, it should be 1 or num_head(%d), "
|
||
"but got %d",
|
||
num_head,
|
||
src_mask->dims()[1]));
|
||
}
|
||
params.attn_mask = src_mask->data<T>();
|
||
timestep = src_mask->dims()[3] - 1;
|
||
}
|
||
|
||
if (out_scale > 0) {
|
||
dev_ctx.template Alloc<int8_t>(out);
|
||
} else {
|
||
dev_ctx.template Alloc<T>(out);
|
||
}
|
||
|
||
if (sequence_lengths) {
|
||
params.sequence_lengths = sequence_lengths->data<int>();
|
||
}
|
||
|
||
if (cum_offsets) {
|
||
PADDLE_THROW(common::errors::PermissionDenied(
|
||
"Current mmha kernel does not support cum_offsets param."));
|
||
}
|
||
|
||
if (rotary_emb_dims > 0) {
|
||
params.rotary_emb = rotary_tensor->data<float>();
|
||
} else {
|
||
params.rotary_emb = nullptr;
|
||
}
|
||
|
||
if (beam_cache_offset) {
|
||
params.beam_cache_offset = beam_cache_offset->data<int>();
|
||
params.beam_width = beam_cache_offset->dims()[1];
|
||
}
|
||
|
||
params.mask_broadcast_num_heads = mask_broadcast_num_heads;
|
||
params.cache_kv = const_cast<T *>(cache_kv_out->data<T>());
|
||
params.neox_rotary_style = use_neox_rotary_style;
|
||
params.batch_size = bsz;
|
||
params.cache_batch_size = cache_bsz;
|
||
params.num_head = num_head;
|
||
params.kv_num_head = k_num_head;
|
||
params.timestep = timestep;
|
||
params.max_seq_length = max_seq_len;
|
||
params.inv_sqrt_dh = inv_sqrt_dh;
|
||
params.rotary_emb_dims = rotary_emb_dims;
|
||
|
||
params.steps_per_block = timestep; // if not SPLIT, this is useless.
|
||
params.split_seq = 1; // if not SPLIT, grid.x==1
|
||
|
||
bool SPLIT = false;
|
||
if (bsz <= 4 && timestep >= 512) {
|
||
SPLIT = true;
|
||
}
|
||
if (SPLIT) {
|
||
const int steps_per_block = 128;
|
||
params.steps_per_block = steps_per_block;
|
||
params.split_seq = (timestep - 1) / steps_per_block + 1;
|
||
int split_seq = params.split_seq;
|
||
|
||
DenseTensor qk_sum_max_split_seq;
|
||
// 2 means sum and max.
|
||
qk_sum_max_split_seq.Resize({bsz, num_head, split_seq, 2});
|
||
dev_ctx.template Alloc<float>(&qk_sum_max_split_seq,
|
||
qk_sum_max_split_seq.numel() * sizeof(float));
|
||
params.qk_sum_max_split_seq = qk_sum_max_split_seq.data<float>();
|
||
|
||
DenseTensor split_out;
|
||
split_out.Resize({bsz, num_head, split_seq, dim_head});
|
||
dev_ctx.template Alloc<float>(&split_out,
|
||
split_out.numel() * sizeof(float));
|
||
params.split_out = split_out.data<float>();
|
||
|
||
if (out_shift) {
|
||
DispatchFMHA<T, true>(dev_ctx,
|
||
x,
|
||
*(out_shift.get_ptr()),
|
||
*(out_smooth.get_ptr()),
|
||
params,
|
||
num_head,
|
||
dim_head,
|
||
out,
|
||
qkv_out_scale.get_ptr(),
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
DispatchFMHA<T, true>(dev_ctx,
|
||
x,
|
||
params,
|
||
num_head,
|
||
dim_head,
|
||
out,
|
||
qkv_out_scale.get_ptr(),
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
}
|
||
} else {
|
||
if (out_shift) {
|
||
DispatchFMHA<T, false>(dev_ctx,
|
||
x,
|
||
*(out_shift.get_ptr()),
|
||
*(out_smooth.get_ptr()),
|
||
params,
|
||
num_head,
|
||
dim_head,
|
||
out,
|
||
qkv_out_scale.get_ptr(),
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
DispatchFMHA<T, false>(dev_ctx,
|
||
x,
|
||
params,
|
||
num_head,
|
||
dim_head,
|
||
out,
|
||
qkv_out_scale.get_ptr(),
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, typename Context>
|
||
void DispatchWithDtype(const Context &dev_ctx,
|
||
const DenseTensor &x,
|
||
const DenseTensor &cache_kv,
|
||
const optional<DenseTensor> &bias,
|
||
const optional<DenseTensor> &src_mask,
|
||
const optional<DenseTensor> &cum_offsets,
|
||
const optional<DenseTensor> &sequence_lengths,
|
||
const optional<DenseTensor> &rotary_tensor,
|
||
const optional<DenseTensor> &beam_cache_offset,
|
||
const optional<DenseTensor> &qkv_out_scale,
|
||
const optional<DenseTensor> &out_shift,
|
||
const optional<DenseTensor> &out_smooth,
|
||
int seq_len,
|
||
int rotary_emb_dims,
|
||
const bool use_neox_rotary_style,
|
||
const float out_scale,
|
||
const int quant_round_type,
|
||
const float quant_max_bound,
|
||
const float quant_min_bound,
|
||
DenseTensor *out,
|
||
DenseTensor *cache_kv_out,
|
||
DenseTensor *beam_cache_offset_out,
|
||
UnusedVersion) {}
|
||
|
||
#endif // PADDLE_WITH_HIP
|
||
|
||
template <typename T, typename Context>
|
||
void MMHAKernel(const Context &dev_ctx,
|
||
const DenseTensor &x,
|
||
const DenseTensor &cache_kv,
|
||
const optional<DenseTensor> &bias,
|
||
const optional<DenseTensor> &src_mask,
|
||
const optional<DenseTensor> &cum_offsets,
|
||
const optional<DenseTensor> &sequence_lengths,
|
||
const optional<DenseTensor> &rotary_tensor,
|
||
const optional<DenseTensor> &beam_cache_offset,
|
||
const optional<DenseTensor> &qkv_out_scale,
|
||
const optional<DenseTensor> &out_shift,
|
||
const optional<DenseTensor> &out_smooth,
|
||
int seq_len,
|
||
int rotary_emb_dims,
|
||
const bool use_neox_rotary_style,
|
||
const std::string &compute_dtype,
|
||
const float out_scale,
|
||
const int quant_round_type,
|
||
const float quant_max_bound,
|
||
const float quant_min_bound,
|
||
DenseTensor *out,
|
||
DenseTensor *cache_kv_out,
|
||
DenseTensor *beam_cache_offset_out) {
|
||
#ifndef PADDLE_WITH_HIP
|
||
if (x.dtype() == phi::DataType::INT32) {
|
||
switch (str2int(compute_dtype.c_str())) {
|
||
case str2int("fp16"):
|
||
DispatchWithDtype<phi::float16, Context>(
|
||
dev_ctx,
|
||
x,
|
||
cache_kv,
|
||
bias,
|
||
src_mask,
|
||
cum_offsets,
|
||
sequence_lengths,
|
||
rotary_tensor,
|
||
beam_cache_offset,
|
||
qkv_out_scale,
|
||
out_shift,
|
||
out_smooth,
|
||
seq_len,
|
||
rotary_emb_dims,
|
||
use_neox_rotary_style,
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound,
|
||
out,
|
||
cache_kv_out,
|
||
beam_cache_offset_out,
|
||
typename DispatchDtypeTrait<phi::float16>::FuncVersion{});
|
||
break;
|
||
#if CUDA_VERSION >= 11000
|
||
case str2int("bf16"):
|
||
DispatchWithDtype<phi::bfloat16, Context>(
|
||
dev_ctx,
|
||
x,
|
||
cache_kv,
|
||
bias,
|
||
src_mask,
|
||
cum_offsets,
|
||
sequence_lengths,
|
||
rotary_tensor,
|
||
beam_cache_offset,
|
||
qkv_out_scale,
|
||
out_shift,
|
||
out_smooth,
|
||
seq_len,
|
||
rotary_emb_dims,
|
||
use_neox_rotary_style,
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound,
|
||
out,
|
||
cache_kv_out,
|
||
beam_cache_offset_out,
|
||
typename DispatchDtypeTrait<phi::bfloat16>::FuncVersion{});
|
||
break;
|
||
#endif
|
||
case str2int("fp32"):
|
||
DispatchWithDtype<float, Context>(
|
||
dev_ctx,
|
||
x,
|
||
cache_kv,
|
||
bias,
|
||
src_mask,
|
||
cum_offsets,
|
||
sequence_lengths,
|
||
rotary_tensor,
|
||
beam_cache_offset,
|
||
qkv_out_scale,
|
||
out_shift,
|
||
out_smooth,
|
||
seq_len,
|
||
rotary_emb_dims,
|
||
use_neox_rotary_style,
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound,
|
||
out,
|
||
cache_kv_out,
|
||
beam_cache_offset_out,
|
||
typename DispatchDtypeTrait<float>::FuncVersion{});
|
||
break;
|
||
default:
|
||
PADDLE_THROW(common::errors::InvalidArgument(
|
||
"In the case of quantization enabled with Input(x) INT32, "
|
||
"Attr(compute_dtype) must be set in (bf16, fp16, fp32), "
|
||
"but get compute_dtype (%s)",
|
||
compute_dtype));
|
||
}
|
||
} else {
|
||
DispatchWithDtype<T, Context>(
|
||
dev_ctx,
|
||
x,
|
||
cache_kv,
|
||
bias,
|
||
src_mask,
|
||
cum_offsets,
|
||
sequence_lengths,
|
||
rotary_tensor,
|
||
beam_cache_offset,
|
||
qkv_out_scale,
|
||
out_shift,
|
||
out_smooth,
|
||
seq_len,
|
||
rotary_emb_dims,
|
||
use_neox_rotary_style,
|
||
out_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound,
|
||
out,
|
||
cache_kv_out,
|
||
beam_cache_offset_out,
|
||
typename DispatchDtypeTrait<T>::FuncVersion{});
|
||
}
|
||
#endif // PADDLE_WITH_HIP
|
||
}
|
||
|
||
} // namespace fusion
|
||
} // namespace phi
|
||
|
||
#if CUDA_VERSION >= 11000
|
||
PD_REGISTER_KERNEL(masked_multihead_attention,
|
||
GPU,
|
||
ALL_LAYOUT,
|
||
phi::fusion::MMHAKernel,
|
||
float,
|
||
phi::float16,
|
||
phi::bfloat16,
|
||
int32_t) {}
|
||
#else
|
||
PD_REGISTER_KERNEL(masked_multihead_attention,
|
||
GPU,
|
||
ALL_LAYOUT,
|
||
phi::fusion::MMHAKernel,
|
||
float,
|
||
phi::float16,
|
||
int32_t) {}
|
||
#endif
|