Skip to content

Commit 6aeb46b

Browse files
cebtenzzrecebtenzzre
committed
CUDA implementation
1 parent 8dec38c commit 6aeb46b

File tree

3 files changed

+115
-42
lines changed

3 files changed

+115
-42
lines changed

‎ggml-cuda.cu

+74-10
Original file line numberDiff line numberDiff line change
@@ -3558,9 +3558,49 @@ static __global__ void cpy_f32_f16(const char * cx, char * cdst, const int ne,
35583558
cpy_1(cx + x_offset, cdst + dst_offset);
35593559
}
35603560

3561+
static __device__ float rope_ntkv2_ramp(const float low, const float high, const int i0) {
3562+
const float y = (i0 / 2 - low) / min(0.001f, high - low);
3563+
return 1.0f - min(1.0f, max(0.0f, y));
3564+
}
3565+
3566+
struct rope_corr_factors {
3567+
float v[4];
3568+
};
3569+
3570+
// NTKv2 algorithm based on LlamaPartNTKScaledRotaryEmbedding.py from https://github.com/jquesnelle/scaled-rope
3571+
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
3572+
static __device__ float rope_ntkv2(
3573+
const float theta_base,
3574+
const float theta_linear,
3575+
const float theta_ntk,
3576+
const rope_corr_factors corr_factors,
3577+
const int64_t i0,
3578+
const float ntk_factor,
3579+
const float ext_factor) {
3580+
float ramp_mix;
3581+
float theta;
3582+
3583+
ramp_mix = rope_ntkv2_ramp(corr_factors.v[0], corr_factors.v[1], i0) * ntk_factor;
3584+
theta = theta_linear * (1 - ramp_mix) + theta_ntk * ramp_mix;
3585+
3586+
ramp_mix = rope_ntkv2_ramp(corr_factors.v[2], corr_factors.v[3], i0) * ext_factor;
3587+
theta = theta * (1 - ramp_mix) + theta_base * ramp_mix;
3588+
return theta;
3589+
}
3590+
35613591
// rope == RoPE == rotary positional embedding
3562-
static __global__ void rope_f32(const float * x, float * dst, const int ncols, const float p0,
3563-
const float p_delta, const int p_delta_rows, const float theta_scale) {
3592+
static __global__ void rope_f32(
3593+
const float * x,
3594+
float * dst,
3595+
const int ncols,
3596+
const float freq_scale,
3597+
const float ntk_factor,
3598+
const float ext_factor,
3599+
const float theta_scale,
3600+
const float theta_ntk_scale,
3601+
const float p0,
3602+
const int p_delta_rows,
3603+
const rope_corr_factors corr_factors) {
35643604
const int col = 2*(blockDim.x*blockIdx.x + threadIdx.x);
35653605

35663606
if (col >= ncols) {
@@ -3570,7 +3610,11 @@ static __global__ void rope_f32(const float * x, float * dst, const int ncols, c
35703610
const int row = blockDim.y*blockIdx.y + threadIdx.y;
35713611
const int i = row*ncols + col;
35723612

3573-
const float theta = (p0 + p_delta * (row/p_delta_rows))*powf(theta_scale, col/2);
3613+
const float p = p0 + row / p_delta_rows;
3614+
const float theta_base = p*powf(theta_scale, col/2);
3615+
const float theta_linear = freq_scale * theta_base;
3616+
const float theta_ntk = p*powf(theta_ntk_scale, col/2);
3617+
const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors, col, ntk_factor, ext_factor);
35743618
const float sin_theta = sinf(theta);
35753619
const float cos_theta = cosf(theta);
35763620

@@ -4234,13 +4278,26 @@ static void scale_f32_cuda(const float * x, float * dst, const float scale, cons
42344278
scale_f32<<<num_blocks, CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, k);
42354279
}
42364280

4237-
static void rope_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, const float p0,
4238-
const float p_delta, const int p_delta_rows, const float theta_scale, cudaStream_t stream) {
4281+
static void rope_f32_cuda(
4282+
const float * x,
4283+
float * dst,
4284+
const int ncols,
4285+
const int nrows,
4286+
const float freq_scale,
4287+
const float ntk_factor,
4288+
const float ext_factor,
4289+
const float theta_scale,
4290+
const float theta_ntk_scale,
4291+
const float p0,
4292+
const int p_delta_rows,
4293+
const rope_corr_factors corr_factors,
4294+
cudaStream_t stream) {
42394295
GGML_ASSERT(nrows % 2 == 0);
42404296
const dim3 block_dims(2*CUDA_ROPE_BLOCK_SIZE, 1, 1);
42414297
const int num_blocks_x = (ncols + 2*CUDA_ROPE_BLOCK_SIZE - 1) / (2*CUDA_ROPE_BLOCK_SIZE);
42424298
const dim3 block_nums(num_blocks_x, nrows, 1);
4243-
rope_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols, p0, p_delta, p_delta_rows, theta_scale);
4299+
rope_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols, freq_scale, ntk_factor, ext_factor, theta_scale,
4300+
theta_ntk_scale, p0, p_delta_rows, corr_factors);
42444301
}
42454302

