nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

metal: implement ROPE_MULTI (mrope/imrope) kernels (#2140) 

The Metal backend lacked GGML_ROPE_TYPE_MROPE/IMROPE support, so models
whose GGUF carries rope_sections (e.g. Qwen 3.5 hybrids) could not run
fully offloaded on Apple Silicon.

Add rope_multi_f32/f16 kernels with section-based position handling
(t/h/w/e blocks, 4 position ids per token), imrope's interleaved section
selection, and the corresponding dispatch in ggml-metal.m. Vision-mode
mrope is not implemented and is asserted out explicitly.

Validated on M2: kernel output matches the CPU backend, and Qwen 3.5-9B
(Q4_K_M, -ngl 99) WikiText-2 perplexity over 145 chunks lands within
0.006 of the same model's CPU baseline.
This commit is contained in:
hchengit 2026-07-17 07:49:13 -07:00 committed by GitHub
parent 7ae6b337a7
commit 2b8d0d5011
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 167 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
ggml/src/ggml-metal.m
View File

@ -337,6 +337,8 @@ enum ggml_metal_kernel_type {
GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16, GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16,
GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F32, GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F32,
GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F16, GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F16,
GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F32,
GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F16,
GGML_METAL_KERNEL_TYPE_IM2COL_F16, GGML_METAL_KERNEL_TYPE_IM2COL_F16,
GGML_METAL_KERNEL_TYPE_IM2COL_F32, GGML_METAL_KERNEL_TYPE_IM2COL_F32,
GGML_METAL_KERNEL_TYPE_UPSCALE_F32, GGML_METAL_KERNEL_TYPE_UPSCALE_F32,
@ -1005,6 +1007,8 @@ static struct ggml_backend_metal_context * ggml_metal_init(int n_cb) {
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16, rope_norm_f16, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16, rope_norm_f16, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F32, rope_neox_f32, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F32, rope_neox_f32, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F16, rope_neox_f16, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F16, rope_neox_f16, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F32, rope_multi_f32, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F16, rope_multi_f16, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_IM2COL_F16, im2col_f16, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_IM2COL_F16, im2col_f16, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_IM2COL_F32, im2col_f32, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_IM2COL_F32, im2col_f32, true);
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_UPSCALE_F32, upscale_f32, true); GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_UPSCALE_F32, upscale_f32, true);
@ -3279,8 +3283,6 @@ static void ggml_metal_encode_node(
} break; } break;
case GGML_OP_ROPE: case GGML_OP_ROPE:
{ {
GGML_ASSERT(ne10 == ne02);
const int nth = MIN(1024, ne00); const int nth = MIN(1024, ne00);
const int n_past = ((int32_t *) dst->op_params)[0]; const int n_past = ((int32_t *) dst->op_params)[0];
@ -3303,11 +3305,33 @@ static void ggml_metal_encode_node(
memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float)); memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float));
memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float));
const bool is_neox = mode & 2; const bool is_neox = mode & GGML_ROPE_TYPE_NEOX;
const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; // also true for imrope
const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE;
const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
GGML_ASSERT(!is_vision); // vision mrope not implemented in metal
int sections[4] = {0, 0, 0, 0};
memcpy(sections, (int32_t *) dst->op_params + 11, sizeof(sections));
if (is_mrope) {
// mrope: 4 position ids per token (t/h/w/e blocks)
GGML_ASSERT(ne10 == ne02*4);
GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0);
} else {
GGML_ASSERT(ne10 == ne02);
}
id<MTLComputePipelineState> pipeline = nil; id<MTLComputePipelineState> pipeline = nil;
if (!is_neox) { if (is_mrope) {
switch (src0->type) {
case GGML_TYPE_F32: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F32].pipeline; break;
case GGML_TYPE_F16: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_MULTI_F16].pipeline; break;
default: GGML_ABORT("fatal error");
};
} else if (!is_neox) {
switch (src0->type) { switch (src0->type) {
case GGML_TYPE_F32: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_NORM_F32].pipeline; break; case GGML_TYPE_F32: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_NORM_F32].pipeline; break;
case GGML_TYPE_F16: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16].pipeline; break; case GGML_TYPE_F16: pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16].pipeline; break;
@ -3355,6 +3379,16 @@ static void ggml_metal_encode_node(
[encoder setBytes:&attn_factor length:sizeof( float) atIndex:26]; [encoder setBytes:&attn_factor length:sizeof( float) atIndex:26];
[encoder setBytes:&beta_fast length:sizeof( float) atIndex:27]; [encoder setBytes:&beta_fast length:sizeof( float) atIndex:27];
[encoder setBytes:&beta_slow length:sizeof( float) atIndex:28]; [encoder setBytes:&beta_slow length:sizeof( float) atIndex:28];
if (is_mrope) {
// must match ggml_metal_mrope_args in ggml-metal.metal
struct {
int32_t sect[4];
int32_t is_imrope;
} margs;
memcpy(margs.sect, sections, sizeof(sections));
margs.is_imrope = is_imrope ? 1 : 0;
[encoder setBytes:&margs length:sizeof(margs) atIndex:29];
}
[encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
} break; } break;

