create helper function instead of re-using params

Author: am17an

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -3211,17 +3211,10 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
                     }
 
                     if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) {
-                        ggml_tensor * src3 = cgraph->nodes[i + 2]->src[1];
                         ggml_tensor * rope = cgraph->nodes[i];
-                        ggml_tensor * dst  = cgraph->nodes[i + 2];
+                        ggml_tensor * set_rows = cgraph->nodes[i + 2];
 
-                        rope->src[3] = src3;
-                        rope->data   = dst->data;
-                        rope->type   = dst->type;
-
-                        const size_t set_rows_stride = dst->nb[1] / ggml_type_size(dst->type);
-                        ggml_set_op_params_i32(rope, 15, set_rows_stride);
-                        ggml_cuda_op_rope(*cuda_ctx, rope);
+                        ggml_cuda_op_rope_fused(*cuda_ctx, rope, set_rows);
                         i += 2;
                         continue;
                     }

ggml/src/ggml-cuda/rope.cu

@@ -1,3 +1,5 @@
+#include "ggml-cuda/common.cuh"
+#include "ggml.h"
 #include "rope.cuh"
 
 struct rope_corr_dims {
@@ -399,16 +401,28 @@ static void rope_vision_cuda(
 }
 
 template <bool forward>
-void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
+                            ggml_tensor *               dst,
+                            const ggml_tensor *         set_rows = nullptr) {
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
     const ggml_tensor * src2 = dst->src[2];
-    const ggml_tensor * src3 = dst->src[3];
 
     const float * src0_d = (const float *)src0->data;
     const float * src1_d = (const float *)src1->data;
 
-    float * dst_d = (float *)dst->data;
+    void *          dst_d           = dst->data;
+    const int64_t * row_indices     = nullptr;
+    ggml_type       dst_type        = dst->type;
+    int             set_rows_stride = 0;
+
+    if (set_rows != nullptr) {
+        GGML_ASSERT(forward);
+        dst_d           = set_rows->data;
+        row_indices     = (const int64_t *) set_rows->src[1]->data;
+        dst_type        = set_rows->type;
+        set_rows_stride = set_rows->nb[1] / ggml_type_size(set_rows->type);
+    }
     cudaStream_t stream = ctx.stream();
 
     GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
@@ -468,29 +482,20 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
         freq_factors = (const float *) src2->data;
     }
 
-    // Row indices for fused ROPE + VIEW + SET_ROWS
-    const int64_t * row_indices     = nullptr;
-    int             set_rows_stride = 0;
-    if (src3 != nullptr) {
-        GGML_ASSERT(src3->type == GGML_TYPE_I64);
-        row_indices     = (const int64_t *) src3->data;
-        set_rows_stride = ggml_get_op_params_i32(dst, 15);
-    }
-
     rope_corr_dims corr_dims;
     ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v);
 
     // compute
     if (is_neox) {
-        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
             rope_neox_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims,
                                                   nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                   freq_factors, row_indices, set_rows_stride, stream);
-        } else if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
+        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
             rope_neox_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims,
                                                  nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                  freq_factors, row_indices, set_rows_stride, stream);
-        } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
+        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
             rope_neox_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr,
                                                 pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                 freq_factors, row_indices, set_rows_stride, stream);
@@ -522,15 +527,15 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
             GGML_ABORT("fatal error");
         }
     } else {
-        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
             rope_norm_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims,
                                                   nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                   freq_factors, row_indices, set_rows_stride, stream);
-        } else if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
+        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
             rope_norm_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims,
                                                  nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                  freq_factors, row_indices, set_rows_stride, stream);
-        } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
+        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
             rope_norm_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr,
                                                 pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
                                                 freq_factors, row_indices, set_rows_stride, stream);
@@ -547,3 +552,7 @@ void ggml_cuda_op_rope(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     ggml_cuda_op_rope_impl<false>(ctx, dst);
 }
+
+void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope, ggml_tensor * set_rows) {
+    ggml_cuda_op_rope_impl<true>(ctx, rope, set_rows);
+}

ggml/src/ggml-cuda/rope.cuh

@@ -5,3 +5,5 @@
 void ggml_cuda_op_rope(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
 void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows);