huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-20 12:06:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Contiguous variant of the rope kernel. (#1929)

Browse Source 
* Contiguous variant of the rope kernel.

* Add the cuda kernel.

* Metal kernel.
This commit is contained in:
Laurent Mazare
2024-03-25 09:11:20 +01:00
committed by GitHub
parent 1b98f84a2b
commit e7f8e72588
5 changed files with 389 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
candle-kernels/src/reduce.cu
View File

@ -160,6 +160,24 @@ __device__ void ropei(const T * src, const T * cos, const T * sin, T * dst, cons
dst[2 * idx + 1] = src[2 * idx] * s + src[2 * idx + 1] * c;
}
template <typename T>
__device__ void rope(const T * src, const T * cos, const T * sin, T * dst, const uint32_t bh, const uint32_t td, const uint32_t d) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (2 * idx > bh * td) return;
uint32_t i_bh = idx / (td / 2);
uint32_t i_td = idx - (td / 2) * i_bh;
uint32_t i_t = i_td / (d / 2);
uint32_t i_d = i_td - (d / 2) * i_t;
uint32_t i1 = i_bh * td + i_t * d + i_d;
uint32_t i2 = i1 + d / 2;
uint32_t i_cs = i_t * (d / 2) + i_d;
T c = cos[i_cs];
T s = sin[i_cs];
dst[i1] = src[i1] * c - src[i2] * s;
dst[i2] = src[i1] * s + src[i2] * c;
}
template <typename T>
__device__ void
@ -416,8 +434,8 @@ fast_argmax(const size_t src_numel, const size_t el_to_sum_per_block,
rmsnorm<TYPENAME>(src, dst, alpha, n_cols, eps); \
} \
#define ROPEI_OP(TYPENAME, FN_NAME) \
extern "C" __global__ void FN_NAME( \
#define ROPE_OP(TYPENAME, FN_NAME, FN_NAME_I) \
extern "C" __global__ void FN_NAME_I( \
const TYPENAME *src, \
const TYPENAME *cos, \
const TYPENAME *sin, \
@ -426,11 +444,21 @@ fast_argmax(const size_t src_numel, const size_t el_to_sum_per_block,
const uint32_t td) { \
ropei<TYPENAME>(src, cos, sin, dst, bh, td); \
} \
extern "C" __global__ void FN_NAME( \
const TYPENAME *src, \
const TYPENAME *cos, \
const TYPENAME *sin, \
TYPENAME *dst, \
const uint32_t bh, \
const uint32_t td, \
const uint32_t d) { \
rope<TYPENAME>(src, cos, sin, dst, bh, td, d); \
} \
#if __CUDA_ARCH__ >= 800
SOFTMAX_OP(__nv_bfloat16, float, softmax_bf16)
RMSNORM_OP(__nv_bfloat16, rmsnorm_bf16)
ROPEI_OP(__nv_bfloat16, rope_i_bf16)
ROPE_OP(__nv_bfloat16, rope_bf16, rope_i_bf16)
SUM_OP(__nv_bfloat16, sum_bf16)
FAST_OP(__nv_bfloat16, fast_min_bf16, fast_max_bf16, fast_argmin_bf16, fast_argmax_bf16, fast_sum_bf16)
#endif
@ -438,7 +466,7 @@ FAST_OP(__nv_bfloat16, fast_min_bf16, fast_max_bf16, fast_argmin_bf16, fast_argm
#if __CUDA_ARCH__ >= 530
SOFTMAX_OP(__half, float, softmax_f16)
RMSNORM_OP(__half, rmsnorm_f16)
ROPEI_OP(__half, rope_i_f16)
ROPE_OP(__half, rope_f16, rope_i_f16)
SUM_OP(__half, sum_f16)
FAST_OP(__half, fast_min_f16, fast_max_f16, fast_argmin_f16, fast_argmax_f16, fast_sum_f16)
#endif
@ -450,8 +478,8 @@ SOFTMAX_OP(float, float, softmax_f32)
SOFTMAX_OP(double, double, softmax_f64)
RMSNORM_OP(float, rmsnorm_f32)
RMSNORM_OP(double, rmsnorm_f64)
ROPEI_OP(float, rope_i_f32)
ROPEI_OP(double, rope_i_f64)
ROPE_OP(float, rope_f32, rope_i_f32)
ROPE_OP(double, rope_f64, rope_i_f64)
FAST_OP(float, fast_min_f32, fast_max_f32, fast_argmin_f32, fast_argmax_f32, fast_sum_f32)
FAST_OP(double, fast_min_f64, fast_max_f64, fast_argmin_f64, fast_argmax_f64, fast_sum_f64)