wehub-resource-sync
339 lines
13 KiB
C++
339 lines
13 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/enforce.h"
|
|
#include "paddle/phi/kernels/funcs/aligned_vector.h"
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
template <typename T, typename IndexT>
|
|
__device__ void set_sin_cos_shared_mem(const T* sin,
|
|
const T* cos,
|
|
const int64_t* position_ids,
|
|
const bool flag_sin_cos,
|
|
const float rotary_emb_base,
|
|
const IndexT seq_len,
|
|
const IndexT s_id,
|
|
const IndexT b_id,
|
|
const IndexT d,
|
|
const IndexT d2,
|
|
float* shared_mem_sin,
|
|
float* shared_mem_cos) {
|
|
IndexT tid = static_cast<IndexT>(threadIdx.x) * blockDim.y + threadIdx.y;
|
|
for (IndexT i = tid; i < d2; i += blockDim.x * blockDim.y) {
|
|
int64_t pos = s_id;
|
|
if (position_ids) {
|
|
pos = position_ids[b_id * seq_len + s_id];
|
|
}
|
|
|
|
if (flag_sin_cos) {
|
|
shared_mem_sin[i] = static_cast<float>(sin[pos * d2 + i]);
|
|
shared_mem_cos[i] = static_cast<float>(cos[pos * d2 + i]);
|
|
} else {
|
|
float idx = static_cast<float>((i / 2) * 2);
|
|
float inv_freq =
|
|
1.0f / powf(rotary_emb_base, idx / static_cast<float>(d2));
|
|
float freq = static_cast<float>(pos) * inv_freq;
|
|
sincosf(freq, &shared_mem_sin[i], &shared_mem_cos[i]);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
template <typename T, typename IndexT>
|
|
__global__ void FusedRopeKernelImpl(const T* src,
|
|
const T* sin,
|
|
const T* cos,
|
|
T* dst,
|
|
const int64_t* position_ids,
|
|
const bool flag_sin_cos,
|
|
const bool use_neox_rotary_style,
|
|
const IndexT h,
|
|
const IndexT d,
|
|
const IndexT d2,
|
|
const IndexT stride_s,
|
|
const IndexT stride_b,
|
|
const IndexT stride_h,
|
|
const IndexT stride_d,
|
|
const IndexT o_stride_s,
|
|
const IndexT o_stride_b,
|
|
const IndexT o_stride_h,
|
|
const IndexT o_stride_d,
|
|
const float rotary_emb_base,
|
|
const IndexT seq_len) {
|
|
IndexT s_id = blockIdx.x;
|
|
IndexT b_id = blockIdx.y;
|
|
|
|
IndexT offset_block = s_id * stride_s + b_id * stride_b;
|
|
IndexT offset_block_dst = s_id * o_stride_s + b_id * o_stride_b;
|
|
|
|
extern __shared__ float shared_mem[];
|
|
float* shared_mem_cos = shared_mem;
|
|
float* shared_mem_sin = shared_mem + d2;
|
|
|
|
set_sin_cos_shared_mem<T>(sin,
|
|
cos,
|
|
position_ids,
|
|
flag_sin_cos,
|
|
rotary_emb_base,
|
|
seq_len,
|
|
s_id,
|
|
b_id,
|
|
d,
|
|
d2,
|
|
shared_mem_sin,
|
|
shared_mem_cos);
|
|
|
|
#pragma unroll
|
|
for (IndexT h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
|
|
#pragma unroll
|
|
for (IndexT d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {
|
|
float v_cos = shared_mem_cos[d_id];
|
|
float v_sin = shared_mem_sin[d_id];
|
|
IndexT offset_src = offset_block + h_id * stride_h + d_id * stride_d;
|
|
IndexT offset_dst =
|
|
offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
|
|
float v_src = static_cast<float>(src[offset_src]);
|
|
float v_src_rotate;
|
|
if (!use_neox_rotary_style) {
|
|
v_src_rotate =
|
|
(d_id + d2 / 2 < d2)
|
|
? -static_cast<float>(src[offset_src + (d2 / 2) * stride_d])
|
|
: static_cast<float>(
|
|
src[offset_src + (d2 / 2 - d2) * stride_d]);
|
|
} else {
|
|
v_src_rotate = (d_id % 2 == 0)
|
|
? -static_cast<float>(src[offset_src + stride_d])
|
|
: static_cast<float>(src[offset_src - stride_d]);
|
|
}
|
|
dst[offset_dst] = static_cast<T>(v_src * v_cos + v_src_rotate * v_sin);
|
|
}
|
|
}
|
|
|
|
// copy the rest
|
|
if (d > d2) {
|
|
#pragma unroll
|
|
for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {
|
|
#pragma unroll
|
|
for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
|
|
int offset_src = offset_block + h_id * stride_h + d_id * stride_d;
|
|
int offset_dst =
|
|
offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
|
|
dst[offset_dst] = src[offset_src];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename IndexT>
|
|
__global__ void FusedRopeGradKernelImpl(const T* src,
|
|
const T* sin,
|
|
const T* cos,
|
|
T* dst,
|
|
const int64_t* position_ids,
|
|
const bool flag_sin_cos,
|
|
const bool use_neox_rotary_style,
|
|
const IndexT h,
|
|
const IndexT d,
|
|
const IndexT d2,
|
|
const IndexT stride_s,
|
|
const IndexT stride_b,
|
|
const IndexT stride_h,
|
|
const IndexT stride_d,
|
|
const IndexT o_stride_s,
|
|
const IndexT o_stride_b,
|
|
const IndexT o_stride_h,
|
|
const IndexT o_stride_d,
|
|
const float rotary_emb_base,
|
|
const IndexT seq_len) {
|
|
IndexT s_id = blockIdx.x;
|
|
IndexT b_id = blockIdx.y;
|
|
|
|
IndexT offset_block = s_id * stride_s + b_id * stride_b;
|
|
IndexT offset_block_dst = s_id * o_stride_s + b_id * o_stride_b;
|
|
|
|
extern __shared__ float shared_mem[];
|
|
float* shared_mem_cos = shared_mem;
|
|
float* shared_mem_sin = shared_mem + d2;
|
|
|
|
set_sin_cos_shared_mem<T>(sin,
|
|
cos,
|
|
position_ids,
|
|
flag_sin_cos,
|
|
rotary_emb_base,
|
|
seq_len,
|
|
s_id,
|
|
b_id,
|
|
d,
|
|
d2,
|
|
shared_mem_sin,
|
|
shared_mem_cos);
|
|
|
|
#pragma unroll
|
|
for (IndexT h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
|
|
#pragma unroll
|
|
for (IndexT d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {
|
|
IndexT offset_src = offset_block + h_id * stride_h + d_id * stride_d;
|
|
IndexT offset_dst =
|
|
offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
|
|
float v_src = static_cast<float>(src[offset_src]);
|
|
float v_cos = shared_mem_cos[d_id];
|
|
float v_src_rotate, v_sin;
|
|
if (!use_neox_rotary_style) {
|
|
if (d_id + d2 / 2 < d2) {
|
|
v_src_rotate =
|
|
static_cast<float>(src[offset_src + (d2 / 2) * stride_d]);
|
|
v_sin = shared_mem_sin[d_id + d2 / 2];
|
|
} else {
|
|
v_src_rotate =
|
|
static_cast<float>(src[offset_src + (d2 / 2 - d2) * stride_d]);
|
|
v_sin = -shared_mem_sin[d_id + d2 / 2 - d2];
|
|
}
|
|
} else {
|
|
if (d_id % 2 == 0) {
|
|
v_src_rotate = static_cast<float>(src[offset_src + stride_d]);
|
|
v_sin = shared_mem_sin[d_id + 1];
|
|
} else {
|
|
v_src_rotate = static_cast<float>(src[offset_src - stride_d]);
|
|
v_sin = -shared_mem_sin[d_id - 1];
|
|
}
|
|
}
|
|
dst[offset_dst] = static_cast<T>(v_src * v_cos + v_src_rotate * v_sin);
|
|
}
|
|
}
|
|
|
|
// copy the rest
|
|
if (d > d2) {
|
|
#pragma unroll
|
|
for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {
|
|
#pragma unroll
|
|
for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
|
|
int offset_src = offset_block + h_id * stride_h + d_id * stride_d;
|
|
int offset_dst =
|
|
offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
|
|
dst[offset_dst] = src[offset_src];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename IndexT>
|
|
using FusedRopeKernelFunc = void (*)(const T*,
|
|
const T*,
|
|
const T*,
|
|
T*,
|
|
const int64_t*,
|
|
const bool,
|
|
const bool,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const IndexT,
|
|
const float,
|
|
const IndexT);
|
|
|
|
template <typename T>
|
|
void FusedRopeKernelLauncher(const T* src,
|
|
const T* sin,
|
|
const T* cos,
|
|
T* dst,
|
|
FusedRopeKernelFunc<T, int> kernel_int32,
|
|
FusedRopeKernelFunc<T, int64_t> kernel_int64,
|
|
const int64_t* position_ids,
|
|
const bool flag_sin_cos,
|
|
const bool use_neox_rotary_style,
|
|
const int64_t h,
|
|
const int64_t d,
|
|
const int64_t d2,
|
|
const int64_t stride_s,
|
|
const int64_t stride_b,
|
|
const int64_t stride_h,
|
|
const int64_t stride_d,
|
|
const int64_t o_stride_s,
|
|
const int64_t o_stride_b,
|
|
const int64_t o_stride_h,
|
|
const int64_t o_stride_d,
|
|
const float rotary_emb_base,
|
|
const int64_t seq_len,
|
|
const int64_t batch_size,
|
|
const int64_t numel,
|
|
gpuStream_t stream) {
|
|
if (numel <= 0) return;
|
|
const int warps_per_block = h < 16 ? 4 : 8;
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(seq_len, "fused_rope CUDA launch grid.x");
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(batch_size, "fused_rope CUDA launch grid.y");
|
|
dim3 grid(static_cast<uint32_t>(seq_len), static_cast<uint32_t>(batch_size));
|
|
dim3 block(32, warps_per_block); // 32 threads per warp
|
|
size_t shared_mem_size = 2 * d2 * sizeof(float);
|
|
|
|
if (numel <= std::numeric_limits<int>::max()) {
|
|
kernel_int32<<<grid, block, shared_mem_size, stream>>>(
|
|
src,
|
|
sin,
|
|
cos,
|
|
dst,
|
|
position_ids,
|
|
flag_sin_cos,
|
|
use_neox_rotary_style,
|
|
h,
|
|
d,
|
|
d2,
|
|
stride_s,
|
|
stride_b,
|
|
stride_h,
|
|
stride_d,
|
|
o_stride_s,
|
|
o_stride_b,
|
|
o_stride_h,
|
|
o_stride_d,
|
|
rotary_emb_base,
|
|
seq_len);
|
|
} else {
|
|
kernel_int64<<<grid, block, shared_mem_size, stream>>>(
|
|
src,
|
|
sin,
|
|
cos,
|
|
dst,
|
|
position_ids,
|
|
flag_sin_cos,
|
|
use_neox_rotary_style,
|
|
h,
|
|
d,
|
|
d2,
|
|
stride_s,
|
|
stride_b,
|
|
stride_h,
|
|
stride_d,
|
|
o_stride_s,
|
|
o_stride_b,
|
|
o_stride_h,
|
|
o_stride_d,
|
|
rotary_emb_base,
|
|
seq_len);
|
|
}
|
|
}
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|