42464303
static void rope_glm_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, const float p, const float block_p, const float theta_scale, cudaStream_t stream) {
@@ -4941,11 +4998,13 @@ inline void ggml_cuda_op_rope(
49414998
const int n_dims = ((int32_t *) dst->op_params)[1];
49424999
const int mode = ((int32_t *) dst->op_params)[2];
49435000
const int n_ctx = ((int32_t *) dst->op_params)[3];
4944-
// RoPE alteration for extended context
49455001

4946-
float freq_base, freq_scale;
5002+
// RoPE alteration for extended context
5003+
float freq_base, freq_scale, ntk_factor, ext_factor;
49475004
memcpy(&freq_base, (int32_t *) dst->op_params + 4, sizeof(float));
49485005
memcpy(&freq_scale, (int32_t *) dst->op_params + 5, sizeof(float));
5006+
memcpy(&ntk_factor, (int32_t *) dst->op_params + 6, sizeof(float));
5007+
memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float));
49495008

49505009
const float theta_scale = powf(freq_base, -2.0f/n_dims);
49515010

@@ -4958,8 +5017,13 @@ inline void ggml_cuda_op_rope(
49585017
const float block_p = max(p - (n_ctx - 2.f), 0.f);
49595018
rope_glm_f32_cuda(src0_ddf_i, dst_ddf_i, ne00, i01_diff, id_p, block_p, theta_scale, cudaStream_main);
49605019
} else {
4961-
const float p0 = (((mode & 1) == 0 ? n_past : 0)) * freq_scale;
4962-
rope_f32_cuda(src0_ddf_i, dst_ddf_i, ne00, i01_diff, p0, freq_scale, ne01, theta_scale, cudaStream_main);
5020+
const float p0 = (mode & 1) == 0 ? n_past : 0;
5021+
const float theta_ntk_scale = powf(freq_base * powf(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims);
5022+
rope_corr_factors corr_factors;
5023+
ggml_rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors.v);
5024+
5025+
rope_f32_cuda(src0_ddf_i, dst_ddf_i, ne00, i01_diff, freq_scale, ntk_factor, ext_factor, theta_scale,
5026+
theta_ntk_scale, p0, ne01, corr_factors, cudaStream_main);
49635027
}
49645028

49655029
(void) src1;

‎ggml.c

