Commit a55c16e4e for llama.cpp
commit a55c16e4ee532aa36d55f5756f3ba28a21a2b62b
Author: Aman Gupta <amangupta052@gmail.com>
Date: Sun Oct 25 17:15:45 2026 +0800
ggml-cuda: fuse qwen MRoPE + set-rows (#31025)
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 454d17d78..b6efc8736 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -2881,9 +2881,10 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
return false;
}
- // Only norm/neox shaders have the fusion code
+ // Only norm/neox/multi shaders have the fusion code
const int mode = ((const int32_t *) rope->op_params)[2];
- if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
+ if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX &&
+ mode != GGML_ROPE_TYPE_MROPE && mode != GGML_ROPE_TYPE_IMROPE) {
return false;
}
diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu
index b4fca778f..45e813406 100644
--- a/ggml/src/ggml-cuda/rope.cu
+++ b/ggml/src/ggml-cuda/rope.cu
@@ -196,9 +196,9 @@ static __global__ void rope_neox(const T * x,
dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast<D>(x0 * sin_theta + x1 * cos_theta);
}
-template <bool forward, bool has_ff, typename T>
+template <bool forward, bool has_ff, typename T, typename D>
static __global__ void rope_multi(const T * x,
- T * dst,
+ D * dst,
const int ne00,
const int ne01,
const int ne02,
@@ -219,6 +219,8 @@ static __global__ void rope_multi(const T * x,
const float * freq_factors,
const mrope_sections sections,
const bool is_imrope,
+ const int64_t * row_indices,
+ const int set_rows_stride,
const bool inplace) {
const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y);
@@ -236,12 +238,20 @@ static __global__ void rope_multi(const T * x,
const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
ggml_cuda_pdl_sync();
+
+ // Fusion optimization: ROPE + VIEW + SET_ROWS.
+ // The rope output is viewed as a 1D tensor and offset based on a row index in row_indices.
+ if (set_rows_stride != 0) {
+ idst = i1 * s1 + i0 / 2;
+ idst += row_indices[i2] * set_rows_stride;
+ }
+
if (i0 < n_offs || i0 >= n_offs + n_dims) {
if (inplace) {
return;
}
- dst[idst + i0/2 + 0] = x[ix + i0/2 + 0];
- dst[idst + i0/2 + 1] = x[ix + i0/2 + 1];
+ dst[idst + i0/2 + 0] = ggml_cuda_cast<D>(x[ix + i0/2 + 0]);
+ dst[idst + i0/2 + 1] = ggml_cuda_cast<D>(x[ix + i0/2 + 1]);
return;
}
@@ -286,8 +296,8 @@ static __global__ void rope_multi(const T * x,
const float x0 = x[ix + n_offs/2 + 0];
const float x1 = x[ix + n_offs/2 + n_dims/2];
- dst[idst + n_offs/2 + 0] = x0*cos_theta - x1*sin_theta;
- dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta;
+ dst[idst + n_offs/2 + 0] = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
+ dst[idst + n_offs/2 + n_dims/2] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
}
template <bool forward, bool has_ff, typename T>
@@ -443,9 +453,9 @@ static void rope_neox_cuda(const T * x,
}
}
-template <bool forward, typename T>
+template <bool forward, typename T, typename D>
static void rope_multi_cuda(const T * x,
- T * dst,
+ D * dst,
const int ne00,
const int ne01,
const int ne02,
@@ -467,6 +477,8 @@ static void rope_multi_cuda(const T * x,
const float * freq_factors,
const mrope_sections sections,
const bool is_imrope,
+ const int64_t * row_indices,
+ const int set_rows_stride,
const bool inplace,
cudaStream_t stream) {
GGML_ASSERT(ne00 % 2 == 0);
@@ -478,14 +490,14 @@ static void rope_multi_cuda(const T * x,
if (freq_factors == nullptr) {
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
- ggml_cuda_kernel_launch(rope_multi<forward, false, T>, launch_params,
+ ggml_cuda_kernel_launch(rope_multi<forward, false, T, D>, launch_params,
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
- attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+ attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
} else {
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
- ggml_cuda_kernel_launch(rope_multi<forward, true, T>, launch_params,
+ ggml_cuda_kernel_launch(rope_multi<forward, true, T, D>, launch_params,
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
- attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+ attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
}
}
@@ -648,14 +660,21 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
GGML_ABORT("fatal error");
}
} else if (is_mrope && !is_vision) {
- if (src0->type == GGML_TYPE_F32) {
- rope_multi_cuda<forward>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
- s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
- corr_dims, freq_factors, sections, is_imrope, inplace, stream);
- } else if (src0->type == GGML_TYPE_F16) {
- rope_multi_cuda<forward>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
- s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
- corr_dims, freq_factors, sections, is_imrope, inplace, stream);
+ if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
+ rope_multi_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02,
+ s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+ ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+ row_indices, set_rows_stride, inplace, stream);
+ } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
+ rope_multi_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+ s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+ ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+ row_indices, set_rows_stride, inplace, stream);
+ } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
+ rope_multi_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+ s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+ ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+ row_indices, set_rows_stride, inplace, stream);
} else {
GGML_ABORT("fatal error");
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 1eee812e6..cb68651dc 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -2798,11 +2798,11 @@ struct test_rope_set_rows : public test_case {
const ggml_type type_idx;
const std::array<int64_t, 4> ne_a;
int mode;
+ const int n_dims;
const int n_ctx{512};
- const int n_dims{128};
std::string vars() override {
- return VARS_TO_STR4(type, type_idx, ne_a, mode);
+ return VARS_TO_STR5(type, type_idx, ne_a, mode, n_dims);
}
std::string op_desc(ggml_tensor * t) override {
@@ -2815,8 +2815,9 @@ struct test_rope_set_rows : public test_case {
test_rope_set_rows(ggml_type type,
ggml_type type_idx,
std::array<int64_t, 4> ne_a,
- int mode)
- : type(type), type_idx(type_idx), ne_a(ne_a), mode(mode) {}
+ int mode,
+ int n_dims = 128)
+ : type(type), type_idx(type_idx), ne_a(ne_a), mode(mode), n_dims(n_dims) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne_a[0], ne_a[1], ne_a[2], 1);
@@ -2850,7 +2851,7 @@ struct test_rope_set_rows : public test_case {
rope = ggml_rope_multi(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f);
}
} else {
- rope = ggml_rope(ctx, a, pos, ne_a[0], mode);
+ rope = ggml_rope(ctx, a, pos, n_dims, mode);
}
ggml_tensor * view = ggml_view_2d(ctx, rope, ne_a[0] * ne_a[1], ne_a[2], rope->nb[2], 0);
@@ -10233,6 +10234,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
test_cases.emplace_back(new test_rope_set_rows(GGML_TYPE_F32, GGML_TYPE_I32, { 128, 32, 8, 1 }, GGML_ROPE_TYPE_IMROPE));
+ // partial rotation
+ for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_MROPE, GGML_ROPE_TYPE_IMROPE }) {
+ for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+ test_cases.emplace_back(new test_rope_set_rows(type, GGML_TYPE_I64, { 256, 4, 8, 1 }, mode, 64));
+ }
+ }
for (ggml_type type_input : {GGML_TYPE_F32}) {
for (ggml_op_pool pool_type : {GGML_OP_POOL_AVG, GGML_OP_POOL_MAX}) {