From 2b8d0d501195a29fea602e52ae6e9cb194da5a24 Mon Sep 17 00:00:00 2001 From: hchengit <95317477+hchengit@users.noreply.github.com> Date: Fri, 17 Jul 2026 07:49:13 -0700 Subject: [PATCH] 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. --- ggml/src/ggml-metal.m | 42 ++++++++++-- ggml/src/ggml-metal.metal | 131 +++++++++++++++++++++++++++++++++++++- 2 files changed, 167 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-metal.m b/ggml/src/ggml-metal.m index cfba2e5f..323d4e45 100644 --- a/ggml/src/ggml-metal.m +++ b/ggml/src/ggml-metal.m @@ -337,6 +337,8 @@ enum ggml_metal_kernel_type { GGML_METAL_KERNEL_TYPE_ROPE_NORM_F16, GGML_METAL_KERNEL_TYPE_ROPE_NEOX_F32, 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_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_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_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_F32, im2col_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; case GGML_OP_ROPE: { - GGML_ASSERT(ne10 == ne02); - const int nth = MIN(1024, ne00); 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_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 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) { 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; @@ -3355,6 +3379,16 @@ static void ggml_metal_encode_node( [encoder setBytes:&attn_factor length:sizeof( float) atIndex:26]; [encoder setBytes:&beta_fast length:sizeof( float) atIndex:27]; [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)]; } break; diff --git a/ggml/src/ggml-metal.metal b/ggml/src/ggml-metal.metal index f700e6f7..0a3229df 100644 --- a/ggml/src/ggml-metal.metal +++ b/ggml/src/ggml-metal.metal @@ -2294,8 +2294,132 @@ kernel void kernel_rope_neox( //} } -typedef decltype(kernel_rope_norm) kernel_rope_norm_t; -typedef decltype(kernel_rope_neox) kernel_rope_neox_t; +// multi-section RoPE (mrope/imrope) — positions are laid out as 4 blocks of +// 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 +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) kernel_rope_norm_t; +typedef decltype(kernel_rope_neox) kernel_rope_neox_t; +typedef decltype(kernel_rope_multi) kernel_rope_multi_t; template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm; template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm; @@ -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; template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox; +template [[host_name("kernel_rope_multi_f32")]] kernel kernel_rope_multi_t kernel_rope_multi; +template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kernel_rope_multi; + typedef void (im2col_t)( device const float * x, device char * dst,