+38-32
Original file line numberDiff line numberDiff line change
@@ -12012,11 +12012,6 @@ static void ggml_compute_forward_clamp(
1201212012

1201312013
// ggml_compute_forward_rope
1201412014

12015-
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
12016-
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
12017-
#define NTKV2_MAX_POS_EMB 2048
12018-
#define NTKV2_CORRECTION_FACTOR(n_rot) (__builtin_logf(NTKV2_MAX_POS_EMB / ((n_rot) * 2 * (float)M_PI)) / 2)
12019-
1202012015
static inline float rope_ntkv2_ramp(const float low, const float high, const int i0) {
1202112016
const float y = (i0 / 2 - low) / MIN(0.001f, high - low);
1202212017
return 1 - MIN(1, MAX(0, y));
@@ -12026,36 +12021,43 @@ static inline float rope_ntkv2_ramp(const float low, const float high, const int
1202612021
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
1202712022
static float rope_ntkv2(
1202812023
const float theta_base,
12024+
const float theta_linear,
1202912025
const float theta_ntk,
12030-
const float dims_over_base,
12031-
const float freq_scale,
12026+
const float corr_factors[4],
1203212027
const int64_t i0,
1203312028
const float ntk_factor,
12034-
const float ext_factor,
12035-
const int n_dims) {
12029+
const float ext_factor) {
12030+
float ramp_mix;
12031+
float theta;
12032+
12033+
ramp_mix = rope_ntkv2_ramp(corr_factors[0], corr_factors[1], i0) * ntk_factor;
12034+
theta = theta_linear * (1 - ramp_mix) + theta_ntk * ramp_mix;
12035+
12036+
ramp_mix = rope_ntkv2_ramp(corr_factors[2], corr_factors[3], i0) * ext_factor;
12037+
theta = theta * (1 - ramp_mix) + theta_base * ramp_mix;
12038+
return theta;
12039+
}
12040+
12041+
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
12042+
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
12043+
static float ggml_rope_ntkv2_corr_factor(const int n_dims, const float n_rot, const float base) {
12044+
static const float max_pos_emb = 2048;
12045+
return n_dims * logf(max_pos_emb / (n_rot * 2 * (float)M_PI)) / (2 * logf(base));
12046+
}
12047+
12048+
void ggml_rope_ntkv2_corr_factors(int n_dims, const float freq_base, float factors[4]) {
1203612049
// Interpolation constants found experimentally for LLaMA (might not be totally optimal though)
1203712050
// Do not change unless there is a good reason for doing so!
12038-
static const float BETA_0 = 1.75f;
12039-
static const float BETA_1 = 1.25f;
12051+
static const float BETA_0 = 1.75f;
12052+
static const float BETA_1 = 1.25f;
1204012053
static const float GAMMA_0 = 16.0f;
1204112054
static const float GAMMA_1 = 2.0f;
1204212055

12043-
static const float low_1p = NTKV2_CORRECTION_FACTOR(BETA_0);
12044-
static const float high_1p = NTKV2_CORRECTION_FACTOR(BETA_1);
12045-
static const float low_2p = NTKV2_CORRECTION_FACTOR(GAMMA_0);
12046-
static const float high_2p = NTKV2_CORRECTION_FACTOR(GAMMA_1);
12047-
1204812056
// start and end correction factors
12049-
const float low_1 = MAX(0, floorf(low_1p * dims_over_base));
12050-
const float high_1 = MIN(n_dims - 1, ceilf(high_1p * dims_over_base));
12051-
const float low_2 = MAX(0, floorf(low_2p * dims_over_base));
12052-
const float high_2 = MIN(n_dims - 1, ceilf(high_2p * dims_over_base));
12053-
12054-
const float theta_linear = freq_scale * theta_base;
12055-
const float ramp_mix = rope_ntkv2_ramp(low_1, high_1, i0) * ntk_factor;
12056-
const float theta_mix = theta_linear * (1 - ramp_mix) + theta_ntk * ramp_mix;
12057-
const float ramp_final = rope_ntkv2_ramp(low_2, high_2, i0) * ext_factor;
12058-
return theta_mix * (1 - ramp_final) + theta_base * ramp_final;
12057+
factors[0] = MAX(0, floorf(ggml_rope_ntkv2_corr_factor(n_dims, BETA_0, freq_base)));
12058+
factors[1] = MIN(n_dims - 1, ceilf(ggml_rope_ntkv2_corr_factor(n_dims, BETA_1, freq_base)));
12059+
factors[2] = MAX(0, floorf(ggml_rope_ntkv2_corr_factor(n_dims, GAMMA_0, freq_base)));
12060+
factors[3] = MIN(n_dims - 1, ceilf(ggml_rope_ntkv2_corr_factor(n_dims, GAMMA_1, freq_base)));
1205912061
}
1206012062

1206112063
static void ggml_compute_forward_rope_f32(
@@ -12110,7 +12112,8 @@ static void ggml_compute_forward_rope_f32(
1211012112

1211112113
const float theta_scale = powf(freq_base, -2.0f/n_dims);
1211212114
const float theta_ntk_scale = powf(freq_base * powf(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims);
12113-
const float dims_over_base = n_dims / logf(freq_base);
12115+
float corr_factors[4];
12116+
ggml_rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors);
1211412117

1211512118
const bool is_neox = mode & 2;
1211612119
const bool is_glm = mode & 4;
@@ -12152,8 +12155,9 @@ static void ggml_compute_forward_rope_f32(
1215212155
}
1215312156
} else if (!is_neox) {
1215412157
for (int64_t i0 = 0; i0 < ne0; i0 += 2) {
12155-
const float theta = rope_ntkv2(theta_base, theta_ntk, dims_over_base,
12156-
freq_scale, i0, ntk_factor, ext_factor, n_dims);
12158+
const float theta_linear = freq_scale * theta_base;
12159+
const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors,
12160+
i0, ntk_factor, ext_factor);
1215712161
const float cos_theta = cosf(theta);
1215812162
const float sin_theta = sinf(theta);
1215912163

@@ -12250,7 +12254,8 @@ static void ggml_compute_forward_rope_f16(
1225012254

1225112255
const float theta_scale = powf(freq_base, -2.0f/n_dims);
1225212256
const float theta_ntk_scale = powf(freq_base * powf(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims);
12253-
const float dims_over_base = n_dims / logf(freq_base);
12257+
float corr_factors[4];
12258+
ggml_rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors);
1225412259

1225512260
const bool is_neox = mode & 2;
1225612261
const bool is_glm = mode & 4;
@@ -12292,8 +12297,9 @@ static void ggml_compute_forward_rope_f16(
1229212297
}
1229312298
} if (!is_neox) {
1229412299
for (int64_t i0 = 0; i0 < ne0; i0 += 2) {
12295-
const float theta = rope_ntkv2(theta_base, theta_ntk, dims_over_base,
12296-
freq_scale, i0, ntk_factor, ext_factor, n_dims);
12300+
const float theta_linear = freq_scale * theta_base;
12301+
const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors,
12302+
i0, ntk_factor, ext_factor);
1229712303
const float cos_theta = cosf(theta);
1229812304
const float sin_theta = sinf(theta);
1229912305

‎ggml.h

+3
Original file line numberDiff line numberDiff line change
@@ -1211,6 +1211,9 @@ extern "C" {
12111211
float ntk_factor,
12121212
float ext_factor);
12131213

1214+
// compute correction factors for NTKv2 RoPE scaling
1215+
void ggml_rope_ntkv2_corr_factors(int n_dims, const float freq_base, float factors[4]);
1216+
12141217
// rotary position embedding backward, i.e compute dx from dy
12151218
// a - dy
12161219
GGML_API struct ggml_tensor * ggml_rope_back(

0 commit comments

Comments
 (0)