Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-11 04:56:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

metal kernel

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-15 12:07:24 +02:00
parent b5464b9310
commit 53845a366e
4 changed files with 45 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-impl.h
+2
View File
@@ -329,6 +329,7 @@ typedef struct {
uint64_t nb3;
int32_t n_past;
int32_t n_dims;
int32_t n_offs;
int32_t n_ctx_orig;
float freq_base;
float freq_scale;
@@ -341,6 +342,7 @@ typedef struct {
int32_t sect_2;
int32_t sect_3;
bool src2;
bool inplace;
} ggml_metal_kargs_rope;
typedef struct {
ggml/src/ggml-metal/ggml-metal-ops.cpp
+7
View File
@@ -3884,6 +3884,11 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
const int sect_2 = ((const int32_t *) op->op_params)[13];
const int sect_3 = ((const int32_t *) op->op_params)[14];
const int n_offs = ((const int32_t *) op->op_params)[15];
// when dst aliases src0, the channels outside the rotated window already hold the correct data
const bool inplace = op->data == op->src[0]->data;
ggml_metal_kargs_rope args = {
/*.ne00 =*/ ne00,
/*.ne01 =*/ ne01,
@@ -3903,6 +3908,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
/*.nb3 =*/ nb3,
/*.n_past =*/ n_past,
/*.n_dims =*/ n_dims,
/*.n_offs =*/ n_offs,
/*.n_ctx_orig =*/ n_ctx_orig,
/*.freq_base =*/ freq_base,
/*.freq_scale =*/ freq_scale,
@@ -3915,6 +3921,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
/* sect_2 =*/ sect_2,
/* sect_3 =*/ sect_3,
/* src2 =*/ op->src[2] != nullptr,
/* inplace =*/ inplace,
};
auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op);
ggml/src/ggml-metal/ggml-metal.metal
+31 -16
View File
@@ -4688,14 +4688,15 @@ kernel void kernel_rope_norm(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4706,6 +4707,10 @@ kernel void kernel_rope_norm(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[1] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4741,17 +4746,18 @@ kernel void kernel_rope_neox(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0);
const float x0 = src[0];
const float x1 = src[args.n_dims/2];
@@ -4759,6 +4765,10 @@ kernel void kernel_rope_neox(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4793,8 +4803,9 @@ kernel void kernel_rope_multi(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
// mrope theta calculations
// note: the rest is the same as kernel_rope_neox
@@ -4827,14 +4838,14 @@ kernel void kernel_rope_multi(
}
// end of mrope
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0);
const float x0 = src[0];
const float x1 = src[args.n_dims/2];
@@ -4842,6 +4853,10 @@ kernel void kernel_rope_multi(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
tests/test-backend-ops.cpp
+5 -3
View File
@@ -5331,17 +5331,19 @@ struct test_rope : public test_case {
int v; // view (1 : non-contiguous a)
bool forward;
bool inplace;
int n_offs; // offset of the rotated dims window, set via ggml_rope_set_offset()
std::string vars() override {
// forward can be inferred from the op, does not need to be printed
return VARS_TO_STR11(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace);
return VARS_TO_STR12(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace, n_offs);
}
test_rope(ggml_type type = GGML_TYPE_F32,
std::array<int64_t, 4> ne_a = {10, 5, 3, 1},
int n_dims = 10, int mode = GGML_ROPE_TYPE_NORMAL, int n_ctx = 512, float fs = 1.0f,
float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false)
: type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace) {}
float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false,
int n_offs = 0)
: type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace), n_offs(n_offs) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a;