Proposal: ggml_rope_v2

[ngxson]

Motivation

While working with Devstral 2 and more recently, GLM4V, I notice that both the RoPE kernels and the API ggml_rope_ext is quite difficult to work with. I also spotted some opportunities to improve a bit the performance, so I'm writing this proposal

Problems with the existing kernel:

• It does some redundant calculations repeatedly, for example: theta_scale, or the mscale calculation in yarn code path. They can be calculated once on the host device then pass the result into the kernel
• M-RoPE code paths is difficult to track (it got worse after interleaved M-RoPE was added along side with Qwen3-VL)

Problems with the API:

• ggml_rope_ext has too many arguments
• ggml_rope_multi has even more arguments and half of them are hardly used in practice
• mode was supposed to be a bit field, but got messy with M-RoPE
• attn_factor is mis-aligned with transformers which made the recent Devstral 2 debugging quite tricky

Proposal

My proposal contains 2 parts: The API and the kernel

For the API, I propose an array of calls like this:

// rope Qcur with pos
struct ggml_tensor * roped = ggml_rope_v2(ctx, Qcur, pos, n_dims, freq_base);
// with yarn
ggml_rope_v2_set_yarn(ctx, roped, ext_factor, mscale, beta_fast, beta_slow);
// with m-rope
ggml_rope_v2_set_mrope(ctx, roped, sections);
// with ordering (NEOX or NORMAL)
ggml_rope_v2_set_ordering(ctx, roped, GGML_ROPE_ORDERING_NEOX);

Behind the scene, these call will pre-calculate as many things as possible before passing args to the kernel. For example theta_scale = powf(freq_base, -2.0f/n_dims) will be calculated inside ggml_rope_v2's implementation

For the kernel, we can statically compile templated kernels with the combination of:

• 2 x direction: forward or backward
• 2 x types: f16, f32
• 2 x ordering: normal, neox --> the indexing will be controlled via an input arg
• 2 x modes: normal, M-RoPE

Note: mode vision (aka 2D-RoPE with Neox) will become a composed op of 2 x ggml_rope_v2 instead

So in total, we will have 8 statically compiled kernels

[ggerganov]

The proposed v2 API seems OK.

Regarding the kernels - do I understand correctly that this is independent from the API change? Basically, you propose that we progressively templetify the implementations in the different backends, even with the existing ggml_rope/ggml_rope_ext API.
If yes, I think this is also OK. We have to validate each change very carefully. We can start with the Metal backend?

[ngxson]

Regarding the kernels - do I understand correctly that this is independent from the API change?

Yes, these 2 changes are independent one from another. But for testing purposes, I think we can have ggml_rope_v2 API+kernel on Metal as you suggested, to serve as a demo.

The short-term goal is to have ggml_rope_v2 and an array of minimalist kernels.

The, the long-term goal will be to completely remove the kernels using by ggml_rope_ext and replace this API with composed calls to ggml_rope_v2

[ngxson]

Just to clarify, I think there won't be a public API called ggml_rope_v2 very soon, but my goal is to firstly make it as an internal API such that the current ggml_rope_ext can be built on top

In later stage, we will make any adjustments needed then make it a public API. I'm even a bit doubted about the naming ggml_rope_v2, probably something like ggml_rope_comp (comp = composable) will be more suitable - but we can decide this later

[ngxson]

I made a very minimal demo with CPU-only here: https://github.com/ngxson/llama.cpp/pull/59/changes

• Yarn logic is simplified, no more mscale calculation, high/low factors are calculated outside of the kernel
• Ordering is translate into index scaling and offset (probably can be improved further?)

Thinking along the line, I'm wondering if we should take this opportunity to allow float as position input. I remember there was a discussion about this a while ago but probably there is no use case for it. But still, the kernel internally converts it to f32, so probably accepting f32 input can be a tiny improvement

[ggerganov]

Looks good. The changes would still need to be done in llama.cpp as we have more test coverage there.

[ngxson]

FYI, I added these into the mentioned PR:

• use F32 for input pos
• metal kernel
• add a series of test in test-rope to make sure the output of rope_comp is the same as rope_ext

I still need to work on m-rope before publishing the PR to upstream llama.cpp. But in the meantime, feel free to have a look

[ngxson]

Maybe a bit related to ggml-org/llama.cpp#18191 , but I think the m-rope mode should accept input 2D tensor [n_tokens, 4] instead of 1D [n_tokens * 4]. Not sure if doing having any side effects, but I'll give a try.