wehub-resource-sync
4559 lines
178 KiB
C++
4559 lines
178 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/common/flags.h"
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/kernels/funcs/quant_dequant.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/mmha_util.cu.h"
|
|
|
|
COMMON_DECLARE_bool(use_xqa_optim);
|
|
COMMON_DECLARE_bool(blha_use_fp32_qk_sum);
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#define GPU(str) hip##str
|
|
#define GPUMultiProcessorCount hipDeviceAttributeMultiprocessorCount
|
|
#else
|
|
#define GPU(str) cuda##str
|
|
#define GPUMultiProcessorCount cudaDevAttrMultiProcessorCount
|
|
#endif
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
constexpr int VEC_16B = 16;
|
|
|
|
template <typename T>
|
|
struct Block_AttN_params {
|
|
// output buffer, [B, 1(seq_len), q_num_head * dim_head]
|
|
T *out;
|
|
// qkv_out, [B, 1(seq_len), (2 * kv_num_head + q_num_head) * dim_head]
|
|
const T *qkv;
|
|
// bias, [(2 * kv_num_head + q_num_head), dim_head]
|
|
const T *qkv_bias;
|
|
// [bsz, seq_len]
|
|
const int *cum_offsets;
|
|
// TODO(wangxi): optimize with input_lengths and max_input_len?
|
|
// [bsz, 1, 1, time_step(cache_seq_length)+1]
|
|
const T *attn_mask;
|
|
|
|
// mask_length is the 3th dimension of attn_mask.
|
|
int mask_length;
|
|
bool mask_broadcast_num_heads;
|
|
|
|
// k_cache [max_block_num, num_head, block_size, head_size]
|
|
// v_cache [max_block_num, num_head, block_size, head_size]
|
|
T *k_cache;
|
|
T *v_cache;
|
|
|
|
uint8_t *k_cache_I;
|
|
uint8_t *v_cache_I;
|
|
|
|
const int *block_tables;
|
|
|
|
const int *sequence_lengths{nullptr};
|
|
|
|
const float *rotary_emb = nullptr;
|
|
int rotary_emb_dims;
|
|
int rope_stride;
|
|
float inv_compression_ratio = 1.0f; // Default as 1.0
|
|
|
|
int batch_size; // batch * beam
|
|
int beam_width = 1;
|
|
int max_num_blocks_per_seq;
|
|
int block_size;
|
|
int q_num_head;
|
|
int kv_num_head;
|
|
int gqa_num_per_partitions;
|
|
int timestep;
|
|
int seq_len;
|
|
|
|
int pre_cache_length = 0;
|
|
|
|
// 1.f / sqrt(Dh)
|
|
float inv_sqrt_dh;
|
|
|
|
bool add_qkv_bias;
|
|
bool neox_rotary_style;
|
|
|
|
const float *cache_k_quant_scales = nullptr;
|
|
const float *cache_v_quant_scales = nullptr;
|
|
const float *cache_k_dequant_scales = nullptr;
|
|
const float *cache_v_dequant_scales = nullptr;
|
|
|
|
float rope_theta = 10000.0f;
|
|
};
|
|
|
|
template <typename T,
|
|
typename SUM_T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int THREADS_PER_KEY,
|
|
int THREADS_PER_VALUE,
|
|
int THREADS_PER_BLOCK,
|
|
int BLOCK_SIZE,
|
|
CacheType CACHE_TYPE,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
__global__ __launch_bounds__(THREADS_PER_BLOCK) void block_attention_kernel(
|
|
Block_AttN_params<T> params, LoadFunc load_func, StoreFunc store_func) {
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
const int bi = blockIdx.y;
|
|
int act_time_step = params.sequence_lengths[bi];
|
|
if (act_time_step == 0) {
|
|
return;
|
|
}
|
|
|
|
act_time_step += params.pre_cache_length;
|
|
|
|
const int *block_table =
|
|
params.block_tables + bi * params.max_num_blocks_per_seq;
|
|
|
|
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_[];
|
|
|
|
int block_smem_offset =
|
|
div_up(params.max_num_blocks_per_seq, 4) * 4 * sizeof(int);
|
|
float *qk_smem = reinterpret_cast<float *>(smem_ + block_smem_offset);
|
|
|
|
float *logits_smem = reinterpret_cast<float *>(smem_ + block_smem_offset);
|
|
|
|
/*
|
|
Note: This shared memory is used in Final Reduce, so we do not
|
|
need to add offset. Since in Final Reduce Phase, the block tables is not used.
|
|
*/
|
|
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_sum_type = typename Qk_vec_<SUM_T, Dh_MAX>::Type;
|
|
using Qk_vec_RoPE = typename Qk_vec_RoPE_<T, float, Dh_MAX>::Type;
|
|
using QK_Packed_Int8_t = typename Packed_Int8_<Qk_vec, CACHE_TYPE>::Type;
|
|
|
|
// Each block has a head's q value
|
|
__shared__ __align__(sizeof(Qk_vec)) T q_smem[Dh_MAX];
|
|
|
|
const int tid = threadIdx.x;
|
|
const int hi = blockIdx.x; // head index
|
|
const int kv_hi = hi / params.gqa_num_per_partitions;
|
|
|
|
float k_quant_scale;
|
|
float v_quant_scale;
|
|
float k_dequant_scale;
|
|
float v_dequant_scale;
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
k_quant_scale = static_cast<float>(params.cache_k_quant_scales[kv_hi]);
|
|
v_quant_scale = static_cast<float>(params.cache_v_quant_scales[kv_hi]);
|
|
k_dequant_scale = static_cast<float>(params.cache_k_dequant_scales[kv_hi]);
|
|
v_dequant_scale = static_cast<float>(params.cache_v_dequant_scales[kv_hi]);
|
|
}
|
|
|
|
const int bhi = bi * params.q_num_head + hi;
|
|
const int ti =
|
|
params.cum_offsets ? bi * params.seq_len - params.cum_offsets[bi] : -1;
|
|
const int thi = params.cum_offsets ? ti * params.q_num_head + hi : -1;
|
|
int *block_table_smem = reinterpret_cast<int *>(smem_);
|
|
|
|
for (int local_id = tid; local_id < params.max_num_blocks_per_seq;
|
|
local_id += blockDim.x) {
|
|
block_table_smem[local_id] = block_table[local_id];
|
|
}
|
|
__syncthreads();
|
|
|
|
float qk_max = -FLT_MAX;
|
|
float qk = 0;
|
|
|
|
const int block_idx = act_time_step / BLOCK_SIZE;
|
|
const int block_offset = act_time_step % BLOCK_SIZE;
|
|
const int physical_block_number = block_table_smem[block_idx];
|
|
|
|
// cache offset of current token
|
|
const int base_cache_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh;
|
|
|
|
// qkv [B, S=1, num_head + 2 * (kv_num_head), head_dim]
|
|
int qkv_base_offset = bi * (params.q_num_head + params.kv_num_head * 2) * Dh;
|
|
|
|
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_bias_base = nullptr;
|
|
const T *k_bias_base = nullptr;
|
|
|
|
if (params.add_qkv_bias) {
|
|
q_bias_base = params.qkv_bias;
|
|
k_bias_base = params.qkv_bias + params.q_num_head * Dh;
|
|
}
|
|
// Load the current timestep's Q and K, then compute q*k,
|
|
// with each block computing one head.
|
|
if (tid < QK_VECS_PER_WARP) {
|
|
const int qk_offset = qkv_base_offset + tid * QK_VEC_SIZE;
|
|
|
|
Qk_vec q;
|
|
zero(q);
|
|
if (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(q, qk_offset + hi * Dh);
|
|
}
|
|
|
|
Qk_vec k;
|
|
zero(k);
|
|
if (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(
|
|
k, params.q_num_head * Dh + qk_offset + kv_hi * Dh);
|
|
}
|
|
|
|
if (params.add_qkv_bias) {
|
|
const int q_bias_offset = hi * Dh + tid * QK_VEC_SIZE;
|
|
const int k_bias_offset = kv_hi * Dh + tid * QK_VEC_SIZE;
|
|
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);
|
|
k = add(k, k_bias);
|
|
}
|
|
|
|
if (params.rotary_emb_dims != 0) {
|
|
if (!params.neox_rotary_style) {
|
|
apply_rotary_embedding(q,
|
|
k,
|
|
tid,
|
|
Dh,
|
|
act_time_step - params.pre_cache_length,
|
|
params.inv_compression_ratio,
|
|
params.rope_theta);
|
|
} else {
|
|
int last_dim = Dh / params.rotary_emb_dims;
|
|
int half_lastdim = last_dim / 2;
|
|
int rotary_offset = act_time_step * Dh + tid * QK_VEC_SIZE;
|
|
const float *cos_base = params.rotary_emb;
|
|
const float *sin_base = params.rotary_emb + params.rope_stride;
|
|
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 q_right_offset = qkv_base_offset + hi * Dh + right_id * QK_VEC_SIZE;
|
|
int k_right_offset = qkv_base_offset + params.q_num_head * Dh +
|
|
kv_hi * Dh + right_id * QK_VEC_SIZE;
|
|
|
|
Qk_vec q_right;
|
|
zero(q_right);
|
|
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(q_right, q_right_offset);
|
|
}
|
|
Qk_vec k_right;
|
|
zero(k_right);
|
|
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(k_right, k_right_offset);
|
|
}
|
|
|
|
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 (Dh == Dh_MAX || tid * QK_VEC_SIZE < Dh) {
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
const int offset = base_cache_offset + tid * QK_VEC_SIZE;
|
|
QK_Packed_Int8_t k_tmp =
|
|
round_tmp<QK_Packed_Int8_t, Qk_vec, CACHE_TYPE>(
|
|
mul<Qk_vec, float, Qk_vec>(k_quant_scale, k));
|
|
*reinterpret_cast<QK_Packed_Int8_t *>(¶ms.k_cache_I[offset]) =
|
|
k_tmp;
|
|
} else {
|
|
const int offset = base_cache_offset + tid * QK_VEC_SIZE;
|
|
*reinterpret_cast<Qk_vec *>(¶ms.k_cache[offset]) = k;
|
|
}
|
|
}
|
|
|
|
qk = dot<Qk_sum_type, Qk_vec>(q, k);
|
|
|
|
if (QK_VECS_PER_WARP <= WARP_SIZE) {
|
|
#pragma unroll
|
|
for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk += __shfl_xor(qk, mask);
|
|
#else
|
|
qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
if (QK_VECS_PER_WARP > WARP_SIZE) {
|
|
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);
|
|
}
|
|
if (tid == 0) {
|
|
qk *= params.inv_sqrt_dh;
|
|
if (params.attn_mask) {
|
|
auto mask_bhi = bhi;
|
|
if (params.mask_broadcast_num_heads) {
|
|
mask_bhi = bi;
|
|
}
|
|
T mask = params.attn_mask[mask_bhi * params.mask_length + act_time_step];
|
|
qk += static_cast<float>(mask);
|
|
}
|
|
qk_max = qk;
|
|
qk_smem[act_time_step] = qk;
|
|
}
|
|
__syncthreads();
|
|
using K_vec = typename K_vec_bttn_<T>::Type;
|
|
using K_vec_I = typename K_vec_I_bttn_<T, CACHE_TYPE>::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;
|
|
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]);
|
|
}
|
|
|
|
// threads in one block can process 'K_PER_ITER' keys
|
|
constexpr int K_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_KEY;
|
|
// threads in one warp can process 'K_PER_ITER' keys
|
|
constexpr int K_PER_WARP = WARP_SIZE / THREADS_PER_KEY;
|
|
|
|
int ti_end = div_up(act_time_step, K_PER_WARP) * K_PER_WARP;
|
|
|
|
K_vec k[K_VECS_PER_THREAD];
|
|
K_vec k_vec_zero;
|
|
zero(k_vec_zero);
|
|
// ti is the timestep index
|
|
for (int ti = ko; ti < ti_end; ti += K_PER_ITER) {
|
|
const int physical_block_number = block_table_smem[ti / BLOCK_SIZE];
|
|
const int block_offset = ti % BLOCK_SIZE;
|
|
const int k_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh + ki;
|
|
#pragma unroll
|
|
for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) {
|
|
if (ti < act_time_step) {
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
mul_pointer_v2<K_vec, float, K_vec_I, CACHE_TYPE>(
|
|
&k[ii],
|
|
k_dequant_scale,
|
|
reinterpret_cast<K_vec_I *>(params.k_cache_I + k_offset +
|
|
ii * THREADS_PER_KEY * K_VEC_SIZE));
|
|
} else {
|
|
k[ii] = (Dh == Dh_MAX || ii * THREADS_PER_KEY * K_VEC_SIZE < Dh)
|
|
? *reinterpret_cast<const K_vec *>(
|
|
params.k_cache + k_offset +
|
|
ii * THREADS_PER_KEY * K_VEC_SIZE)
|
|
: k_vec_zero;
|
|
}
|
|
} else {
|
|
k[ii] = k_vec_zero;
|
|
}
|
|
}
|
|
|
|
float qk = Qk_dot<T, THREADS_PER_KEY>::dot(q, k, params.inv_sqrt_dh);
|
|
|
|
if (params.attn_mask) {
|
|
auto mask_bhi = bhi;
|
|
if (params.mask_broadcast_num_heads) {
|
|
mask_bhi = bi;
|
|
}
|
|
T mask = params.attn_mask[mask_bhi * params.mask_length + ti];
|
|
qk += static_cast<float>(mask);
|
|
}
|
|
|
|
if (ti < act_time_step && tid % THREADS_PER_KEY == 0) {
|
|
qk_max = fmaxf(qk_max, qk);
|
|
qk_smem[ti] = qk;
|
|
}
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int mask = WARP_SIZE / 2; mask >= THREADS_PER_KEY; mask /= 2) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask));
|
|
#else
|
|
qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask));
|
|
#endif
|
|
}
|
|
|
|
const int warp = tid / WARP_SIZE;
|
|
const int lane = tid % WARP_SIZE;
|
|
|
|
if (lane == 0) {
|
|
red_smem[warp] = qk_max;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_max = -FLT_MAX;
|
|
qk_max = red_smem[lane];
|
|
#else
|
|
qk_max = lane < WARPS_PER_BLOCK ? red_smem[lane] : -FLT_MAX;
|
|
#endif
|
|
|
|
#pragma unroll
|
|
for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask));
|
|
#else
|
|
qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask));
|
|
#endif
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_max = __shfl(qk_max, 0);
|
|
#else
|
|
qk_max = __shfl_sync(uint32_t(-1), qk_max, 0);
|
|
#endif
|
|
|
|
float sum = 0.f;
|
|
for (int ti = tid; ti <= act_time_step; ti += THREADS_PER_BLOCK) {
|
|
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);
|
|
|
|
float inv_sum = __fdividef(1.f, sum + 1.e-6f);
|
|
|
|
for (int ti = tid; ti <= act_time_step; 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;
|
|
using V_Packed_Int8_t = typename Packed_Int8_<V_vec, CACHE_TYPE>::Type;
|
|
|
|
int vo = tid / THREADS_PER_VALUE;
|
|
int vi = (tid % THREADS_PER_VALUE) * V_VEC_SIZE;
|
|
|
|
#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);
|
|
|
|
constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE;
|
|
if (Dh == Dh_MAX || vi < Dh) {
|
|
for (int ti = vo; ti < act_time_step; ti += V_PER_ITER) {
|
|
int physical_block_number = block_table_smem[ti / BLOCK_SIZE];
|
|
|
|
const int block_offset = ti % BLOCK_SIZE;
|
|
const int v_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh + vi;
|
|
V_vec v;
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
mul_pointer_v2<V_vec, float, V_Packed_Int8_t, CACHE_TYPE>(
|
|
&v,
|
|
v_dequant_scale,
|
|
reinterpret_cast<V_Packed_Int8_t *>(params.v_cache_I + v_offset));
|
|
} else {
|
|
v = *reinterpret_cast<const V_vec *>(params.v_cache + v_offset);
|
|
}
|
|
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
|
float logit = logits_smem[ti];
|
|
out = fma(logit, cast_to_float(v), out);
|
|
#else
|
|
DataType_ logit = static_cast<DataType_>(logits_smem[ti]);
|
|
// Update the partial sums.
|
|
out = fma(logit, v, out);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
V_vec v_bias;
|
|
zero(v_bias);
|
|
if (vo == (act_time_step % V_PER_ITER) && (Dh == Dh_MAX || vi < Dh)) {
|
|
V_vec v;
|
|
load_func.template load<V_vec>(
|
|
v,
|
|
(params.q_num_head + params.kv_num_head) * Dh + qkv_base_offset +
|
|
kv_hi * Dh + vi);
|
|
if (params.add_qkv_bias) {
|
|
v_bias = *reinterpret_cast<const V_vec *>(
|
|
¶ms.qkv_bias[(params.q_num_head + params.kv_num_head) * Dh +
|
|
kv_hi * Dh + vi]);
|
|
v = add(v, v_bias);
|
|
}
|
|
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
V_Packed_Int8_t v_tmp = round_tmp<V_Packed_Int8_t, V_vec, CACHE_TYPE>(
|
|
mul<V_vec, float, V_vec>(v_quant_scale, v));
|
|
*reinterpret_cast<V_Packed_Int8_t *>(params.v_cache_I +
|
|
base_cache_offset + vi) = v_tmp;
|
|
} else {
|
|
*reinterpret_cast<V_vec *>(params.v_cache + base_cache_offset + vi) = v;
|
|
}
|
|
|
|
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
|
out = fma(logits_smem[act_time_step], cast_to_float(v), out);
|
|
#else
|
|
out = fma(logits_smem[act_time_step], v, out);
|
|
#endif
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (Dh == Dh_MAX || vi < Dh) {
|
|
#pragma unroll
|
|
for (int active_groups = V_PER_ITER; active_groups >= 2;
|
|
active_groups /= 2) {
|
|
int midpoint = active_groups / 2;
|
|
|
|
if (vo >= midpoint && vo < active_groups && (Dh == Dh_MAX || vi < Dh)) {
|
|
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
|
convert_from_float(
|
|
*reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]),
|
|
out);
|
|
#else
|
|
*reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]) = out;
|
|
#endif
|
|
}
|
|
__syncthreads();
|
|
if (vo < midpoint && (Dh == Dh_MAX || vi < Dh)) {
|
|
out =
|
|
add(*reinterpret_cast<const V_vec *>(&out_smem[vo * Dh + vi]), out);
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
|
|
if (vo == 0 && (Dh == Dh_MAX || vi < Dh)) {
|
|
#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,
|
|
thi != -1 ? thi * Dh + vi : bhi * Dh + vi);
|
|
#else
|
|
store_func.template store<V_vec>(out,
|
|
thi != -1 ? thi * Dh + vi : bhi * Dh + vi);
|
|
#endif
|
|
}
|
|
#else
|
|
assert(false);
|
|
#endif
|
|
}
|
|
|
|
template <typename T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int THREADS_PER_KEY,
|
|
int THREADS_PER_VALUE,
|
|
int THREADS_PER_BLOCK,
|
|
int BLOCK_SIZE,
|
|
CacheType CACHE_TYPE,
|
|
int GQA_PARTITION_SIZE,
|
|
int GQA_SUB_PARTITION_SIZE,
|
|
int GQA_NUM_SUB_PARTITIONS,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
__global__ __launch_bounds__(THREADS_PER_BLOCK) void gqa_block_attention_kernel(
|
|
Block_AttN_params<T> params, LoadFunc load_func, StoreFunc store_func) {
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
const int bi = blockIdx.y;
|
|
const int act_time_step = params.sequence_lengths[bi];
|
|
if (act_time_step == 0) {
|
|
return;
|
|
}
|
|
|
|
const int *block_table =
|
|
params.block_tables + bi * params.max_num_blocks_per_seq;
|
|
|
|
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_[];
|
|
|
|
int block_smem_offset =
|
|
div_up(params.max_num_blocks_per_seq, 4) * 4 * sizeof(int);
|
|
int q_smem_offset =
|
|
div_up(Dh_MAX * GQA_SUB_PARTITION_SIZE, 4) * 4 * sizeof(T);
|
|
|
|
T *q_smem = reinterpret_cast<T *>(smem_ + block_smem_offset);
|
|
float *qk_smem =
|
|
reinterpret_cast<float *>(smem_ + block_smem_offset + q_smem_offset);
|
|
float *logits_smem =
|
|
reinterpret_cast<float *>(smem_ + block_smem_offset + q_smem_offset);
|
|
|
|
/*
|
|
Note: This shared memory is used in Final Reduce, so we do not
|
|
need to add offset. Since in Final Reduce Phase, the block tables is not used.
|
|
*/
|
|
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;
|
|
using QK_Packed_Int8_t = typename Packed_Int8_<Qk_vec, CACHE_TYPE>::Type;
|
|
|
|
const int tid = threadIdx.x;
|
|
const int kv_hi = blockIdx.x / GQA_NUM_SUB_PARTITIONS;
|
|
const int gqa_sub_partition_id = blockIdx.x % GQA_NUM_SUB_PARTITIONS;
|
|
|
|
float k_quant_scale;
|
|
float v_quant_scale;
|
|
float k_dequant_scale;
|
|
float v_dequant_scale;
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
k_quant_scale = static_cast<float>(params.cache_k_quant_scales[kv_hi]);
|
|
v_quant_scale = static_cast<float>(params.cache_v_quant_scales[kv_hi]);
|
|
k_dequant_scale = static_cast<float>(params.cache_k_dequant_scales[kv_hi]);
|
|
v_dequant_scale = static_cast<float>(params.cache_v_dequant_scales[kv_hi]);
|
|
}
|
|
|
|
const int ti =
|
|
params.cum_offsets ? bi * params.seq_len - params.cum_offsets[bi] : -1;
|
|
|
|
int *block_table_smem = reinterpret_cast<int *>(smem_);
|
|
|
|
for (int local_id = tid; local_id < params.max_num_blocks_per_seq;
|
|
local_id += blockDim.x) {
|
|
block_table_smem[local_id] = block_table[local_id];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
float qk_max = -FLT_MAX;
|
|
float qk = 0;
|
|
|
|
const int block_idx = act_time_step / BLOCK_SIZE;
|
|
const int block_offset = act_time_step % BLOCK_SIZE;
|
|
const int physical_block_number = block_table_smem[block_idx];
|
|
|
|
// cache offset of current token
|
|
const int base_cache_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh;
|
|
|
|
// qkv [B, S=1, num_head + 2 * (kv_num_head), head_dim]
|
|
int qkv_base_offset = bi * (params.q_num_head + params.kv_num_head * 2) * Dh;
|
|
|
|
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_bias_base = nullptr;
|
|
const T *k_bias_base = nullptr;
|
|
|
|
if (params.add_qkv_bias) {
|
|
q_bias_base = params.qkv_bias;
|
|
k_bias_base = params.qkv_bias + params.q_num_head * Dh;
|
|
}
|
|
// Load the current timestep's Q and K, then compute q*k,
|
|
// GQA use one warp calculate one head
|
|
assert(QK_VECS_PER_WARP < WARP_SIZE);
|
|
const int warp_id = tid / WARP_SIZE;
|
|
const int lane_id = tid % WARP_SIZE;
|
|
|
|
if (warp_id < GQA_SUB_PARTITION_SIZE && lane_id < QK_VECS_PER_WARP) {
|
|
const int hi = kv_hi * GQA_PARTITION_SIZE +
|
|
gqa_sub_partition_id * GQA_SUB_PARTITION_SIZE + warp_id;
|
|
const int qk_offset = qkv_base_offset + lane_id * QK_VEC_SIZE;
|
|
|
|
Qk_vec q;
|
|
zero(q);
|
|
if (Dh == Dh_MAX || lane_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(q, qk_offset + hi * Dh);
|
|
}
|
|
|
|
Qk_vec k;
|
|
zero(k);
|
|
if (Dh == Dh_MAX || lane_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(
|
|
k, params.q_num_head * Dh + qk_offset + kv_hi * Dh);
|
|
}
|
|
|
|
if (params.add_qkv_bias) {
|
|
const int q_bias_offset = hi * Dh + lane_id * QK_VEC_SIZE;
|
|
const int k_bias_offset = kv_hi * Dh + lane_id * QK_VEC_SIZE;
|
|
Qk_vec q_bias;
|
|
zero(q_bias);
|
|
Qk_vec k_bias;
|
|
zero(k_bias);
|
|
|
|
q_bias =
|
|
(Dh == Dh_MAX || lane_id * QK_VEC_SIZE < Dh)
|
|
? *reinterpret_cast<const Qk_vec *>(&q_bias_base[q_bias_offset])
|
|
: q_bias;
|
|
k_bias =
|
|
(Dh == Dh_MAX || lane_id * QK_VEC_SIZE < Dh)
|
|
? *reinterpret_cast<const Qk_vec *>(&k_bias_base[k_bias_offset])
|
|
: k_bias;
|
|
|
|
q = add(q, q_bias);
|
|
k = add(k, k_bias);
|
|
}
|
|
|
|
if (params.rotary_emb_dims != 0) {
|
|
if (!params.neox_rotary_style) {
|
|
apply_rotary_embedding(q,
|
|
k,
|
|
lane_id,
|
|
Dh,
|
|
act_time_step,
|
|
params.inv_compression_ratio,
|
|
params.rope_theta);
|
|
} else {
|
|
int last_dim = Dh / params.rotary_emb_dims;
|
|
int half_lastdim = last_dim / 2;
|
|
int rotary_offset = act_time_step * Dh + lane_id * QK_VEC_SIZE;
|
|
const float *cos_base = params.rotary_emb;
|
|
const float *sin_base = params.rotary_emb + params.rope_stride;
|
|
int stride = half_lastdim / QK_VEC_SIZE;
|
|
int stride_all_lastdim = 2 * stride;
|
|
int right_id = lane_id / stride_all_lastdim * stride_all_lastdim +
|
|
(lane_id + stride) % (stride_all_lastdim);
|
|
int q_right_offset = qkv_base_offset + hi * Dh + right_id * QK_VEC_SIZE;
|
|
int k_right_offset = qkv_base_offset + params.q_num_head * Dh +
|
|
kv_hi * Dh + right_id * QK_VEC_SIZE;
|
|
|
|
Qk_vec q_right;
|
|
zero(q_right);
|
|
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(q_right, q_right_offset);
|
|
}
|
|
Qk_vec k_right;
|
|
zero(k_right);
|
|
if (Dh == Dh_MAX || right_id * QK_VEC_SIZE < Dh) {
|
|
load_func.template load<Qk_vec>(k_right, k_right_offset);
|
|
}
|
|
|
|
Qk_vec_RoPE cos_emb;
|
|
zero(cos_emb);
|
|
cos_emb = (Dh == Dh_MAX || lane_id * 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 || lane_id * QK_VEC_SIZE < Dh)
|
|
? *reinterpret_cast<const Qk_vec_RoPE *>(
|
|
&sin_base[rotary_offset])
|
|
: sin_emb;
|
|
float alpha = (lane_id % 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[warp_id * Dh_MAX + lane_id * QK_VEC_SIZE]) = q;
|
|
|
|
if (Dh == Dh_MAX || lane_id * QK_VEC_SIZE < Dh) {
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
const int offset = base_cache_offset + lane_id * QK_VEC_SIZE;
|
|
QK_Packed_Int8_t k_tmp =
|
|
round_tmp<QK_Packed_Int8_t, Qk_vec, CACHE_TYPE>(
|
|
mul<Qk_vec, float, Qk_vec>(k_quant_scale, k));
|
|
*reinterpret_cast<QK_Packed_Int8_t *>(¶ms.k_cache_I[offset]) =
|
|
k_tmp;
|
|
} else {
|
|
const int offset = base_cache_offset + lane_id * QK_VEC_SIZE;
|
|
*reinterpret_cast<Qk_vec *>(¶ms.k_cache[offset]) = k;
|
|
}
|
|
}
|
|
|
|
qk = dot<Qk_vec, Qk_vec>(q, k);
|
|
|
|
if (QK_VECS_PER_WARP <= WARP_SIZE) {
|
|
#pragma unroll
|
|
for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk += __shfl_xor(qk, mask);
|
|
#else
|
|
qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
|
|
if (lane_id == 0) {
|
|
qk *= params.inv_sqrt_dh;
|
|
const int hi = kv_hi * GQA_PARTITION_SIZE +
|
|
gqa_sub_partition_id * GQA_SUB_PARTITION_SIZE + warp_id;
|
|
const int bhi = bi * params.q_num_head + hi;
|
|
if (params.attn_mask) {
|
|
auto mask_bhi = bhi;
|
|
if (params.mask_broadcast_num_heads) {
|
|
mask_bhi = bi;
|
|
}
|
|
T mask = params.attn_mask[mask_bhi * params.mask_length + act_time_step];
|
|
qk += static_cast<float>(mask);
|
|
}
|
|
|
|
qk_max = qk;
|
|
qk_smem[act_time_step * GQA_SUB_PARTITION_SIZE + warp_id] = qk;
|
|
}
|
|
__syncthreads();
|
|
using K_vec = typename K_vec_bttn_<T>::Type;
|
|
using K_vec_I = typename K_vec_I_bttn_<T, CACHE_TYPE>::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;
|
|
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]);
|
|
// }
|
|
|
|
float qk_maxs[GQA_SUB_PARTITION_SIZE];
|
|
#pragma unroll
|
|
for (int i = 0; i < GQA_SUB_PARTITION_SIZE; i++) {
|
|
// qk_maxs[i] = -FLT_MAX;
|
|
// initialize qk_maxs!!!
|
|
qk_maxs[i] = qk_smem[act_time_step * GQA_SUB_PARTITION_SIZE + i];
|
|
}
|
|
|
|
// threads in one block can process 'K_PER_ITER' keys
|
|
constexpr int K_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_KEY;
|
|
// threads in one warp can process 'K_PER_ITER' keys
|
|
constexpr int K_PER_WARP = WARP_SIZE / THREADS_PER_KEY;
|
|
|
|
int ti_end = div_up(act_time_step, K_PER_WARP) * K_PER_WARP;
|
|
|
|
K_vec k[K_VECS_PER_THREAD];
|
|
K_vec k_vec_zero;
|
|
zero(k_vec_zero);
|
|
// ti is the timestep index
|
|
for (int ti = ko; ti < ti_end; ti += K_PER_ITER) {
|
|
int physical_block_number = block_table_smem[ti / BLOCK_SIZE];
|
|
|
|
const int block_offset = ti % BLOCK_SIZE;
|
|
const int k_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh + ki;
|
|
#pragma unroll
|
|
for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) {
|
|
if (ti < act_time_step) {
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
mul_pointer_v2<K_vec, float, K_vec_I, CACHE_TYPE>(
|
|
&k[ii],
|
|
k_dequant_scale,
|
|
reinterpret_cast<K_vec_I *>(params.k_cache_I + k_offset +
|
|
ii * THREADS_PER_KEY * K_VEC_SIZE));
|
|
} else {
|
|
k[ii] = (Dh == Dh_MAX || ii * THREADS_PER_KEY * K_VEC_SIZE < Dh)
|
|
? *reinterpret_cast<const K_vec *>(
|
|
params.k_cache + k_offset +
|
|
ii * THREADS_PER_KEY * K_VEC_SIZE)
|
|
: k_vec_zero;
|
|
}
|
|
} else {
|
|
k[ii] = k_vec_zero;
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
const int hi = kv_hi * GQA_PARTITION_SIZE +
|
|
gqa_sub_partition_id * GQA_SUB_PARTITION_SIZE + local_hi;
|
|
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[local_hi * Dh_MAX + ki + i * THREADS_PER_KEY * K_VEC_SIZE]);
|
|
}
|
|
|
|
float qk = Qk_dot<T, THREADS_PER_KEY>::dot(q, k, params.inv_sqrt_dh);
|
|
|
|
if (params.attn_mask) {
|
|
const int bhi = bi * params.q_num_head + hi;
|
|
auto mask_bhi = bhi;
|
|
if (params.mask_broadcast_num_heads) {
|
|
mask_bhi = bi;
|
|
}
|
|
T mask = params.attn_mask[mask_bhi * params.mask_length + ti];
|
|
qk += static_cast<float>(mask);
|
|
}
|
|
|
|
if (ti < act_time_step && tid % THREADS_PER_KEY == 0) {
|
|
qk_maxs[local_hi] = fmaxf(qk_maxs[local_hi], qk);
|
|
qk_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi] = qk;
|
|
}
|
|
}
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
#pragma unroll
|
|
for (int mask = WARP_SIZE / 2; mask >= THREADS_PER_KEY; mask /= 2) {
|
|
qk_maxs[local_hi] = fmaxf(qk_maxs[local_hi],
|
|
#ifdef PADDLE_WITH_HIP
|
|
__shfl_xor(qk_maxs[local_hi], mask));
|
|
#else
|
|
__shfl_xor_sync(
|
|
uint32_t(-1), qk_maxs[local_hi], mask));
|
|
|
|
#endif
|
|
}
|
|
|
|
if (lane_id == 0) {
|
|
red_smem[warp_id] = qk_maxs[local_hi];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_maxs[local_hi] = -FLT_MAX;
|
|
qk_maxs[local_hi] = red_smem[lane_id];
|
|
#else
|
|
qk_maxs[local_hi] =
|
|
lane_id < WARPS_PER_BLOCK ? red_smem[lane_id] : -FLT_MAX;
|
|
#endif
|
|
|
|
#pragma unroll
|
|
for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) {
|
|
qk_maxs[local_hi] = fmaxf(qk_maxs[local_hi],
|
|
#ifdef PADDLE_WITH_HIP
|
|
__shfl_xor(qk_maxs[local_hi], mask));
|
|
#else
|
|
__shfl_xor_sync(
|
|
uint32_t(-1), qk_maxs[local_hi], mask));
|
|
#endif
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
qk_maxs[local_hi] = __shfl(qk_maxs[local_hi], 0);
|
|
#else
|
|
qk_maxs[local_hi] = __shfl_sync(uint32_t(-1), qk_maxs[local_hi], 0);
|
|
#endif
|
|
|
|
float sum = 0.f;
|
|
for (int ti = tid; ti <= act_time_step; ti += THREADS_PER_BLOCK) {
|
|
float logit = __expf(qk_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi] -
|
|
qk_maxs[local_hi]);
|
|
sum += logit;
|
|
qk_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi] = logit;
|
|
}
|
|
|
|
sum = block_sum<WARPS_PER_BLOCK>(&red_smem[WARPS_PER_BLOCK], sum);
|
|
|
|
float inv_sum = __fdividef(1.f, sum + 1.e-6f);
|
|
|
|
for (int ti = tid; ti <= act_time_step; ti += THREADS_PER_BLOCK) {
|
|
convert_from_float(
|
|
logits_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi],
|
|
qk_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi] * inv_sum);
|
|
}
|
|
} // for (int local_hi
|
|
__syncthreads();
|
|
|
|
constexpr int V_VEC_SIZE = Dh_MAX / THREADS_PER_VALUE;
|
|
using V_vec = typename V_vec_<T, V_VEC_SIZE>::Type;
|
|
using V_Packed_Int8_t = typename Packed_Int8_<V_vec, CACHE_TYPE>::Type;
|
|
|
|
int vo = tid / THREADS_PER_VALUE;
|
|
int vi = (tid % THREADS_PER_VALUE) * V_VEC_SIZE;
|
|
|
|
#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[GQA_SUB_PARTITION_SIZE];
|
|
#pragma unroll
|
|
for (int i = 0; i < GQA_SUB_PARTITION_SIZE; i++) {
|
|
zero(out[i]);
|
|
}
|
|
|
|
constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE;
|
|
if (Dh == Dh_MAX || vi < Dh) {
|
|
for (int ti = vo; ti < act_time_step; ti += V_PER_ITER) {
|
|
int physical_block_number = block_table_smem[ti / BLOCK_SIZE];
|
|
|
|
const int block_offset = ti % BLOCK_SIZE;
|
|
const int v_offset =
|
|
physical_block_number * params.kv_num_head * BLOCK_SIZE * Dh +
|
|
kv_hi * BLOCK_SIZE * Dh + block_offset * Dh + vi;
|
|
V_vec v;
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
mul_pointer_v2<V_vec, float, V_Packed_Int8_t, CACHE_TYPE>(
|
|
&v,
|
|
v_dequant_scale,
|
|
reinterpret_cast<V_Packed_Int8_t *>(params.v_cache_I + v_offset));
|
|
} else {
|
|
v = *reinterpret_cast<const V_vec *>(params.v_cache + v_offset);
|
|
}
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
|
float logit = logits_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi];
|
|
out[local_hi] = fma(logit, cast_to_float(v), out[local_hi]);
|
|
#else
|
|
DataType_ logit = static_cast<DataType_>(
|
|
logits_smem[ti * GQA_SUB_PARTITION_SIZE + local_hi]);
|
|
// Update the partial sums.
|
|
out[local_hi] = fma(logit, v, out[local_hi]);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
|
|
V_vec v_bias;
|
|
zero(v_bias);
|
|
if (vo == (act_time_step % V_PER_ITER) && (Dh == Dh_MAX || vi < Dh)) {
|
|
V_vec v;
|
|
// Load the current v into the v cache
|
|
load_func.template load<V_vec>(
|
|
v,
|
|
(params.q_num_head + params.kv_num_head) * Dh + qkv_base_offset +
|
|
kv_hi * Dh + vi);
|
|
if (params.add_qkv_bias) {
|
|
v_bias = *reinterpret_cast<const V_vec *>(
|
|
¶ms.qkv_bias[(params.q_num_head + params.kv_num_head) * Dh +
|
|
kv_hi * Dh + vi]);
|
|
v = add(v, v_bias);
|
|
}
|
|
|
|
if (CACHE_TYPE == CacheType::INT8) {
|
|
V_Packed_Int8_t v_tmp = round_tmp<V_Packed_Int8_t, V_vec, CACHE_TYPE>(
|
|
mul<V_vec, float, V_vec>(v_quant_scale, v));
|
|
*reinterpret_cast<V_Packed_Int8_t *>(params.v_cache_I +
|
|
base_cache_offset + vi) = v_tmp;
|
|
} else {
|
|
*reinterpret_cast<V_vec *>(params.v_cache + base_cache_offset + vi) = v;
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
|
|
out[local_hi] =
|
|
fma(logits_smem[act_time_step * GQA_SUB_PARTITION_SIZE + local_hi],
|
|
cast_to_float(v),
|
|
out[local_hi]);
|
|
#else
|
|
out[local_hi] =
|
|
fma(logits_smem[act_time_step * GQA_SUB_PARTITION_SIZE + local_hi],
|
|
v,
|
|
out[local_hi]);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (Dh == Dh_MAX || vi < Dh) {
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
#pragma unroll
|
|
for (int active_groups = V_PER_ITER; active_groups >= 2;
|
|
active_groups /= 2) {
|
|
int midpoint = active_groups / 2;
|
|
|
|
if (vo >= midpoint && vo < active_groups && (Dh == Dh_MAX || vi < Dh)) {
|
|
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
|
convert_from_float(
|
|
*reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]),
|
|
out[local_hi]);
|
|
#else
|
|
*reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]) =
|
|
out[local_hi];
|
|
#endif
|
|
}
|
|
__syncthreads();
|
|
if (vo < midpoint && (Dh == Dh_MAX || vi < Dh)) {
|
|
out[local_hi] =
|
|
add(*reinterpret_cast<const V_vec *>(&out_smem[vo * Dh + vi]),
|
|
out[local_hi]);
|
|
}
|
|
__syncthreads();
|
|
}
|
|
} // for (int local_hi = 0
|
|
}
|
|
|
|
if (vo == 0 && (Dh == Dh_MAX || vi < Dh)) {
|
|
#pragma unroll
|
|
for (int local_hi = 0; local_hi < GQA_SUB_PARTITION_SIZE; local_hi++) {
|
|
const int hi = kv_hi * GQA_PARTITION_SIZE +
|
|
gqa_sub_partition_id * GQA_SUB_PARTITION_SIZE + local_hi;
|
|
const int ti = params.cum_offsets
|
|
? bi * params.seq_len - params.cum_offsets[bi]
|
|
: -1;
|
|
const int thi = params.cum_offsets ? ti * params.q_num_head + hi : -1;
|
|
const int bhi = bi * params.q_num_head + hi;
|
|
#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
|
|
V_vec tmp_out;
|
|
convert_from_float(tmp_out, out[local_hi]);
|
|
store_func.template store<V_vec>(
|
|
tmp_out, thi != -1 ? thi * Dh + vi : bhi * Dh + vi);
|
|
#else
|
|
store_func.template store<V_vec>(
|
|
out[local_hi], thi != -1 ? thi * Dh + vi : bhi * Dh + vi);
|
|
#endif
|
|
}
|
|
}
|
|
#else
|
|
assert(false);
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
inline size_t smem_size_in_bytes(const Block_AttN_params<T> ¶ms,
|
|
int dim_head,
|
|
int threads_per_value,
|
|
int threads_per_block) {
|
|
size_t qk_sz = div_up(params.timestep + 1, 4) * 4 * sizeof(int);
|
|
|
|
size_t block_table_sz =
|
|
div_up(params.max_num_blocks_per_seq, 4) * 4 * sizeof(int);
|
|
size_t cache_scale_sz = 0;
|
|
|
|
size_t logits_table_sz = qk_sz + block_table_sz + cache_scale_sz;
|
|
|
|
int rows_per_red = threads_per_block / threads_per_value;
|
|
|
|
// Note: Because half threads in ThreadBlock use out_smem to do reduce.
|
|
size_t red_sz = rows_per_red * dim_head * sizeof(T) / 2;
|
|
|
|
return max(logits_table_sz, red_sz);
|
|
}
|
|
|
|
template <typename T, int GQA_SUB_PARTITION_SIZE>
|
|
inline size_t gqa_smem_size_in_bytes(const Block_AttN_params<T> ¶ms,
|
|
int dim_head,
|
|
int threads_per_value,
|
|
int threads_per_block) {
|
|
size_t qk_sz = div_up(dim_head * GQA_SUB_PARTITION_SIZE +
|
|
params.timestep * GQA_SUB_PARTITION_SIZE + 1,
|
|
4) *
|
|
4 * sizeof(int);
|
|
size_t block_table_sz =
|
|
div_up(params.max_num_blocks_per_seq, 4) * 4 * sizeof(int);
|
|
size_t logits_table_sz = qk_sz + block_table_sz;
|
|
|
|
int rows_per_red = threads_per_block / threads_per_value;
|
|
size_t red_sz = rows_per_red * dim_head * sizeof(T) / 2;
|
|
|
|
return max(logits_table_sz, red_sz);
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#define BLHAG_LAUNCH_KERNEL(T, \
|
|
SUM_T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
stream, \
|
|
load_func, \
|
|
store_func) \
|
|
size_t smem_sz = \
|
|
smem_size_in_bytes<T>(params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
|
|
constexpr auto kernel_fn = block_attention_kernel<T, \
|
|
SUM_T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
decltype(load_func), \
|
|
decltype(store_func)>; \
|
|
if (smem_sz > 0xc000) { \
|
|
hipError_t result = \
|
|
hipFuncSetAttribute((const void *)kernel_fn, \
|
|
hipFuncAttributeMaxDynamicSharedMemorySize, \
|
|
smem_sz); \
|
|
if (result != hipSuccess) { \
|
|
result = hipGetLastError(); \
|
|
PADDLE_THROW(::common::errors::Unavailable( \
|
|
" hipFuncSetAttribute() returned error %s", \
|
|
hipGetErrorString(result))); \
|
|
} \
|
|
} \
|
|
dim3 grid(params.q_num_head, params.batch_size); \
|
|
kernel_fn<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>( \
|
|
params, load_func, store_func);
|
|
|
|
// For GQA specific optimization (use xqa)
|
|
#define BLHA_LAUNCH_GQA_KERNEL(T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
GQA_PARTITION_SIZE, \
|
|
GQA_SUB_PARTITION_SIZE, \
|
|
stream, \
|
|
load_func, \
|
|
store_func) \
|
|
size_t smem_sz = gqa_smem_size_in_bytes<T, GQA_SUB_PARTITION_SIZE>( \
|
|
params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
|
|
constexpr int GQA_NUM_SUB_PARTITIONS = \
|
|
GQA_PARTITION_SIZE / GQA_SUB_PARTITION_SIZE; \
|
|
constexpr auto kernel_fn = \
|
|
gqa_block_attention_kernel<T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
GQA_PARTITION_SIZE, \
|
|
GQA_SUB_PARTITION_SIZE, \
|
|
GQA_NUM_SUB_PARTITIONS, \
|
|
decltype(load_func), \
|
|
decltype(store_func)>; \
|
|
if (smem_sz > 0xc000) { \
|
|
hipError_t result = \
|
|
hipFuncSetAttribute((const void *)kernel_fn, \
|
|
hipFuncAttributeMaxDynamicSharedMemorySize, \
|
|
smem_sz); \
|
|
if (result != hipSuccess) { \
|
|
result = hipGetLastError(); \
|
|
PADDLE_THROW(::common::errors::Unavailable( \
|
|
" hipFuncSetAttribute() returned error %s", \
|
|
hipGetErrorString(result))); \
|
|
} \
|
|
} \
|
|
dim3 grid(params.kv_num_head *GQA_NUM_SUB_PARTITIONS, params.batch_size); \
|
|
kernel_fn<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>( \
|
|
params, load_func, store_func);
|
|
#else
|
|
#define BLHAG_LAUNCH_KERNEL(T, \
|
|
SUM_T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
stream, \
|
|
load_func, \
|
|
store_func) \
|
|
size_t smem_sz = \
|
|
smem_size_in_bytes<T>(params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
|
|
constexpr auto kernel_fn = block_attention_kernel<T, \
|
|
SUM_T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
decltype(load_func), \
|
|
decltype(store_func)>; \
|
|
if (smem_sz > 0xc000) { \
|
|
cudaError_t result = cudaFuncSetAttribute( \
|
|
kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz); \
|
|
if (result != cudaSuccess) { \
|
|
result = cudaGetLastError(); \
|
|
PADDLE_THROW(::common::errors::Unavailable( \
|
|
" cudaFuncSetAttribute() returned error %s", \
|
|
cudaGetErrorString(result))); \
|
|
} \
|
|
} \
|
|
dim3 grid(params.q_num_head, params.batch_size); \
|
|
kernel_fn<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>( \
|
|
params, load_func, store_func);
|
|
|
|
// For GQA specific optimization (use xqa)
|
|
#define BLHA_LAUNCH_GQA_KERNEL(T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
GQA_PARTITION_SIZE, \
|
|
GQA_SUB_PARTITION_SIZE, \
|
|
stream, \
|
|
load_func, \
|
|
store_func) \
|
|
size_t smem_sz = gqa_smem_size_in_bytes<T, GQA_SUB_PARTITION_SIZE>( \
|
|
params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
|
|
constexpr int GQA_NUM_SUB_PARTITIONS = \
|
|
GQA_PARTITION_SIZE / GQA_SUB_PARTITION_SIZE; \
|
|
constexpr auto kernel_fn = \
|
|
gqa_block_attention_kernel<T, \
|
|
Dh, \
|
|
Dh_MAX, \
|
|
THDS_PER_KEY, \
|
|
THDS_PER_VALUE, \
|
|
THDS_PER_BLOCK, \
|
|
BLOCK_SIZE, \
|
|
CACHE_TYPE, \
|
|
GQA_PARTITION_SIZE, \
|
|
GQA_SUB_PARTITION_SIZE, \
|
|
GQA_NUM_SUB_PARTITIONS, \
|
|
decltype(load_func), \
|
|
decltype(store_func)>; \
|
|
if (smem_sz > 0xc000) { \
|
|
cudaError_t result = cudaFuncSetAttribute( \
|
|
kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz); \
|
|
if (result != cudaSuccess) { \
|
|
result = cudaGetLastError(); \
|
|
PADDLE_THROW(::common::errors::Unavailable( \
|
|
" cudaFuncSetAttribute() returned error %s", \
|
|
cudaGetErrorString(result))); \
|
|
} \
|
|
} \
|
|
dim3 grid(params.kv_num_head *GQA_NUM_SUB_PARTITIONS, params.batch_size); \
|
|
kernel_fn<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>( \
|
|
params, load_func, store_func);
|
|
#endif
|
|
|
|
template <typename T,
|
|
typename SUM_T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int BlockSize,
|
|
int THREADS_PER_VALUE,
|
|
int THREADS_PER_KEY,
|
|
int THREADS_PER_BLOCK,
|
|
CacheType CACHE_TYPE,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
void dispatch_blha_impl_kernel(const Block_AttN_params<T> ¶ms,
|
|
const GPU(Stream_t) & stream,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func) {
|
|
VLOG(1) << "group wise";
|
|
BLHAG_LAUNCH_KERNEL(T,
|
|
SUM_T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THREADS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
}
|
|
|
|
template <typename T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int BlockSize,
|
|
int THREADS_PER_VALUE,
|
|
int THREADS_PER_KEY,
|
|
int THREADS_PER_BLOCK,
|
|
CacheType CACHE_TYPE,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
void dispatch_blha_gqa_kernel(const Block_AttN_params<T> ¶ms,
|
|
const GPU(Stream_t) & stream,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func) {
|
|
if (params.gqa_num_per_partitions == 1 || !FLAGS_use_xqa_optim) {
|
|
auto dispatch_blha_kernel = [&](auto kernel_type, auto qk_sum_type) {
|
|
using Kernel_T = decltype(kernel_type);
|
|
using SUM_T = decltype(qk_sum_type);
|
|
dispatch_blha_impl_kernel<Kernel_T,
|
|
SUM_T,
|
|
Dh,
|
|
Dh_MAX,
|
|
BlockSize,
|
|
THREADS_PER_VALUE,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_BLOCK,
|
|
CACHE_TYPE>(
|
|
params, stream, load_func, store_func);
|
|
};
|
|
if (FLAGS_blha_use_fp32_qk_sum) {
|
|
if constexpr (std::is_same_v<T, float16>) {
|
|
dispatch_blha_kernel(float16{}, float{});
|
|
} else {
|
|
dispatch_blha_kernel(T{}, T{});
|
|
}
|
|
} else {
|
|
dispatch_blha_kernel(T{}, T{});
|
|
}
|
|
|
|
} else if (params.gqa_num_per_partitions == 2) {
|
|
constexpr int THDS_PER_BLOCK = 1024;
|
|
BLHA_LAUNCH_GQA_KERNEL(T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THDS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
2,
|
|
2,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
} else if (params.gqa_num_per_partitions == 4) {
|
|
constexpr int THDS_PER_BLOCK = 1024;
|
|
BLHA_LAUNCH_GQA_KERNEL(T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THDS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
4,
|
|
4,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
} else if (params.gqa_num_per_partitions == 6) {
|
|
constexpr int THDS_PER_BLOCK = 1024;
|
|
BLHA_LAUNCH_GQA_KERNEL(T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THDS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
6,
|
|
2,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
} else if (params.gqa_num_per_partitions == 7) {
|
|
constexpr int THDS_PER_BLOCK = 1024;
|
|
BLHA_LAUNCH_GQA_KERNEL(T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THDS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
7,
|
|
1,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
} else if (params.gqa_num_per_partitions == 8) {
|
|
constexpr int THDS_PER_BLOCK = 1024;
|
|
BLHA_LAUNCH_GQA_KERNEL(T,
|
|
Dh,
|
|
Dh_MAX,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_VALUE,
|
|
THDS_PER_BLOCK,
|
|
BlockSize,
|
|
CACHE_TYPE,
|
|
8,
|
|
4,
|
|
stream,
|
|
load_func,
|
|
store_func)
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"gqa_num_per_partitions = %d is unsupported!",
|
|
params.gqa_num_per_partitions));
|
|
}
|
|
}
|
|
|
|
template <typename T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int BlockSize,
|
|
int THREADS_PER_KEY,
|
|
int THREADS_PER_BLOCK,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
void dispatch_blha_impl(const Block_AttN_params<T> ¶ms,
|
|
const GPU(Stream_t) & stream,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func,
|
|
const int use_cachekv_int8) {
|
|
constexpr int THREADS_PER_VALUE = Dh_MAX * sizeof(T) / 16;
|
|
|
|
if (use_cachekv_int8 > 0) {
|
|
VLOG(1) << "cache_int8";
|
|
dispatch_blha_gqa_kernel<T,
|
|
Dh,
|
|
Dh_MAX,
|
|
BlockSize,
|
|
THREADS_PER_VALUE,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_BLOCK,
|
|
CacheType::INT8>(
|
|
params, stream, load_func, store_func);
|
|
} else {
|
|
VLOG(1) << "normal cache";
|
|
dispatch_blha_gqa_kernel<T,
|
|
Dh,
|
|
Dh_MAX,
|
|
BlockSize,
|
|
THREADS_PER_VALUE,
|
|
THREADS_PER_KEY,
|
|
THREADS_PER_BLOCK,
|
|
CacheType::NORMAL>(
|
|
params, stream, load_func, store_func);
|
|
}
|
|
}
|
|
|
|
template <typename T,
|
|
int Dh,
|
|
int Dh_MAX,
|
|
int BlockSize,
|
|
typename LoadFunc,
|
|
typename StoreFunc>
|
|
void dispatch_blha_impl_key_and_thread(const Block_AttN_params<T> ¶ms,
|
|
const GPU(Stream_t) & stream,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func,
|
|
const int use_cachekv_int8) {
|
|
dispatch_blha_impl<T, Dh, Dh_MAX, BlockSize, 4, 256>(
|
|
params, stream, load_func, store_func, use_cachekv_int8);
|
|
}
|
|
|
|
template <typename T, int Dh, int Dh_MAX, typename LoadFunc, typename StoreFunc>
|
|
void dispatch_blha_impl_blocksize(const Block_AttN_params<T> ¶ms,
|
|
const GPU(Stream_t) & stream,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func,
|
|
const int use_cachekv_int8) {
|
|
switch (params.block_size) {
|
|
case 1:
|
|
dispatch_blha_impl_key_and_thread<T, Dh, Dh_MAX, 1>(
|
|
params, stream, load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 32:
|
|
dispatch_blha_impl_key_and_thread<T, Dh, Dh_MAX, 32>(
|
|
params, stream, load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 64:
|
|
dispatch_blha_impl_key_and_thread<T, Dh, Dh_MAX, 64>(
|
|
params, stream, load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 128:
|
|
dispatch_blha_impl_key_and_thread<T, Dh, Dh_MAX, 128>(
|
|
params, stream, load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
default:
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"block_size = %d is unsupported!", params.block_size));
|
|
}
|
|
}
|
|
|
|
template <typename T, typename LoadFunc, typename StoreFunc>
|
|
void dispatch_blha_impl_headsize(const GPUContext &dev_ctx,
|
|
const Block_AttN_params<T> ¶ms,
|
|
int dim_head,
|
|
LoadFunc load_func,
|
|
StoreFunc store_func,
|
|
const int use_cachekv_int8) {
|
|
switch (dim_head) {
|
|
case 32:
|
|
dispatch_blha_impl_blocksize<T, 32, 32>(
|
|
params, dev_ctx.stream(), load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 64:
|
|
dispatch_blha_impl_blocksize<T, 64, 64>(
|
|
params, dev_ctx.stream(), load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 96:
|
|
dispatch_blha_impl_blocksize<T, 96, 128>(
|
|
params, dev_ctx.stream(), load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
case 128:
|
|
dispatch_blha_impl_blocksize<T, 128, 128>(
|
|
params, dev_ctx.stream(), load_func, store_func, use_cachekv_int8);
|
|
break;
|
|
default:
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Dim_head = %d is unsupported!", dim_head));
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void DispatchBLHA(const GPUContext &dev_ctx,
|
|
const DenseTensor &qkv_tensor,
|
|
const Block_AttN_params<T> ¶ms,
|
|
int q_num_head,
|
|
int kv_num_head,
|
|
int dim_head,
|
|
int use_cachekv_int8,
|
|
DenseTensor *out_tensor) {
|
|
MMHALoad<T> load_func(qkv_tensor.data<T>());
|
|
MMHAStore<T> store_func(out_tensor->data<T>());
|
|
dispatch_blha_impl_headsize(
|
|
dev_ctx, params, dim_head, load_func, store_func, use_cachekv_int8);
|
|
}
|
|
|
|
template <typename T>
|
|
void blha(const GPUContext &dev_ctx,
|
|
const DenseTensor &qkv_tensor,
|
|
const DenseTensor *qkv_bias_tensor,
|
|
const DenseTensor *block_tables,
|
|
const DenseTensor *src_mask_tensor,
|
|
const DenseTensor *cum_offsets_tensor,
|
|
const DenseTensor *sequence_lengths_tensor,
|
|
const DenseTensor *rotary_tensor,
|
|
DenseTensor *k_cache,
|
|
DenseTensor *v_cache,
|
|
DenseTensor *out_tensor,
|
|
const int batch_size,
|
|
const int max_num_blocks_per_seq,
|
|
const int block_size,
|
|
const int seq_len,
|
|
const int pre_cache_length,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int dim_head,
|
|
const int timestep,
|
|
const int rotary_emb_dims,
|
|
float inv_sqrt_dh,
|
|
const float rope_theta,
|
|
const bool add_qkv_bias = true,
|
|
const bool neox_rotary_style = false,
|
|
const int quant_round_type = 1,
|
|
const float quant_max_bound = 127.0f,
|
|
const float quant_min_bound = -127.0f,
|
|
const DenseTensor *cache_k_quant_scales = nullptr,
|
|
const DenseTensor *cache_v_quant_scales = nullptr,
|
|
const DenseTensor *cache_k_dequant_scales = nullptr,
|
|
const DenseTensor *cache_v_dequant_scales = nullptr,
|
|
const DenseTensor *dequant_qkv_scales = nullptr,
|
|
const DenseTensor *shift = nullptr,
|
|
const DenseTensor *smooth = nullptr,
|
|
const float quant_fmha_out_scale = -1,
|
|
int use_cachekv_int8 = 0) {
|
|
Block_AttN_params<T> params;
|
|
|
|
if (cache_k_quant_scales) {
|
|
VLOG(1) << "blha quant cachekv";
|
|
params.k_cache_I = k_cache->data<uint8_t>();
|
|
params.v_cache_I = v_cache->data<uint8_t>();
|
|
params.cache_k_quant_scales = cache_k_quant_scales->data<float>();
|
|
params.cache_v_quant_scales = cache_v_quant_scales->data<float>();
|
|
params.cache_k_dequant_scales = cache_k_dequant_scales->data<float>();
|
|
params.cache_v_dequant_scales = cache_v_dequant_scales->data<float>();
|
|
} else {
|
|
VLOG(1) << "blha not quant cachekv";
|
|
params.k_cache = k_cache->data<T>();
|
|
params.v_cache = v_cache->data<T>();
|
|
}
|
|
|
|
params.max_num_blocks_per_seq = max_num_blocks_per_seq;
|
|
params.neox_rotary_style = neox_rotary_style;
|
|
params.attn_mask = nullptr;
|
|
bool mask_broadcast_num_heads = false;
|
|
if (src_mask_tensor) {
|
|
if (src_mask_tensor->dims()[1] == 1) {
|
|
// all head share a mask.
|
|
mask_broadcast_num_heads = true;
|
|
} else if (src_mask_tensor->dims()[1] == q_num_head) {
|
|
mask_broadcast_num_heads = false;
|
|
} else {
|
|
PADDLE_THROW(errors::InvalidArgument(
|
|
"Unknown dimension for attn_mask, the q_num_head(2nd) "
|
|
"dimension is invalid, it should be 1 or q_num_head(%d), "
|
|
"but got %d",
|
|
q_num_head,
|
|
src_mask_tensor->dims()[1]));
|
|
}
|
|
params.attn_mask = src_mask_tensor->data<T>();
|
|
params.mask_broadcast_num_heads = mask_broadcast_num_heads;
|
|
params.mask_length = src_mask_tensor->dims()[3];
|
|
} else {
|
|
params.attn_mask = nullptr;
|
|
}
|
|
params.block_tables = block_tables->data<int>();
|
|
|
|
if (sequence_lengths_tensor) {
|
|
params.sequence_lengths = sequence_lengths_tensor->data<int>();
|
|
}
|
|
|
|
if (cum_offsets_tensor) {
|
|
params.cum_offsets = cum_offsets_tensor->data<int>();
|
|
} else {
|
|
params.cum_offsets = nullptr;
|
|
}
|
|
params.seq_len = seq_len;
|
|
params.pre_cache_length = pre_cache_length;
|
|
|
|
if (rotary_tensor) {
|
|
params.rotary_emb = rotary_tensor->data<float>();
|
|
params.rope_stride = rotary_tensor->dims()[2] * rotary_tensor->dims()[4];
|
|
} else {
|
|
params.rotary_emb = nullptr;
|
|
}
|
|
|
|
params.add_qkv_bias = add_qkv_bias;
|
|
if (add_qkv_bias) {
|
|
params.qkv_bias = qkv_bias_tensor->data<T>();
|
|
}
|
|
|
|
params.batch_size = batch_size;
|
|
params.block_size = block_size;
|
|
params.q_num_head = q_num_head;
|
|
params.kv_num_head = kv_num_head;
|
|
params.gqa_num_per_partitions = q_num_head / kv_num_head;
|
|
|
|
params.timestep = timestep + pre_cache_length;
|
|
params.inv_sqrt_dh = inv_sqrt_dh;
|
|
params.rotary_emb_dims = rotary_emb_dims;
|
|
params.rope_theta = rope_theta;
|
|
|
|
VLOG(3) << "batch_size: " << batch_size << " q_num_head: " << q_num_head
|
|
<< " kv_num_head: " << kv_num_head << " block_size: " << block_size
|
|
<< " timestep: " << timestep;
|
|
|
|
DispatchBLHA<T>(dev_ctx,
|
|
qkv_tensor,
|
|
params,
|
|
q_num_head,
|
|
kv_num_head,
|
|
dim_head,
|
|
use_cachekv_int8,
|
|
out_tensor);
|
|
}
|
|
|
|
inline GPU(Error_t) GetNumBlocks(int64_t n, int *num_blocks) {
|
|
constexpr int kBlockSize = 128;
|
|
constexpr int kNumWaves = 16;
|
|
|
|
const int device_id = backends::gpu::GetCurrentDeviceId();
|
|
const int sm_count = backends::gpu::GetGPUMultiProcessors(device_id);
|
|
const int max_thread_per_multiprocessor =
|
|
backends::gpu::GetGPUMaxThreadsPerMultiProcessor(device_id);
|
|
|
|
*num_blocks =
|
|
std::max<int>(1,
|
|
std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
|
|
sm_count * max_thread_per_multiprocessor /
|
|
kBlockSize * kNumWaves));
|
|
return GPU(Success);
|
|
}
|
|
|
|
template <class Func>
|
|
inline GPU(Error_t) GetNumBlocks(Func func,
|
|
int64_t block_size,
|
|
size_t dynamic_smem_size,
|
|
int64_t max_blocks,
|
|
int64_t waves,
|
|
int *num_blocks) {
|
|
int dev;
|
|
{
|
|
GPU(Error_t) err = GPU(GetDevice)(&dev);
|
|
if (err != GPU(Success)) {
|
|
return err;
|
|
}
|
|
}
|
|
int sm_count;
|
|
{
|
|
GPU(Error_t)
|
|
err = GPU(DeviceGetAttribute)(&sm_count, GPUMultiProcessorCount, dev);
|
|
if (err != GPU(Success)) {
|
|
return err;
|
|
}
|
|
}
|
|
int max_active_blocks;
|
|
{
|
|
GPU(Error_t)
|
|
err = GPU(OccupancyMaxActiveBlocksPerMultiprocessor)(
|
|
&max_active_blocks, func, block_size, dynamic_smem_size);
|
|
}
|
|
*num_blocks = std::max<int>(
|
|
1, std::min<int64_t>(max_blocks, sm_count * max_active_blocks * waves));
|
|
return GPU(Success);
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void cache_int8_kernel(
|
|
const T *__restrict__ qkv, // [num_tokens, num_heads + 2 * kv_num_heads,
|
|
// head_size]
|
|
uint8_t *__restrict__ key_cache, // [num_blocks, kv_num_heads,
|
|
// block_size, head_size]
|
|
uint8_t *__restrict__ value_cache, // [num_blocks, kv_num_heads,
|
|
// block_size, head_size]
|
|
const int *__restrict__ block_tables, // [bsz, max_blocks_per_seq]
|
|
const int *__restrict__ padding_offsets, // [num_tokens]
|
|
const int *__restrict__ seq_lens, // [bsz]
|
|
const float *cache_k_scales,
|
|
const float *cache_v_scales,
|
|
const int max_seq_len,
|
|
const int max_blocks_per_seq,
|
|
const int q_num_heads,
|
|
const int kv_num_heads,
|
|
const int head_size,
|
|
const int block_size,
|
|
const int pre_cache_length,
|
|
const int elem_cnt,
|
|
const int round_type,
|
|
const float max_bound,
|
|
const float min_bound) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadKVT = AlignedVector<uint8_t, VecSize>;
|
|
LoadT src_vec;
|
|
LoadKVT cache_vec;
|
|
|
|
uint32_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
const uint32_t hidden_size = kv_num_heads * head_size;
|
|
const uint32_t offset = 2 * hidden_size;
|
|
for (uint32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const uint32_t token_idx = linear_index / offset;
|
|
const uint32_t bias = linear_index % offset;
|
|
const uint32_t qkv_id = bias / hidden_size; // skip q
|
|
const uint32_t qkv_bias = bias % hidden_size;
|
|
const uint32_t hi = qkv_bias / head_size;
|
|
const uint32_t h_bias = qkv_bias % head_size;
|
|
const uint32_t ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const uint32_t ori_bi = ori_token_idx / max_seq_len;
|
|
if (seq_lens[ori_bi] == 0) continue;
|
|
const uint32_t ori_seq_id = ori_token_idx % max_seq_len + pre_cache_length;
|
|
|
|
const int32_t *block_table_now = block_tables + ori_bi * max_blocks_per_seq;
|
|
const uint32_t block_idx = block_table_now[ori_seq_id / block_size];
|
|
const uint32_t block_offset = ori_seq_id % block_size;
|
|
|
|
const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * head_size +
|
|
hi * block_size * head_size +
|
|
block_offset * head_size + h_bias;
|
|
const uint32_t ori_idx =
|
|
token_idx * (q_num_heads + 2 * kv_num_heads) * head_size +
|
|
q_num_heads * head_size + qkv_id * hidden_size + hi * head_size +
|
|
h_bias;
|
|
Load<T, VecSize>(&qkv[ori_idx], &src_vec);
|
|
|
|
const uint32_t cache_idx = hi;
|
|
#ifdef PADDLE_WITH_HIP
|
|
T scale;
|
|
if constexpr (kernel_dtype_is_same<T, half>::value) {
|
|
scale = qkv_id == 0 ? __float2half(cache_k_scales[cache_idx])
|
|
: __float2half(cache_v_scales[cache_idx]);
|
|
} else {
|
|
scale = qkv_id == 0 ? static_cast<T>(cache_k_scales[cache_idx])
|
|
: static_cast<T>(cache_v_scales[cache_idx]);
|
|
}
|
|
#else
|
|
const T scale =
|
|
qkv_id == 0 ? cache_k_scales[cache_idx] : cache_v_scales[cache_idx];
|
|
#endif
|
|
#pragma unroll
|
|
for (uint32_t i = 0; i < VecSize; i++) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
float quant_value;
|
|
if constexpr (kernel_dtype_is_same<T, half>::value) {
|
|
quant_value = __half2float(scale * src_vec[i]);
|
|
} else {
|
|
quant_value = static_cast<float>(scale * src_vec[i]);
|
|
}
|
|
#else
|
|
float quant_value = static_cast<float>(scale * src_vec[i]);
|
|
#endif
|
|
if (round_type == 0) {
|
|
quant_value = static_cast<float>(roundWithTiesToEven(quant_value));
|
|
} else {
|
|
quant_value = static_cast<float>(round(quant_value));
|
|
}
|
|
quant_value = quant_value > max_bound ? max_bound : quant_value;
|
|
quant_value = quant_value < min_bound ? min_bound : quant_value;
|
|
cache_vec[i] = static_cast<uint8_t>(quant_value + 128.0f);
|
|
}
|
|
|
|
if (qkv_id == 0) {
|
|
Store<uint8_t, VecSize>(cache_vec, &key_cache[tgt_idx]);
|
|
} else {
|
|
Store<uint8_t, VecSize>(cache_vec, &value_cache[tgt_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void cache_kernel(
|
|
const T *__restrict__ qkv, // [num_tokens, num_heads + 2 * kv_num_heads,
|
|
// head_size]
|
|
T *__restrict__ key_cache, // [num_blocks, kv_num_heads, block_size,
|
|
// head_size]
|
|
T *__restrict__ value_cache, // [num_blocks, kv_num_heads, block_size,
|
|
// head_size]
|
|
const int *__restrict__ block_tables, // [bsz, max_blocks_per_seq]
|
|
const int *__restrict__ padding_offsets, // [num_tokens]
|
|
const int *__restrict__ seq_lens, // [bsz]
|
|
const int max_seq_len,
|
|
const int max_blocks_per_seq,
|
|
const int q_num_heads,
|
|
const int kv_num_heads,
|
|
const int head_size,
|
|
const int block_size,
|
|
const int pre_cache_length,
|
|
const int elem_cnt) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
uint32_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
const uint32_t hidden_size = kv_num_heads * head_size;
|
|
const uint32_t offset = 2 * hidden_size;
|
|
for (uint32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const uint32_t token_idx = linear_index / offset;
|
|
const uint32_t bias = linear_index % offset;
|
|
const uint32_t qkv_id = bias / hidden_size; // skip q
|
|
const uint32_t qkv_bias = bias % hidden_size;
|
|
const uint32_t hi = qkv_bias / head_size;
|
|
const uint32_t h_bias = qkv_bias % head_size;
|
|
const uint32_t ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const uint32_t ori_bi = ori_token_idx / max_seq_len;
|
|
if (seq_lens[ori_bi] == 0) continue;
|
|
const uint32_t ori_seq_id = ori_token_idx % max_seq_len + pre_cache_length;
|
|
|
|
const int32_t *block_table_now = block_tables + ori_bi * max_blocks_per_seq;
|
|
const uint32_t block_idx = block_table_now[ori_seq_id / block_size];
|
|
const uint32_t block_offset = ori_seq_id % block_size;
|
|
|
|
const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * head_size +
|
|
hi * block_size * head_size +
|
|
block_offset * head_size + h_bias;
|
|
const uint32_t ori_idx =
|
|
token_idx * (q_num_heads + 2 * kv_num_heads) * head_size +
|
|
q_num_heads * head_size + qkv_id * hidden_size + hi * head_size +
|
|
h_bias;
|
|
Load<T, VecSize>(&qkv[ori_idx], &src_vec);
|
|
if (qkv_id == 0) {
|
|
Store<T, VecSize>(src_vec, &key_cache[tgt_idx]);
|
|
} else {
|
|
Store<T, VecSize>(src_vec, &value_cache[tgt_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void write_pre_cache_int8_to_cache(
|
|
uint8_t *__restrict__ key_cache, // [num_blocks, num_heads, block_size,
|
|
// head_size]
|
|
uint8_t *__restrict__ value_cache,
|
|
const T *__restrict__ pre_key_cache, // [bsz, pre_cache_len, num_head,
|
|
// head_dim]
|
|
const T *__restrict__ pre_value_cache,
|
|
const int *__restrict__ block_tables, // [bsz, max_blocks_per_seq]
|
|
const int *__restrict__ seq_lens,
|
|
const float *cache_k_scales,
|
|
const float *cache_v_scales,
|
|
const int max_seq_len,
|
|
const int max_blocks_per_seq,
|
|
const int num_heads,
|
|
const int head_size,
|
|
const int block_size,
|
|
const int pre_cache_length,
|
|
const int elem_cnt,
|
|
const int round_type,
|
|
const float max_bound,
|
|
const float min_bound) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadKVT = AlignedVector<uint8_t, VecSize>;
|
|
LoadT src_vec;
|
|
LoadKVT cache_vec;
|
|
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int hidden_size = pre_cache_length * head_size;
|
|
const int cache_hidden_size = num_heads * hidden_size;
|
|
const int offset = 2 * cache_hidden_size;
|
|
|
|
for (int32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int batch_id = linear_index / offset;
|
|
if (seq_lens[batch_id] == 0) continue;
|
|
const int *block_table_now = block_tables + batch_id * max_blocks_per_seq;
|
|
|
|
const int32_t cache_seq_id = (linear_index % hidden_size) / head_size;
|
|
const int32_t head_id = (linear_index % cache_hidden_size) / hidden_size;
|
|
const int32_t size_id = linear_index % head_size;
|
|
|
|
const int32_t kv_id = (linear_index % offset) / cache_hidden_size;
|
|
const int32_t read_id = batch_id * cache_hidden_size +
|
|
head_id * hidden_size + cache_seq_id * head_size +
|
|
size_id;
|
|
if (kv_id == 0) {
|
|
Load<T, VecSize>(&pre_key_cache[read_id], &src_vec);
|
|
} else {
|
|
Load<T, VecSize>(&pre_value_cache[read_id], &src_vec);
|
|
}
|
|
|
|
const int block_idx = block_table_now[cache_seq_id / block_size];
|
|
const int block_offset = cache_seq_id % block_size;
|
|
|
|
const int tgt_idx = block_idx * num_heads * block_size * head_size +
|
|
head_id * block_size * head_size +
|
|
block_offset * head_size + size_id;
|
|
|
|
const float scale =
|
|
kv_id == 0 ? cache_k_scales[head_id] : cache_v_scales[head_id];
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
float quant_value;
|
|
if constexpr (kernel_dtype_is_same<T, half>::value) {
|
|
quant_value = scale * __half2float(src_vec[i]);
|
|
} else {
|
|
quant_value = scale * static_cast<float>(src_vec[i]);
|
|
}
|
|
#else
|
|
float quant_value = scale * static_cast<float>(src_vec[i]);
|
|
#endif
|
|
if (round_type == 0) {
|
|
quant_value = static_cast<float>(roundWithTiesToEven(quant_value));
|
|
} else {
|
|
quant_value = static_cast<float>(round(quant_value));
|
|
}
|
|
quant_value = quant_value > max_bound ? max_bound : quant_value;
|
|
quant_value = quant_value < min_bound ? min_bound : quant_value;
|
|
cache_vec[i] = static_cast<uint8_t>(quant_value + 128.0f);
|
|
}
|
|
|
|
if (kv_id == 0) {
|
|
Store<uint8_t, VecSize>(cache_vec, &key_cache[tgt_idx]);
|
|
} else {
|
|
Store<uint8_t, VecSize>(cache_vec, &value_cache[tgt_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void write_pre_cache_to_cache(
|
|
T *__restrict__ key_cache, // [num_blocks, num_heads, block_size,
|
|
// head_size]
|
|
T *__restrict__ value_cache,
|
|
const T *__restrict__ pre_key_cache, // [bsz, num_heads, pre_cache_len,
|
|
// head_dim]
|
|
const T *__restrict__ pre_value_cache,
|
|
const int *__restrict__ block_tables, // [bsz, max_blocks_per_seq]
|
|
const int *__restrict__ seq_lens,
|
|
const int max_seq_len,
|
|
const int max_blocks_per_seq,
|
|
const int num_heads,
|
|
const int head_size,
|
|
const int block_size,
|
|
const int pre_cache_length,
|
|
const int elem_cnt) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int hidden_size = pre_cache_length * head_size;
|
|
const int cache_hidden_size = num_heads * hidden_size;
|
|
const int offset = 2 * cache_hidden_size;
|
|
|
|
for (int32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int batch_id = linear_index / offset;
|
|
if (seq_lens[batch_id] == 0) continue;
|
|
const int *block_table_now = block_tables + batch_id * max_blocks_per_seq;
|
|
|
|
const int32_t cache_seq_id = (linear_index % hidden_size) / head_size;
|
|
const int32_t head_id = (linear_index % cache_hidden_size) / hidden_size;
|
|
const int32_t size_id = linear_index % head_size;
|
|
|
|
const int32_t kv_id = (linear_index % offset) / cache_hidden_size;
|
|
const int32_t read_id = batch_id * cache_hidden_size +
|
|
head_id * hidden_size + cache_seq_id * head_size +
|
|
size_id;
|
|
if (kv_id == 0) {
|
|
Load<T, VecSize>(&pre_key_cache[read_id], &src_vec);
|
|
} else {
|
|
Load<T, VecSize>(&pre_value_cache[read_id], &src_vec);
|
|
}
|
|
|
|
const int block_idx = block_table_now[cache_seq_id / block_size];
|
|
const int block_offset = cache_seq_id % block_size;
|
|
|
|
const int tgt_idx = block_idx * num_heads * block_size * head_size +
|
|
head_id * block_size * head_size +
|
|
block_offset * head_size + size_id;
|
|
|
|
if (kv_id == 0) {
|
|
Store<T, VecSize>(src_vec, &key_cache[tgt_idx]);
|
|
} else {
|
|
Store<T, VecSize>(src_vec, &value_cache[tgt_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void CacheKernel(const GPUContext &dev_ctx,
|
|
const DenseTensor &qkv, // [token_num, (2 * kv_num_head +
|
|
// q_num_head), head_dim]
|
|
const DenseTensor &block_tables,
|
|
const DenseTensor &padding_offsets,
|
|
const DenseTensor &seq_lens,
|
|
const optional<DenseTensor> &pre_key_cache,
|
|
const optional<DenseTensor> &pre_value_cache,
|
|
const optional<DenseTensor> &cache_k_scales,
|
|
const optional<DenseTensor> &cache_v_scales,
|
|
const int batch_size,
|
|
const int num_tokens,
|
|
const int q_num_heads,
|
|
const int kv_num_heads,
|
|
const int head_size,
|
|
const int max_seq_len,
|
|
const int pre_cache_length,
|
|
DenseTensor *key_cache_out,
|
|
DenseTensor *value_cache_out,
|
|
const int round_type = 0,
|
|
const float max_bound = 127.0,
|
|
const float min_bound = -127.0) {
|
|
typedef PDDataTypeTraits<T> traits_;
|
|
typedef typename traits_::DataType DataType_;
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t max_blocks_per_seq = block_tables.dims()[1];
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t block_size = key_cache_out->dims()[2];
|
|
|
|
// stage 1: write qkv to cache [pre_cache_length:]
|
|
int elem_nums = num_tokens * 2 * kv_num_heads * head_size; // just k and v
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
|
|
if (cache_k_scales) {
|
|
VLOG(1) << "cache kv quant";
|
|
cache_int8_kernel<DataType_, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
reinterpret_cast<DataType_ *>(const_cast<T *>(qkv.data<T>())),
|
|
key_cache_out->data<uint8_t>(),
|
|
value_cache_out->data<uint8_t>(),
|
|
block_tables.data<int>(),
|
|
padding_offsets.data<int>(),
|
|
seq_lens.data<int>(),
|
|
cache_k_scales.get().data<float>(),
|
|
cache_v_scales.get().data<float>(),
|
|
max_seq_len,
|
|
max_blocks_per_seq,
|
|
q_num_heads,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
pre_cache_length,
|
|
elem_nums,
|
|
round_type,
|
|
max_bound,
|
|
min_bound);
|
|
} else {
|
|
VLOG(1) << "cache kv not quant";
|
|
cache_kernel<DataType_, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
reinterpret_cast<DataType_ *>(const_cast<T *>(qkv.data<T>())),
|
|
reinterpret_cast<DataType_ *>(key_cache_out->data<T>()),
|
|
reinterpret_cast<DataType_ *>(value_cache_out->data<T>()),
|
|
block_tables.data<int>(),
|
|
padding_offsets.data<int>(),
|
|
seq_lens.data<int>(),
|
|
max_seq_len,
|
|
max_blocks_per_seq,
|
|
q_num_heads,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
pre_cache_length,
|
|
elem_nums);
|
|
}
|
|
|
|
if (pre_key_cache) {
|
|
// NOTE: not support gqa
|
|
// stage 2: write pre_cache to cache [:pre_cache_length]
|
|
elem_nums = batch_size * kv_num_heads * pre_cache_length * head_size * 2;
|
|
pack_num = elem_nums / PackSize;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (cache_k_scales) {
|
|
write_pre_cache_int8_to_cache<DataType_, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
key_cache_out->data<uint8_t>(),
|
|
value_cache_out->data<uint8_t>(),
|
|
reinterpret_cast<const DataType_ *>(
|
|
pre_key_cache.get().data<T>()),
|
|
reinterpret_cast<const DataType_ *>(
|
|
pre_value_cache.get().data<T>()),
|
|
block_tables.data<int>(),
|
|
seq_lens.data<int>(),
|
|
cache_k_scales->data<float>(),
|
|
cache_v_scales->data<float>(),
|
|
max_seq_len,
|
|
max_blocks_per_seq,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
pre_cache_length,
|
|
elem_nums,
|
|
round_type,
|
|
max_bound,
|
|
min_bound);
|
|
} else {
|
|
write_pre_cache_to_cache<DataType_, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
reinterpret_cast<DataType_ *>(key_cache_out->data<T>()),
|
|
reinterpret_cast<DataType_ *>(value_cache_out->data<T>()),
|
|
reinterpret_cast<const DataType_ *>(
|
|
pre_key_cache.get().data<T>()),
|
|
reinterpret_cast<const DataType_ *>(
|
|
pre_value_cache.get().data<T>()),
|
|
block_tables.data<int>(),
|
|
seq_lens.data<int>(),
|
|
max_seq_len,
|
|
max_blocks_per_seq,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
pre_cache_length,
|
|
elem_nums);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void quant_write_cache_int8_kernel(
|
|
const T *__restrict__ qkv, // [num_tokens, 2 * kv_num_heads + q_num_heads,
|
|
// head_size]
|
|
uint8_t *__restrict__ key_cache, // [num_blocks, kv_num_heads, block_size,
|
|
// head_size]
|
|
uint8_t *__restrict__ value_cache, // [num_blocks, kv_num_heads,
|
|
// block_size, head_size]
|
|
const int *__restrict__ block_tables, // [bsz, max_blocks_per_seq]
|
|
const int *__restrict__ padding_offsets, // [num_tokens]
|
|
const int *__restrict__ seq_lens, // [bsz]
|
|
const int max_seq_len,
|
|
const int pre_cache_length,
|
|
const int max_blocks_per_seq,
|
|
const int num_tokens,
|
|
const int q_num_heads,
|
|
const int kv_num_heads,
|
|
const int head_size,
|
|
const int block_size,
|
|
float *k_quant_scales,
|
|
float *v_quant_scales,
|
|
float *k_dequant_scales,
|
|
float *v_dequant_scales) {
|
|
const int hi = blockIdx.x;
|
|
const int b_id = blockIdx.y;
|
|
if (seq_lens[b_id] <= 0) return;
|
|
const int qkv_id = blockIdx.z;
|
|
const int head_group_size = q_num_heads / kv_num_heads;
|
|
|
|
using InVec = AlignedVector<T, VecSize>;
|
|
using OutVec = AlignedVector<uint8_t, VecSize>;
|
|
|
|
InVec in_vec;
|
|
OutVec out_vec;
|
|
InVec abs_max_vec;
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; ++i) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
if constexpr (kernel_dtype_is_same<T, half>::value) {
|
|
abs_max_vec[i] = __float2half(0.0f);
|
|
} else {
|
|
abs_max_vec[i] = static_cast<T>(0.0f);
|
|
}
|
|
#else
|
|
abs_max_vec[i] = 0.0f;
|
|
#endif
|
|
}
|
|
|
|
uint8_t *dst_ptr;
|
|
float *quant_scales;
|
|
float *dequant_scales;
|
|
if (qkv_id == 0) {
|
|
dst_ptr = key_cache;
|
|
quant_scales = k_quant_scales;
|
|
dequant_scales = k_dequant_scales;
|
|
} else {
|
|
dst_ptr = value_cache;
|
|
quant_scales = v_quant_scales;
|
|
dequant_scales = v_dequant_scales;
|
|
}
|
|
|
|
T local_abs_max;
|
|
|
|
for (int idx = threadIdx.x * VecSize; idx < num_tokens * head_size;
|
|
idx += blockDim.x * VecSize) {
|
|
int token_idx = idx / head_size;
|
|
int h_offset = idx % head_size;
|
|
int linear_idx = token_idx * (2 * kv_num_heads + q_num_heads) * head_size +
|
|
(qkv_id + head_group_size) * kv_num_heads * head_size +
|
|
hi * head_size + h_offset;
|
|
|
|
Load<T, VecSize>(qkv + linear_idx, &in_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; ++i) {
|
|
abs_max_vec[i] = MaxFunc<T>()(abs_max_vec[i], AbsFunc<T>()(in_vec[i]));
|
|
}
|
|
}
|
|
|
|
local_abs_max = LocalReduceMax<T, InVec, VecSize>(abs_max_vec);
|
|
T abs_max_val = BlockReduceAbsMax<T>(local_abs_max, 0xffffffff);
|
|
|
|
__shared__ float quant_scale;
|
|
if (threadIdx.x == 0) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
if constexpr (kernel_dtype_is_same<T, half>::value) {
|
|
quant_scale = 127.0f / __half2float(abs_max_val);
|
|
} else {
|
|
quant_scale = 127.0f / static_cast<float>(abs_max_val);
|
|
}
|
|
#else
|
|
quant_scale = 127.0f / static_cast<float>(abs_max_val);
|
|
#endif
|
|
}
|
|
|
|
__syncthreads();
|
|
for (int idx = threadIdx.x * VecSize; idx < num_tokens * head_size;
|
|
idx += blockDim.x * VecSize) {
|
|
int token_idx = idx / head_size;
|
|
int h_offset = idx % head_size;
|
|
int linear_idx = token_idx * (2 * kv_num_heads + q_num_heads) * head_size +
|
|
(qkv_id + head_group_size) * kv_num_heads * head_size +
|
|
hi * head_size + h_offset;
|
|
|
|
Load<T, VecSize>(qkv + linear_idx, &in_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; ++i) {
|
|
out_vec[i] = QuantFunc<T>()(in_vec[i], quant_scale);
|
|
}
|
|
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / max_seq_len;
|
|
if (ori_bi != b_id) continue;
|
|
const int ori_seq_id = ori_token_idx % max_seq_len + pre_cache_length;
|
|
|
|
const int *block_table_now = block_tables + ori_bi * max_blocks_per_seq;
|
|
const int block_idx = block_table_now[ori_seq_id / block_size];
|
|
const int block_offset = ori_seq_id % block_size;
|
|
// [max_block_num, num_head, block_size, head_dim/x, x]
|
|
Store<uint8_t>(out_vec,
|
|
dst_ptr + block_idx * kv_num_heads * block_size * head_size +
|
|
hi * block_size * head_size + block_offset * head_size +
|
|
h_offset);
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
quant_scales[b_id * kv_num_heads + hi] = quant_scale;
|
|
dequant_scales[b_id * kv_num_heads + hi] = 1.0f / quant_scale;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void DynamicQuantCacheKernel(
|
|
const GPUContext &dev_ctx,
|
|
const DenseTensor
|
|
&qkv, // [token_num, 2 * q_num_head + kv_num_head, head_dim]
|
|
const DenseTensor &block_tables,
|
|
const DenseTensor &padding_offsets,
|
|
const DenseTensor &seq_lens,
|
|
const DenseTensor &k_quant_scales,
|
|
const DenseTensor &v_quant_scales,
|
|
const DenseTensor &k_dequant_scales,
|
|
const DenseTensor &v_dequant_scales,
|
|
const optional<DenseTensor> &pre_key_cache,
|
|
const optional<DenseTensor> &pre_value_cache,
|
|
const int batch_size,
|
|
const int q_num_heads,
|
|
const int kv_num_heads,
|
|
const int head_size,
|
|
const int max_seq_len,
|
|
const int pre_cache_length,
|
|
DenseTensor *key_cache_out,
|
|
DenseTensor *value_cache_out) {
|
|
typedef PDDataTypeTraits<T> traits_;
|
|
typedef typename traits_::DataType DataType_;
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t num_tokens = padding_offsets.dims()[0];
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t max_blocks_per_seq = block_tables.dims()[1];
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t block_size = key_cache_out->dims()[2];
|
|
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
|
|
assert(head_size % PackSize == 0);
|
|
|
|
const DataType_ *qkv_ptr = reinterpret_cast<const DataType_ *>(qkv.data<T>());
|
|
|
|
// [max_block_num, num_head, block_size, head_dim]
|
|
|
|
uint8_t *cache_k_ptr = key_cache_out->data<uint8_t>();
|
|
uint8_t *cache_v_ptr = value_cache_out->data<uint8_t>();
|
|
|
|
float *k_quant_scales_data =
|
|
const_cast<float *>(k_quant_scales.data<float>());
|
|
float *k_dequant_scales_data =
|
|
const_cast<float *>(k_dequant_scales.data<float>());
|
|
|
|
float *v_quant_scales_data =
|
|
const_cast<float *>(v_quant_scales.data<float>());
|
|
float *v_dequant_scales_data =
|
|
const_cast<float *>(v_dequant_scales.data<float>());
|
|
|
|
constexpr int block_sz = 1024;
|
|
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
int64_t bsz = seq_lens.dims()[0];
|
|
|
|
dim3 grid(kv_num_heads, bsz, 2);
|
|
|
|
// [token_num, (2 * kv_num_head + q_num_head), head_dim/x, x]->[max_block_num,
|
|
// kv_num_head, block_size, head_dim/x, x] Quant and Write kv
|
|
quant_write_cache_int8_kernel<DataType_, PackSize>
|
|
<<<grid, block_sz, 0, dev_ctx.stream()>>>(qkv_ptr,
|
|
cache_k_ptr,
|
|
cache_v_ptr,
|
|
block_tables.data<int>(),
|
|
padding_offsets.data<int>(),
|
|
seq_lens.data<int>(),
|
|
max_seq_len,
|
|
pre_cache_length,
|
|
max_blocks_per_seq,
|
|
num_tokens,
|
|
q_num_heads,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
k_quant_scales_data,
|
|
v_quant_scales_data,
|
|
k_dequant_scales_data,
|
|
v_dequant_scales_data);
|
|
|
|
if (pre_key_cache) {
|
|
// stage 2: write pre_cache to cache [:pre_cache_length]
|
|
const int elem_nums =
|
|
batch_size * kv_num_heads * pre_cache_length * head_size * 2;
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
write_pre_cache_int8_to_cache<DataType_, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
key_cache_out->data<uint8_t>(),
|
|
value_cache_out->data<uint8_t>(),
|
|
reinterpret_cast<const DataType_ *>(pre_key_cache.get().data<T>()),
|
|
reinterpret_cast<const DataType_ *>(
|
|
pre_value_cache.get().data<T>()),
|
|
block_tables.data<int>(),
|
|
seq_lens.data<int>(),
|
|
k_quant_scales.data<float>(),
|
|
v_quant_scales.data<float>(),
|
|
max_seq_len,
|
|
max_blocks_per_seq,
|
|
kv_num_heads,
|
|
head_size,
|
|
block_size,
|
|
pre_cache_length,
|
|
elem_nums,
|
|
1,
|
|
127.0f,
|
|
-127.0f);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void VariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * last_dim;
|
|
const int offset = 2 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / last_dim;
|
|
const int h_bias = qkv_bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
const int64_t base_idx = token_idx * 3 * hidden_size +
|
|
qkv_id * hidden_size + hi * last_dim + h_bias;
|
|
Load<T, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
const float input_left = static_cast<float>(src_vec[2 * i]);
|
|
const float input_right = static_cast<float>(src_vec[2 * i + 1]);
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
src_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
src_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
}
|
|
Store<T, VecSize>(src_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void NeoxVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
// [token_num, 2, num_head, dim_head / 2]
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * half_lastdim;
|
|
const int full_hidden_size = num_head * last_dim;
|
|
const int offset = 2 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / half_lastdim;
|
|
const int h_bias = qkv_bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int base_idx_left = token_idx * 3 * full_hidden_size +
|
|
qkv_id * full_hidden_size + hi * last_dim +
|
|
h_bias;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
|
|
Load<T, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<T, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
const float input_left = static_cast<float>(left_vec[i]);
|
|
const float input_right = static_cast<float>(right_vec[i]);
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
}
|
|
Store<T, VecSize>(left_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, num_head, dim_head]
|
|
const T *qkv_input, // qkv
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums = token_num * 2 * head_num * dim_head; // just q and k
|
|
if (use_neox_style) {
|
|
elem_nums = token_num * head_num * dim_head;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
VariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
NeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQAVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + kv_num_head) * last_dim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / last_dim;
|
|
const int h_bias = bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int64_t emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
// [token_num, q_num_head + 2 * kv_num_head, last_dim]
|
|
const int64_t base_idx =
|
|
token_idx * (q_num_head + 2 * kv_num_head) * last_dim + hi * last_dim +
|
|
h_bias;
|
|
Load<T, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
const float input_left = static_cast<float>(src_vec[2 * i]);
|
|
const float input_right = static_cast<float>(src_vec[2 * i + 1]);
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
src_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
src_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
}
|
|
Store<T, VecSize>(src_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQANeoxVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
// [token_num, 2, q_num_head, dim_head / 2]
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + kv_num_head) * half_lastdim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / half_lastdim;
|
|
const int h_bias = bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int base_idx_left =
|
|
token_idx * (q_num_head + 2 * kv_num_head) * last_dim + hi * last_dim +
|
|
h_bias;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
|
|
Load<T, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<T, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
const float input_left = static_cast<float>(left_vec[i]);
|
|
const float input_right = static_cast<float>(right_vec[i]);
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
}
|
|
Store<T, VecSize>(left_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void gqa_rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, q_num_head, dim_head]
|
|
const T *qkv_input, // qkv
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums =
|
|
token_num * (q_head_num + kv_head_num) * dim_head; // just q and k
|
|
if (use_neox_style) {
|
|
elem_nums /= 2;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
GQAVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
GQANeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void VariableLengthRotaryKernel(
|
|
const int *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const float *qkv_out_scales, // [3, num_head, dim_head]
|
|
const T *qkv_biases, // [3, num_head, dim_head]
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<int, VecSize>;
|
|
using LoadBiasT = AlignedVector<T, VecSize>;
|
|
using LoadScaleT = AlignedVector<float, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadBiasT bias_vec;
|
|
LoadScaleT out_scale_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * last_dim;
|
|
const int offset = 3 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / last_dim;
|
|
const int h_bias = qkv_bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
const int bias_idx = qkv_id * hidden_size + hi * last_dim + h_bias;
|
|
const int64_t base_idx = token_idx * 3 * hidden_size + bias_idx;
|
|
Load<int, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx], &out_scale_vec);
|
|
if (qkv_id < 2) {
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
float input_left = static_cast<float>(src_vec[2 * i]);
|
|
float input_right = static_cast<float>(src_vec[2 * i + 1]);
|
|
// dequant + bias_add
|
|
input_left = input_left * out_scale_vec[2 * i] +
|
|
static_cast<float>(bias_vec[2 * i]);
|
|
input_right = input_right * out_scale_vec[2 * i + 1] +
|
|
static_cast<float>(bias_vec[2 * i + 1]);
|
|
if (qkv_id < 2) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
bias_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
bias_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
bias_vec[2 * i] = static_cast<T>(input_left);
|
|
bias_vec[2 * i + 1] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(bias_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void NeoxVariableLengthRotaryKernel(
|
|
const int *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const float *qkv_out_scales, // [3, num_head, dim_head]
|
|
const T *qkv_biases, // [3, num_head, dim_head]
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<int, VecSize>;
|
|
using LoadBiasT = AlignedVector<T, VecSize>;
|
|
using LoadScaleT = AlignedVector<float, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadBiasT left_bias_vec;
|
|
LoadBiasT right_bias_vec;
|
|
LoadScaleT left_out_scale_vec;
|
|
LoadScaleT right_out_scale_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * half_lastdim;
|
|
const int full_hidden_size = num_head * last_dim;
|
|
const int offset = 3 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / half_lastdim;
|
|
const int h_bias = qkv_bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int bias_idx_left =
|
|
qkv_id * full_hidden_size + hi * last_dim + h_bias;
|
|
const int bias_idx_right = bias_idx_left + half_lastdim;
|
|
const int base_idx_left = token_idx * 3 * full_hidden_size + bias_idx_left;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
Load<int, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<int, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx_left], &left_out_scale_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx_right], &right_out_scale_vec);
|
|
if (qkv_id < 2) {
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
float input_left = static_cast<float>(left_vec[i]);
|
|
float input_right = static_cast<float>(right_vec[i]);
|
|
// dequant + bias_add
|
|
input_left = input_left * left_out_scale_vec[i] +
|
|
static_cast<float>(left_bias_vec[i]);
|
|
input_right = input_right * right_out_scale_vec[i] +
|
|
static_cast<float>(right_bias_vec[i]);
|
|
if (qkv_id < 2) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_bias_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_bias_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
left_bias_vec[i] = static_cast<T>(input_left);
|
|
right_bias_vec[i] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(left_bias_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_bias_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, num_head, dim_head]
|
|
const int *qkv_input, // qkv
|
|
const float *qkv_out_scales, // [3, num_head, dim_head]
|
|
const T *qkv_bias,
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums = token_num * 3 * head_num * dim_head; // just q and k
|
|
if (use_neox_style) {
|
|
elem_nums = token_num * 3 * head_num * dim_head / 2;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
VariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_out_scales,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
NeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_out_scales,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQAVariableLengthRotaryKernel(
|
|
const int *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const float *qkv_out_scales, // [3, q_num_head, dim_head]
|
|
const T *qkv_biases, // [3, q_num_head, dim_head]
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<int, VecSize>;
|
|
using LoadBiasT = AlignedVector<T, VecSize>;
|
|
using LoadScaleT = AlignedVector<float, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadBiasT bias_vec;
|
|
LoadScaleT out_scale_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + 2 * kv_num_head) * last_dim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / last_dim;
|
|
const int h_bias = bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int64_t emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
const int64_t bias_idx = hi * last_dim + h_bias;
|
|
const int64_t base_idx = token_idx * offset + bias_idx;
|
|
Load<int, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx], &out_scale_vec);
|
|
if (hi < q_num_head + kv_num_head) {
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
float input_left = static_cast<float>(src_vec[2 * i]);
|
|
float input_right = static_cast<float>(src_vec[2 * i + 1]);
|
|
// dequant + bias_add
|
|
input_left = input_left * out_scale_vec[2 * i] +
|
|
static_cast<float>(bias_vec[2 * i]);
|
|
input_right = input_right * out_scale_vec[2 * i + 1] +
|
|
static_cast<float>(bias_vec[2 * i + 1]);
|
|
if (hi < q_num_head + kv_num_head) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
bias_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
bias_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
bias_vec[2 * i] = static_cast<T>(input_left);
|
|
bias_vec[2 * i + 1] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(bias_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQANeoxVariableLengthRotaryKernel(
|
|
const int *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const float *qkv_out_scales, // [3, q_num_head, dim_head]
|
|
const T *qkv_biases, // [3, q_num_head, dim_head]
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<int, VecSize>;
|
|
using LoadBiasT = AlignedVector<T, VecSize>;
|
|
using LoadScaleT = AlignedVector<float, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadBiasT left_bias_vec;
|
|
LoadBiasT right_bias_vec;
|
|
LoadScaleT left_out_scale_vec;
|
|
LoadScaleT right_out_scale_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + 2 * kv_num_head) * half_lastdim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / half_lastdim;
|
|
const int h_bias = bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int bias_idx_left = hi * last_dim + h_bias;
|
|
const int bias_idx_right = bias_idx_left + half_lastdim;
|
|
const int base_idx_left = token_idx * offset + bias_idx_left;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
Load<int, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<int, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx_left], &left_out_scale_vec);
|
|
Load<float, VecSize>(&qkv_out_scales[bias_idx_right], &right_out_scale_vec);
|
|
if (hi < (q_num_head + kv_num_head)) {
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
float input_left = static_cast<float>(left_vec[i]);
|
|
float input_right = static_cast<float>(right_vec[i]);
|
|
// dequant + bias_add
|
|
input_left = input_left * left_out_scale_vec[i] +
|
|
static_cast<float>(left_bias_vec[i]);
|
|
input_right = input_right * right_out_scale_vec[i] +
|
|
static_cast<float>(right_bias_vec[i]);
|
|
if (hi < (q_num_head + kv_num_head)) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_bias_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_bias_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
left_bias_vec[i] = static_cast<T>(input_left);
|
|
right_bias_vec[i] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(left_bias_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_bias_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void gqa_rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, q_num_head, dim_head]
|
|
const int *qkv_input, // qkv
|
|
const float *qkv_out_scales, // [3, q_num_head, dim_head]
|
|
const T *qkv_bias,
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums =
|
|
token_num * (q_head_num + 2 * kv_head_num) * dim_head; // for all q k v
|
|
if (use_neox_style) {
|
|
elem_nums /= 2;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
GQAVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_out_scales,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
GQANeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_out_scales,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void VariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const T *qkv_biases,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadT bias_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * last_dim;
|
|
const int offset = 3 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / last_dim;
|
|
const int h_bias = qkv_bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
const int bias_idx = qkv_id * hidden_size + hi * last_dim + h_bias;
|
|
const int64_t base_idx = token_idx * 3 * hidden_size + bias_idx;
|
|
Load<T, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
const float input_left =
|
|
static_cast<float>(src_vec[2 * i] + bias_vec[2 * i]);
|
|
const float input_right =
|
|
static_cast<float>(src_vec[2 * i + 1] + bias_vec[2 * i + 1]);
|
|
|
|
if (qkv_id < 2) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
src_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
src_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
src_vec[2 * i] = static_cast<T>(input_left);
|
|
src_vec[2 * i + 1] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(src_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void NeoxVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const T *qkv_biases,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadT left_bias_vec;
|
|
LoadT right_bias_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int hidden_size = num_head * half_lastdim;
|
|
const int full_hidden_size = num_head * last_dim;
|
|
const int offset = 3 * hidden_size;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int qkv_id = bias / hidden_size;
|
|
const int qkv_bias = bias % hidden_size;
|
|
const int hi = qkv_bias / half_lastdim;
|
|
const int h_bias = qkv_bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int bias_idx_left =
|
|
qkv_id * full_hidden_size + hi * last_dim + h_bias;
|
|
const int bias_idx_right = bias_idx_left + half_lastdim;
|
|
const int base_idx_left = token_idx * 3 * full_hidden_size + bias_idx_left;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
Load<int, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<int, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
const float input_left =
|
|
static_cast<float>(left_vec[i] + left_bias_vec[i]);
|
|
const float input_right =
|
|
static_cast<float>(right_vec[i] + right_bias_vec[i]);
|
|
|
|
if (qkv_id < 2) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
left_vec[i] = static_cast<T>(input_left);
|
|
right_vec[i] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(left_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, num_head, dim_head]
|
|
const T *qkv_input, // qkv
|
|
const T *qkv_bias,
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums = token_num * 3 * head_num * dim_head; // just q and k
|
|
if (use_neox_style) {
|
|
elem_nums = token_num * 3 * head_num * dim_head / 2;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
VariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
NeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQAVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const T *qkv_biases,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
constexpr int HalfVecSize = VecSize / 2;
|
|
using LoadEmbT = AlignedVector<float, HalfVecSize>;
|
|
LoadT src_vec;
|
|
LoadT bias_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + 2 * kv_num_head) * last_dim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / last_dim;
|
|
const int h_bias = bias % last_dim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int64_t emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
|
|
const int64_t bias_idx = hi * last_dim + h_bias;
|
|
const int64_t base_idx = token_idx * offset + bias_idx;
|
|
Load<T, VecSize>(&qkv[base_idx], &src_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
|
|
Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < HalfVecSize; i++) {
|
|
const float input_left =
|
|
static_cast<float>(src_vec[2 * i] + bias_vec[2 * i]);
|
|
const float input_right =
|
|
static_cast<float>(src_vec[2 * i + 1] + bias_vec[2 * i + 1]);
|
|
|
|
if (hi < q_num_head + kv_num_head) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
src_vec[2 * i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
src_vec[2 * i + 1] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
src_vec[2 * i] = static_cast<T>(input_left);
|
|
src_vec[2 * i + 1] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(src_vec, &qkv_out[base_idx]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize = 1>
|
|
__global__ void GQANeoxVariableLengthRotaryKernel(
|
|
const T *qkv,
|
|
const float *cos_emb, // [1, 1, seq_len, dim_head / 2]
|
|
const float *sin_emb,
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const T *qkv_biases,
|
|
T *qkv_out,
|
|
const int64_t elem_cnt,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int last_dim) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
using LoadEmbT = AlignedVector<float, VecSize>;
|
|
LoadT left_vec;
|
|
LoadT right_vec;
|
|
LoadT left_bias_vec;
|
|
LoadT right_bias_vec;
|
|
LoadEmbT cos_emb_vec;
|
|
LoadEmbT sin_emb_vec;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int half_lastdim = last_dim / 2;
|
|
const int offset = (q_num_head + 2 * kv_num_head) * half_lastdim;
|
|
for (int64_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / offset;
|
|
const int ori_token_idx = token_idx + padding_offsets[token_idx];
|
|
const int ori_bi = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_bi] == 0) continue;
|
|
const int bias = linear_index % offset;
|
|
const int hi = bias / half_lastdim;
|
|
const int h_bias = bias % half_lastdim;
|
|
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
|
|
const int emb_idx = ori_seq_id * last_dim + h_bias;
|
|
const int bias_idx_left = hi * last_dim + h_bias;
|
|
const int bias_idx_right = bias_idx_left + half_lastdim;
|
|
const int base_idx_left = token_idx * offset + bias_idx_left;
|
|
const int base_idx_right = base_idx_left + half_lastdim;
|
|
Load<int, VecSize>(&qkv[base_idx_left], &left_vec);
|
|
Load<int, VecSize>(&qkv[base_idx_right], &right_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
|
|
Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
|
|
Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
|
|
Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
const float input_left =
|
|
static_cast<float>(left_vec[i] + left_bias_vec[i]);
|
|
const float input_right =
|
|
static_cast<float>(right_vec[i] + right_bias_vec[i]);
|
|
|
|
if (hi < (q_num_head + kv_num_head)) { // qk rope
|
|
const float cos_tmp = cos_emb_vec[i];
|
|
const float sin_tmp = sin_emb_vec[i];
|
|
left_vec[i] =
|
|
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
|
|
right_vec[i] =
|
|
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
|
|
} else {
|
|
left_vec[i] = static_cast<T>(input_left);
|
|
right_vec[i] = static_cast<T>(input_right);
|
|
}
|
|
}
|
|
Store<T, VecSize>(left_vec, &qkv_out[base_idx_left]);
|
|
Store<T, VecSize>(right_vec, &qkv_out[base_idx_right]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void gqa_rotary_qk_variable(
|
|
const GPUContext &dev_ctx,
|
|
T *qkv, // [token_num, 3, q_num_head, dim_head]
|
|
const T *qkv_input, // qkv
|
|
const T *qkv_bias,
|
|
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
|
|
const int *padding_offsets,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int seq_len,
|
|
const int input_output_len,
|
|
const int dim_head,
|
|
bool use_neox_style = false) {
|
|
int elem_nums =
|
|
token_num * (q_head_num + 2 * kv_head_num) * dim_head; // for all q k v
|
|
if (use_neox_style) {
|
|
elem_nums /= 2;
|
|
}
|
|
constexpr int PackSize = 16 / sizeof(T);
|
|
const int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
if (!use_neox_style) {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
|
|
GQAVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
} else {
|
|
const float *cos_emb = rotary_emb;
|
|
const float *sin_emb = rotary_emb + input_output_len * dim_head;
|
|
GQANeoxVariableLengthRotaryKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(qkv_input,
|
|
cos_emb,
|
|
sin_emb,
|
|
padding_offsets,
|
|
seq_lens,
|
|
qkv_bias,
|
|
qkv,
|
|
elem_nums,
|
|
q_head_num,
|
|
kv_head_num,
|
|
seq_len,
|
|
dim_head);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize, int RoundType>
|
|
__global__ void ShiftSmoothQuant(const T *input,
|
|
const T *shift,
|
|
const T *smooth,
|
|
float scale,
|
|
int8_t *out,
|
|
int num,
|
|
int cols,
|
|
float quant_max_bound,
|
|
float quant_min_bound) {
|
|
AlignedVector<T, VecSize> in_vec;
|
|
AlignedVector<T, VecSize> shift_vec;
|
|
AlignedVector<T, VecSize> smooth_vec;
|
|
AlignedVector<int8_t, VecSize> out_vec;
|
|
|
|
for (int64_t linear_id =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
linear_id * VecSize < num;
|
|
linear_id += gridDim.x * blockDim.x) {
|
|
int idx = linear_id * VecSize;
|
|
Load<T, VecSize>(input + idx, &in_vec);
|
|
Load<T, VecSize>(shift + (idx % cols), &shift_vec);
|
|
Load<T, VecSize>(smooth + (idx % cols), &smooth_vec);
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
float quant_value =
|
|
quant_max_bound *
|
|
static_cast<float>((in_vec[i] + shift_vec[i]) * smooth_vec[i]) *
|
|
scale;
|
|
quant_value = static_cast<float>(RoundType == 1 ? round(quant_value)
|
|
: rintf(quant_value));
|
|
quant_value =
|
|
quant_value > quant_max_bound ? quant_max_bound : quant_value;
|
|
quant_value =
|
|
quant_value < quant_min_bound ? quant_min_bound : quant_value;
|
|
out_vec[i] = static_cast<int8_t>(quant_value);
|
|
}
|
|
Store<int8_t, VecSize>(out_vec, out + idx);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize, int RoundType>
|
|
__global__ void ShiftSmooth(const T *input,
|
|
const T *shift,
|
|
const T *smooth,
|
|
T *out,
|
|
int num,
|
|
int cols) {
|
|
AlignedVector<T, VecSize> in_vec;
|
|
AlignedVector<T, VecSize> shift_vec;
|
|
AlignedVector<T, VecSize> smooth_vec;
|
|
AlignedVector<T, VecSize> out_vec;
|
|
|
|
for (int64_t linear_id =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
linear_id * VecSize < num;
|
|
linear_id += gridDim.x * blockDim.x) {
|
|
int idx = linear_id * VecSize;
|
|
Load<T, VecSize>(input + idx, &in_vec);
|
|
Load<T, VecSize>(shift + (idx % cols), &shift_vec);
|
|
Load<T, VecSize>(smooth + (idx % cols), &smooth_vec);
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < VecSize; i++) {
|
|
out_vec[i] = (in_vec[i] + shift_vec[i]) * smooth_vec[i];
|
|
}
|
|
Store<T, VecSize>(out_vec, out + idx);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void shift_smooth_quant(const GPUContext &dev_ctx,
|
|
DenseTensor *fmha_out,
|
|
const DenseTensor &fmha_in,
|
|
const DenseTensor &out_linear_shift,
|
|
const DenseTensor &out_linear_smooth,
|
|
float out_linear_in_scale,
|
|
const int num_head,
|
|
const int dim_head,
|
|
const int quant_round_type,
|
|
const float quant_max_bound,
|
|
const float quant_min_bound) {
|
|
constexpr int block_size = 512;
|
|
constexpr int waves = 32;
|
|
constexpr int vec_size = 16 / sizeof(T);
|
|
|
|
int max_blocks = fmha_out->numel() / vec_size;
|
|
int num_blocks = 0;
|
|
if (out_linear_in_scale > 0) {
|
|
if (quant_round_type == 0) {
|
|
GetNumBlocks(ShiftSmoothQuant<T, vec_size, 0>,
|
|
block_size,
|
|
0,
|
|
max_blocks,
|
|
waves,
|
|
&num_blocks);
|
|
ShiftSmoothQuant<T, vec_size, 0>
|
|
<<<num_blocks, block_size, 0, dev_ctx.stream()>>>(
|
|
fmha_in.data<T>(),
|
|
out_linear_shift.data<T>(),
|
|
out_linear_smooth.data<T>(),
|
|
out_linear_in_scale,
|
|
fmha_out->data<int8_t>(),
|
|
fmha_out->numel(),
|
|
num_head * dim_head,
|
|
quant_max_bound,
|
|
quant_min_bound);
|
|
} else {
|
|
GetNumBlocks(ShiftSmoothQuant<T, vec_size, 1>,
|
|
block_size,
|
|
0,
|
|
max_blocks,
|
|
waves,
|
|
&num_blocks);
|
|
ShiftSmoothQuant<T, vec_size, 1>
|
|
<<<num_blocks, block_size, 0, dev_ctx.stream()>>>(
|
|
fmha_in.data<T>(),
|
|
out_linear_shift.data<T>(),
|
|
out_linear_smooth.data<T>(),
|
|
out_linear_in_scale,
|
|
fmha_out->data<int8_t>(),
|
|
fmha_out->numel(),
|
|
num_head * dim_head,
|
|
quant_max_bound,
|
|
quant_min_bound);
|
|
}
|
|
} else {
|
|
if (quant_round_type == 0) {
|
|
GetNumBlocks(ShiftSmooth<T, vec_size, 0>,
|
|
block_size,
|
|
0,
|
|
max_blocks,
|
|
waves,
|
|
&num_blocks);
|
|
ShiftSmooth<T, vec_size, 0>
|
|
<<<num_blocks, block_size, 0, dev_ctx.stream()>>>(
|
|
fmha_in.data<T>(),
|
|
out_linear_shift.data<T>(),
|
|
out_linear_smooth.data<T>(),
|
|
fmha_out->data<T>(),
|
|
fmha_out->numel(),
|
|
num_head * dim_head);
|
|
} else {
|
|
GetNumBlocks(ShiftSmooth<T, vec_size, 1>,
|
|
block_size,
|
|
0,
|
|
max_blocks,
|
|
waves,
|
|
&num_blocks);
|
|
ShiftSmooth<T, vec_size, 1>
|
|
<<<num_blocks, block_size, 0, dev_ctx.stream()>>>(
|
|
fmha_in.data<T>(),
|
|
out_linear_shift.data<T>(),
|
|
out_linear_smooth.data<T>(),
|
|
fmha_out->data<T>(),
|
|
fmha_out->numel(),
|
|
num_head * dim_head);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void fusedQKV_transpose_split_kernel(T *q_buf,
|
|
T *k_buf,
|
|
T *v_buf,
|
|
const T *qkv,
|
|
const int *padding_offset,
|
|
const int *seq_lens,
|
|
const int32_t elem_cnt,
|
|
const int batch_size,
|
|
const int seq_len,
|
|
const int token_num,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int size_per_head) {
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
const int fused_hidden_size = (q_head_num + 2 * kv_head_num) * size_per_head;
|
|
|
|
for (int32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
Load<T, VecSize>(&qkv[linear_index], &src_vec);
|
|
int32_t bias_idx = linear_index % fused_hidden_size;
|
|
const int32_t token_idx = linear_index / fused_hidden_size;
|
|
const int32_t ori_token_idx =
|
|
token_idx + (padding_offset == nullptr ? 0 : padding_offset[token_idx]);
|
|
const int32_t target_batch_id = ori_token_idx / seq_len;
|
|
if (seq_lens[target_batch_id] == 0) continue;
|
|
|
|
const int32_t head_id = bias_idx / size_per_head;
|
|
const int32_t size_id = linear_index % size_per_head;
|
|
|
|
// [token_num, q_head_num or kv_head_num, size_per_head]
|
|
if (head_id < q_head_num) {
|
|
const int32_t write_idx = token_idx * q_head_num * size_per_head +
|
|
head_id * size_per_head + size_id;
|
|
Store<T, VecSize>(src_vec, &q_buf[write_idx]);
|
|
} else {
|
|
if (head_id < q_head_num + kv_head_num) {
|
|
const int32_t write_idx = token_idx * kv_head_num * size_per_head +
|
|
(head_id - q_head_num) * size_per_head +
|
|
size_id;
|
|
Store<T, VecSize>(src_vec, &k_buf[write_idx]);
|
|
} else {
|
|
const int32_t write_idx =
|
|
token_idx * kv_head_num * size_per_head +
|
|
(head_id - q_head_num - kv_head_num) * size_per_head + size_id;
|
|
Store<T, VecSize>(src_vec, &v_buf[write_idx]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void qkv_transpose_split(const GPUContext &dev_ctx,
|
|
T *q_buf,
|
|
T *k_buf,
|
|
T *v_buf,
|
|
const T *qkv,
|
|
const int *padding_offset,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int batch_size,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int seq_len,
|
|
const int size_per_head) {
|
|
const int32_t elem_cnt =
|
|
token_num * (q_head_num + kv_head_num * 2) * size_per_head;
|
|
constexpr int PackSize = VEC_16B / sizeof(T);
|
|
PADDLE_ENFORCE_EQ(size_per_head % PackSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"dim_head=%d must be divisible by vec_size=%d",
|
|
size_per_head,
|
|
PackSize));
|
|
const int32_t pack_num = elem_cnt / PackSize;
|
|
const int32_t blocksize = 128;
|
|
int32_t grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
fusedQKV_transpose_split_kernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(q_buf,
|
|
k_buf,
|
|
v_buf,
|
|
qkv,
|
|
padding_offset,
|
|
seq_lens,
|
|
elem_cnt,
|
|
batch_size,
|
|
seq_len,
|
|
token_num,
|
|
q_head_num,
|
|
kv_head_num,
|
|
size_per_head);
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void write_pre_cache_to_kv_buffer(
|
|
T *k_buf, // [bsz, num_head, seq_len + pre_cache_length, head_dim]
|
|
T *v_buf,
|
|
const T *pre_key_cache, // [bsz, num_head, pre_cache_length, head_dim]
|
|
const T *pre_value_cache,
|
|
const int *seq_lens,
|
|
const int batch_size,
|
|
const int pre_cache_length,
|
|
const int num_head,
|
|
const int head_dim,
|
|
const int max_len_this_time,
|
|
const int elem_cnt) {
|
|
const int32_t hidden_size = pre_cache_length * head_dim;
|
|
const int32_t cache_hidden_size = num_head * hidden_size;
|
|
const int32_t fused_hidden_size = 2 * cache_hidden_size;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
for (int32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int32_t batch_id = linear_index / fused_hidden_size;
|
|
if (seq_lens[batch_id] == 0) continue;
|
|
|
|
const int32_t cache_seq_id = (linear_index % hidden_size) / head_dim;
|
|
const int32_t head_id = (linear_index % cache_hidden_size) / hidden_size;
|
|
const int32_t size_id = linear_index % head_dim;
|
|
const int32_t kv_id =
|
|
(linear_index % fused_hidden_size) / cache_hidden_size;
|
|
|
|
const int32_t read_id = batch_id * cache_hidden_size +
|
|
head_id * hidden_size + cache_seq_id * head_dim +
|
|
size_id;
|
|
if (kv_id == 0) {
|
|
Load<T, VecSize>(&pre_key_cache[read_id], &src_vec);
|
|
} else {
|
|
Load<T, VecSize>(&pre_value_cache[read_id], &src_vec);
|
|
}
|
|
|
|
const int tmp_max_len_this_time = max_len_this_time + pre_cache_length;
|
|
const int32_t write_idx =
|
|
batch_id * num_head * tmp_max_len_this_time * head_dim +
|
|
head_id * tmp_max_len_this_time * head_dim + cache_seq_id * head_dim +
|
|
size_id;
|
|
if (kv_id == 0) {
|
|
Store<T, VecSize>(src_vec, &k_buf[write_idx]);
|
|
} else {
|
|
Store<T, VecSize>(src_vec, &v_buf[write_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void fusedQKV_transpose_split_kernel(T *q_buf,
|
|
T *k_buf,
|
|
T *v_buf,
|
|
const T *qkv,
|
|
const int *padding_offset,
|
|
const int *seq_lens,
|
|
const int32_t elem_cnt,
|
|
const int batch_size,
|
|
const int max_len_this_time,
|
|
const int seq_len,
|
|
const int pre_cache_length,
|
|
const int token_num,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int size_per_head) {
|
|
const int fused_hidden_size = (q_head_num + 2 * kv_head_num) * size_per_head;
|
|
int64_t global_thread_idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
for (int32_t linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
Load<T, VecSize>(&qkv[linear_index], &src_vec);
|
|
int32_t bias_idx = linear_index % fused_hidden_size;
|
|
const int32_t token_idx = linear_index / fused_hidden_size;
|
|
const int32_t ori_token_idx =
|
|
token_idx + (padding_offset == nullptr ? 0 : padding_offset[token_idx]);
|
|
const int32_t target_batch_id = ori_token_idx / seq_len;
|
|
if (seq_lens[target_batch_id] == 0) continue;
|
|
const int32_t seq_id = ori_token_idx % seq_len;
|
|
|
|
const int32_t head_id = bias_idx / size_per_head;
|
|
const int32_t size_id = linear_index % size_per_head;
|
|
|
|
const int tmp_max_len_this_time =
|
|
max_len_this_time + (head_id < q_head_num ? 0 : pre_cache_length);
|
|
const int tmp_seq_id =
|
|
head_id < q_head_num ? seq_id : seq_id + pre_cache_length;
|
|
|
|
if (head_id < q_head_num) {
|
|
const int write_idx =
|
|
target_batch_id * q_head_num * tmp_max_len_this_time * size_per_head +
|
|
head_id * tmp_max_len_this_time * size_per_head +
|
|
tmp_seq_id * size_per_head + size_id;
|
|
Store<T, VecSize>(src_vec, &q_buf[write_idx]);
|
|
} else if (head_id < q_head_num + kv_head_num) {
|
|
const int write_idx =
|
|
target_batch_id * kv_head_num * tmp_max_len_this_time *
|
|
size_per_head +
|
|
(head_id - q_head_num) * tmp_max_len_this_time * size_per_head +
|
|
tmp_seq_id * size_per_head + size_id;
|
|
Store<T, VecSize>(src_vec, &k_buf[write_idx]);
|
|
} else {
|
|
const int write_idx = target_batch_id * kv_head_num *
|
|
tmp_max_len_this_time * size_per_head +
|
|
(head_id - q_head_num - kv_head_num) *
|
|
tmp_max_len_this_time * size_per_head +
|
|
tmp_seq_id * size_per_head + size_id;
|
|
Store<T, VecSize>(src_vec, &v_buf[write_idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void qkv_transpose_split(
|
|
const GPUContext &dev_ctx,
|
|
T *q_buf,
|
|
T *k_buf,
|
|
T *v_buf,
|
|
const T *qkv,
|
|
const T *pre_key_cache, // [bsz, num_head, pre_cache_length, head_dim]
|
|
const T *pre_value_cache,
|
|
const int *padding_offset,
|
|
const int *seq_lens,
|
|
const int token_num,
|
|
const int batch_size,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int max_len_this_time,
|
|
const int seq_len,
|
|
const int pre_cache_length,
|
|
const int size_per_head) {
|
|
int32_t elem_cnt = token_num * (q_head_num + kv_head_num * 2) * size_per_head;
|
|
|
|
constexpr int PackSize = VEC_16B / sizeof(T);
|
|
PADDLE_ENFORCE_EQ(size_per_head % PackSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"dim_head=%d must be divisible by vec_size=%d",
|
|
size_per_head,
|
|
PackSize));
|
|
int32_t pack_num = elem_cnt / PackSize;
|
|
const int32_t blocksize = 128;
|
|
int32_t grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
fusedQKV_transpose_split_kernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(q_buf,
|
|
k_buf,
|
|
v_buf,
|
|
qkv,
|
|
padding_offset,
|
|
seq_lens,
|
|
elem_cnt,
|
|
batch_size,
|
|
max_len_this_time,
|
|
seq_len,
|
|
pre_cache_length,
|
|
token_num,
|
|
q_head_num,
|
|
kv_head_num,
|
|
size_per_head);
|
|
if (pre_key_cache) {
|
|
// stage 2: write pre_cache to kv_buf
|
|
elem_cnt = batch_size * q_head_num * pre_cache_length * size_per_head * 2;
|
|
pack_num = elem_cnt / PackSize;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
write_pre_cache_to_kv_buffer<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(k_buf,
|
|
v_buf,
|
|
pre_key_cache,
|
|
pre_value_cache,
|
|
seq_lens,
|
|
batch_size,
|
|
pre_cache_length,
|
|
kv_head_num,
|
|
size_per_head,
|
|
max_len_this_time,
|
|
elem_cnt);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void GetDecoderTensorKernel(const T *qkv_out,
|
|
const int *cum_offsets,
|
|
T *qkv_out_decoder,
|
|
const int token_num,
|
|
const int batch_size,
|
|
const int q_head_num,
|
|
const int kv_head_num,
|
|
const int seq_len,
|
|
const int dim_head,
|
|
const int elem_nums,
|
|
const int qkv_out_nums) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
const int32_t fused_hidden_size = (q_head_num + 2 * kv_head_num) * dim_head;
|
|
const int global_idx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
for (int i = global_idx * VecSize; i < elem_nums;
|
|
i += gridDim.x * blockDim.x * VecSize) {
|
|
const int bi = i / fused_hidden_size;
|
|
const int bias_idx = i % fused_hidden_size;
|
|
const int ori_token_idx = bi * seq_len - cum_offsets[bi];
|
|
const int src_offset = ori_token_idx * fused_hidden_size + bias_idx;
|
|
if (src_offset >= qkv_out_nums) continue;
|
|
Load<T, VecSize>(&qkv_out[src_offset], &src_vec);
|
|
Store<T, VecSize>(src_vec, &qkv_out_decoder[i]);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void GetDecoderRoPEKernel(const T *rope_emb,
|
|
T *rope_out_emb,
|
|
const int rope_bsz,
|
|
const int batch_size,
|
|
const int seq_len,
|
|
const int dim_head,
|
|
const int elem_nums) {
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
const T *rope_cos_emb = rope_emb;
|
|
const T *rope_sin_emb = rope_emb + rope_bsz * seq_len * dim_head;
|
|
T *cos_emb = rope_out_emb;
|
|
T *sin_emb = rope_out_emb + batch_size * dim_head;
|
|
const int global_idx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
for (int i = global_idx * VecSize; i < elem_nums;
|
|
i += gridDim.x * blockDim.x * VecSize) {
|
|
const int bi = i / dim_head;
|
|
const int src_offset = bi * seq_len * dim_head + i % dim_head;
|
|
Load<T, VecSize>(&rope_cos_emb[src_offset], &src_vec);
|
|
Store<T, VecSize>(src_vec, &cos_emb[i]);
|
|
Load<T, VecSize>(&rope_sin_emb[src_offset], &src_vec);
|
|
Store<T, VecSize>(src_vec, &sin_emb[i]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void GetDecoderTensor(const GPUContext &dev_ctx,
|
|
const DenseTensor &qkv_out,
|
|
const DenseTensor *rope_emb,
|
|
const int *cum_offsets,
|
|
DenseTensor *qkv_out_decoder,
|
|
DenseTensor *rope_out_emb,
|
|
const int token_num,
|
|
const int batch_size,
|
|
const int q_num_head,
|
|
const int kv_num_head,
|
|
const int seq_len,
|
|
const int dim_head) {
|
|
// qkv_out: [token_num, 2 * kv_num_head + q_num_head, dim_head] -> [bs, 1, 2 *
|
|
// kv_num_head + q_num_head, dim_head] rope: [2, bsz, 1, seq_len, dim_head] ->
|
|
// [2, bsz, 1, 1, dim_head]
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
int64_t elem_nums = qkv_out_decoder->numel();
|
|
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
int64_t qkv_out_nums = qkv_out.numel();
|
|
|
|
constexpr int PackSize = VEC_16B / sizeof(T);
|
|
PADDLE_ENFORCE_EQ(
|
|
dim_head % PackSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"dim_head=%d must be divisible by vec_size=%d", dim_head, PackSize));
|
|
int pack_num = elem_nums / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
GetDecoderTensorKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
qkv_out.data<T>(),
|
|
cum_offsets,
|
|
qkv_out_decoder->data<T>(),
|
|
token_num,
|
|
batch_size,
|
|
q_num_head,
|
|
kv_num_head,
|
|
seq_len,
|
|
dim_head,
|
|
elem_nums,
|
|
qkv_out_nums);
|
|
if (rope_out_emb) {
|
|
elem_nums = rope_out_emb->numel() / 2;
|
|
pack_num = elem_nums / PackSize;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
GetDecoderRoPEKernel<float, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
rope_emb->data<float>(),
|
|
rope_out_emb->data<float>(),
|
|
rope_emb->dims()[1],
|
|
batch_size,
|
|
seq_len,
|
|
dim_head,
|
|
elem_nums);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
struct MaxOp {
|
|
__device__ __forceinline__ T operator()(const T &a, const T &b) const {
|
|
return max(a, b);
|
|
}
|
|
};
|
|
|
|
template <int THREADBLOCK_SIZE>
|
|
__global__ void GetMaxLenKernel(const int *seq_lens,
|
|
int *max_len,
|
|
const int64_t batch_size) {
|
|
const int64_t tid = threadIdx.x;
|
|
|
|
typedef cub::BlockReduce<int, THREADBLOCK_SIZE> BlockReduce;
|
|
__shared__ typename BlockReduce::TempStorage temp_storage;
|
|
|
|
int max_len_this_thread = 0;
|
|
for (int64_t i = tid; i < batch_size; i += blockDim.x) {
|
|
max_len_this_thread = max(seq_lens[i], max_len_this_thread);
|
|
}
|
|
int total =
|
|
BlockReduce(temp_storage).Reduce(max_len_this_thread, MaxOp<int>());
|
|
if (tid == 0) {
|
|
*max_len = total;
|
|
}
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void InitOutValueKernel(T *output_data,
|
|
const int64_t numel,
|
|
const T init_value) {
|
|
const int tid = threadIdx.x;
|
|
const int bid = blockIdx.x;
|
|
int64_t global_thread_idx = bid * static_cast<int64_t>(blockDim.x) + tid;
|
|
|
|
for (int linear_index = global_thread_idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < numel;
|
|
linear_index += step) {
|
|
for (int i = 0; i < VecSize; i++) {
|
|
output_data[linear_index + i] = init_value;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void InitValue(const GPUContext &dev_ctx,
|
|
T *output_data,
|
|
const int64_t numel,
|
|
const T init_value) {
|
|
constexpr int PackSize = VEC_16B / sizeof(T);
|
|
PADDLE_ENFORCE_EQ(
|
|
numel % PackSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"numel=%d must be divisible by vec_size=%d", numel, PackSize));
|
|
const int pack_num = numel / PackSize;
|
|
const int blocksize = 128;
|
|
int grid_size = 1;
|
|
GetNumBlocks(pack_num, &grid_size);
|
|
InitOutValueKernel<T, PackSize>
|
|
<<<grid_size, blocksize, 0, dev_ctx.stream()>>>(
|
|
output_data, numel, init_value);
|
|
}
|
|
|
|
template <typename T, int VecSize>
|
|
__global__ void TransposeRemovingPadding(const T *input_data,
|
|
const int *seq_lens,
|
|
T *output_data,
|
|
const int batch_size,
|
|
const int num_head,
|
|
const int max_len_this_time,
|
|
const int seq_len,
|
|
const int head_dim,
|
|
const int token_num,
|
|
const int elem_cnt,
|
|
const int *padding_offset) {
|
|
// transpose and remove padding
|
|
// [batch_size, num_head, max_len_this_time, head_dim] -> [token_num,
|
|
// num_head, head_dim]
|
|
int64_t idx =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int dim_embed = num_head * head_dim;
|
|
using LoadT = AlignedVector<T, VecSize>;
|
|
LoadT src_vec;
|
|
|
|
for (int32_t linear_index = idx * VecSize,
|
|
step = gridDim.x * blockDim.x * VecSize;
|
|
linear_index < elem_cnt;
|
|
linear_index += step) {
|
|
const int token_idx = linear_index / dim_embed;
|
|
const int ori_token_idx =
|
|
token_idx + (padding_offset == nullptr ? 0 : padding_offset[token_idx]);
|
|
const int ori_batch_id = ori_token_idx / seq_len;
|
|
if (seq_lens && seq_lens[ori_batch_id] == 0) continue;
|
|
const int ori_seq_id = ori_token_idx % seq_len;
|
|
const int ori_head_id = (linear_index % dim_embed) / head_dim;
|
|
const int ori_head_lane = (linear_index % dim_embed) % head_dim;
|
|
const int ori_idx = ori_batch_id * num_head * max_len_this_time * head_dim +
|
|
ori_head_id * max_len_this_time * head_dim +
|
|
ori_seq_id * head_dim + ori_head_lane;
|
|
Load<T, VecSize>(&input_data[ori_idx], &src_vec);
|
|
Store<T, VecSize>(src_vec, &output_data[linear_index]);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void InvokeTransposeRemovePadding(const GPUContext &dev_ctx,
|
|
const T *input_data,
|
|
const int *seq_lens,
|
|
T *output_data,
|
|
const int batch_size,
|
|
const int num_head,
|
|
const int max_len_this_time,
|
|
const int seq_len,
|
|
const int head_dim,
|
|
const int token_num,
|
|
const int *padding_offset) {
|
|
// [batch_size, num_head, max_len_this_time, head_dim] -> [token_num,
|
|
// num_head, head_dim]
|
|
constexpr int VEC_16B = 16;
|
|
const int elem_cnt = token_num * num_head * head_dim;
|
|
constexpr int PackSize = VEC_16B / sizeof(T);
|
|
PADDLE_ENFORCE_EQ(
|
|
head_dim % PackSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"dim_head=%d must be divisible by vec_size=%d", head_dim, PackSize));
|
|
const int32_t pack_num = elem_cnt / PackSize;
|
|
const int32_t block_size = 128;
|
|
int32_t grid_size = (pack_num + block_size - 1) / block_size;
|
|
TransposeRemovingPadding<T, PackSize>
|
|
<<<grid_size, block_size, 0, dev_ctx.stream()>>>(input_data,
|
|
seq_lens,
|
|
output_data,
|
|
batch_size,
|
|
num_head,
|
|
max_len_this_time,
|
|
seq_len,
|
|
head_dim,
|
|
token_num,
|
|
elem_cnt,
|
|
padding_offset);
|
|
}
|
|
|
|
template <typename T>
|
|
static void VLOGMatrix(const T *mat_d,
|
|
int num,
|
|
std::string name,
|
|
int max_num = 10) {
|
|
num = num < max_num ? num : max_num;
|
|
std::vector<T> tmp(num);
|
|
GPU(Memcpy)(tmp.data(), mat_d, sizeof(T) * num, GPU(MemcpyDeviceToHost));
|
|
|
|
std::stringstream ss;
|
|
|
|
for (int i = 0; i < num; ++i) {
|
|
if (std::is_same<T, int8_t>::value) {
|
|
ss << static_cast<int>(tmp[i]) << ", ";
|
|
} else {
|
|
ss << std::setprecision(8) << (float)(tmp[i]) << ", "; // NOLINT
|
|
}
|
|
}
|
|
LOG(INFO) << "===>: " << name;
|
|
LOG(INFO) << ss.str();
|
|
}
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|