131
ggml/src/ggml-metal.metal
View File

@ -2294,8 +2294,132 @@ kernel void kernel_rope_neox(
//} //}
} }
typedef decltype(kernel_rope_norm<float>) kernel_rope_norm_t; // multi-section RoPE (mrope/imrope) — positions are laid out as 4 blocks of
typedef decltype(kernel_rope_neox<float>) kernel_rope_neox_t; // ne02 entries (t/h/w/e), sector of each dim pair selects which block's
// position drives theta; matches ggml_mrope_cache_init on the CPU side
// packed into one struct: Metal caps kernels at 31 buffer slots
typedef struct {
int32_t sect[4];
int32_t is_imrope;
} ggml_metal_mrope_args;
template<typename T>
kernel void kernel_rope_multi(
device const void * src0,
device const int32_t * src1,
device const float * src2,
device float * dst,
constant int64_t & ne00,
constant int64_t & ne01,
constant int64_t & ne02,
constant int64_t & ne03,
constant uint64_t & nb00,
constant uint64_t & nb01,
constant uint64_t & nb02,
constant uint64_t & nb03,
constant int64_t & ne0,
constant int64_t & ne1,
constant int64_t & ne2,
constant int64_t & ne3,
constant uint64_t & nb0,
constant uint64_t & nb1,
constant uint64_t & nb2,
constant uint64_t & nb3,
constant int & n_past,
constant int & n_dims,
constant int & n_ctx_orig,
constant float & freq_base,
constant float & freq_scale,
constant float & ext_factor,
constant float & attn_factor,
constant float & beta_fast,
constant float & beta_slow,
constant ggml_metal_mrope_args & margs,
uint tiitg[[thread_index_in_threadgroup]],
uint3 tptg[[threads_per_threadgroup]],
uint3 tgpig[[threadgroup_position_in_grid]]) {
const int64_t i3 = tgpig[2];
const int64_t i2 = tgpig[1];
const int64_t i1 = tgpig[0];
float corr_dims[2];
rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims);
device const int32_t * pos = src1;
const float theta_base_t = (float) pos[i2 + ne02*0];
const float theta_base_h = (float) pos[i2 + ne02*1];
const float theta_base_w = (float) pos[i2 + ne02*2];
const float theta_base_e = (float) pos[i2 + ne02*3];
const int sect_0 = margs.sect[0];
const int sect_1 = margs.sect[1];
const int sect_2 = margs.sect[2];
const int is_imrope = margs.is_imrope;
const int sect_dims = sect_0 + sect_1 + sect_2 + margs.sect[3];
const int sec_w = sect_1 + sect_0;
const float theta_scale = pow(freq_base, -2.0f/n_dims);
float cos_theta;
float sin_theta;
int64_t i0 = 2*tiitg;
for ( ; i0 < n_dims; i0 += 2*tptg.x) {
const int64_t ic = i0/2;
const int sector = ic % sect_dims;
float theta_base;
if (is_imrope) { // qwen3vl-style interleaved mrope
if (sector % 3 == 1 && sector < 3*sect_1) {
theta_base = theta_base_h;
} else if (sector % 3 == 2 && sector < 3*sect_2) {
theta_base = theta_base_w;
} else if (sector % 3 == 0 && sector < 3*sect_0) {
theta_base = theta_base_t;
} else {
theta_base = theta_base_e;
}
} else {
if (sector < sect_0) {
theta_base = theta_base_t;
} else if (sector < sec_w) {
theta_base = theta_base_h;
} else if (sector < sec_w + sect_2) {
theta_base = theta_base_w;
} else {
theta_base = theta_base_e;
}
}
const float theta = theta_base * pow(theta_scale, (float) ic);
const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)((device char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
device T * dst_data = (device T *)((device char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0);
const float x0 = src[0];
const float x1 = src[n_dims/2];
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[n_dims/2] = x0*sin_theta + x1*cos_theta;
}
for ( ; i0 < ne0; i0 += 2*tptg.x) {
device const T * const src = (device T *)((device char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
device T * dst_data = (device T *)((device char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
dst_data[0] = src[0];
dst_data[1] = src[1];
}
}
typedef decltype(kernel_rope_norm<float>) kernel_rope_norm_t;
typedef decltype(kernel_rope_neox<float>) kernel_rope_neox_t;
typedef decltype(kernel_rope_multi<float>) kernel_rope_multi_t;
template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm<float>; template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm<float>;
template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm<half>; template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm<half>;
@ -2303,6 +2427,9 @@ template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_
template [[host_name("kernel_rope_neox_f32")]] kernel kernel_rope_neox_t kernel_rope_neox<float>; template [[host_name("kernel_rope_neox_f32")]] kernel kernel_rope_neox_t kernel_rope_neox<float>;
template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox<half>; template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox<half>;
template [[host_name("kernel_rope_multi_f32")]] kernel kernel_rope_multi_t kernel_rope_multi<float>;
template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kernel_rope_multi<half>;
typedef void (im2col_t)( typedef void (im2col_t)(
device const float * x, device const float * x,
device char * dst, device char * dst,