ggml : add scaling to get_rel_pos for different query/key heights

Author: bluebread

ggml/src/ggml-cpu/ops.cpp

@@ -9217,13 +9217,16 @@ static void ggml_compute_forward_get_rel_pos_f32(
     GGML_TENSOR_UNARY_OP_LOCALS
 
     const int64_t kh = ne1;
+    const int64_t qh = ne2;
+    const float k_scale = MAX(qh / kh, 1.0f);
+    const float q_scale = MAX(kh / qh, 1.0f);
 
     float * src0_data = (float *) src0->data;
     float * dst_data  = (float *) dst->data;
 
     for (int64_t i2 = 0; i2 < ne2; ++i2) {
         for (int64_t i1 = 0; i1 < ne1; ++i1) {
-            const int64_t pos = (kh - i1 - 1) + i2;
+            const int pos = int(i2*q_scale - i1*k_scale + (kh - 1)*k_scale);
             for (int64_t i0 = 0; i0 < ne0; ++i0) {
                 dst_data[i2*ne1*ne0 + i1*ne0 + i0] = src0_data[pos*ne00 + i0];
             }
@@ -9243,13 +9246,16 @@ static void ggml_compute_forward_get_rel_pos_f16(
     GGML_TENSOR_UNARY_OP_LOCALS
 
     const int64_t kh = ne1;
+    const int64_t qh = ne2;
+    const float k_scale = MAX(qh / kh, 1.0f);
+    const float q_scale = MAX(kh / qh, 1.0f);
 
     ggml_fp16_t * src0_data = (ggml_fp16_t *) src0->data;
     ggml_fp16_t * dst_data  = (ggml_fp16_t *) dst->data;
 
     for (int64_t i2 = 0; i2 < ne2; ++i2) {
         for (int64_t i1 = 0; i1 < ne1; ++i1) {
-            const int64_t pos = (kh - i1 - 1) + i2;
+            const int pos = int(i2*q_scale - i1*k_scale + (kh - 1)*k_scale);
             for (int64_t i0 = 0; i0 < ne0; ++i0) {
                 dst_data[i2*ne1*ne0 + i1*ne0 + i0] = src0_data[pos*ne00 + i0];
             }

ggml/src/ggml-cuda/rel-pos.cu

@@ -2,82 +2,16 @@
 #include "ggml.h"
 #include "ggml-cuda/rel-pos.cuh"
 
-/*
-
-static void ggml_compute_forward_get_rel_pos_f16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    GGML_UNUSED(params);
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L292-L322
-
-    GGML_TENSOR_UNARY_OP_LOCALS
-
-    const int64_t kh = ne1;
-
-    ggml_fp16_t * src0_data = (ggml_fp16_t *) src0->data;
-    ggml_fp16_t * dst_data  = (ggml_fp16_t *) dst->data;
-
-    for (int64_t i2 = 0; i2 < ne2; ++i2) {
-        for (int64_t i1 = 0; i1 < ne1; ++i1) {
-            const int64_t pos = (kh - i1 - 1) + i2;
-            for (int64_t i0 = 0; i0 < ne0; ++i0) {
-                dst_data[i2*ne1*ne0 + i1*ne0 + i0] = src0_data[pos*ne00 + i0];
-            }
-        }
-    }
-}
-
-
-void ggml_compute_forward_get_rel_pos(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    switch (src0->type) {
-        case GGML_TYPE_F32:
-            {
-                ggml_compute_forward_get_rel_pos_f32(params, dst);
-            } break;
-        case GGML_TYPE_F16:
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_get_rel_pos_f16(params, dst);
-            } break;
-        default:
-            {
-                GGML_ABORT("fatal error");
-            }
-    }
-}
-
-struct ggml_tensor * ggml_get_rel_pos(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   qh,
-        int                   kh) {
-    GGML_ASSERT(qh + kh - 1 <= a->ne[1]);
-
-    const int64_t ne[4] = { a->ne[0], kh, qh, 1, };
-    struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 3, ne);
-
-    result->op     = GGML_OP_GET_REL_POS;
-    result->src[0] = a;
-
-    return result;
-}
-
-*/
 
 template <typename T>
 __global__ static void get_rel_pos_kernel(const void * src, void * dst, int C) {
     int kh = gridDim.x;
+    int qh = gridDim.x;
+    float k_scale = MAX(qh / kh, 1.0f);
+    float q_scale = MAX(kh / qh, 1.0f);
     int ki = blockIdx.x;
     int qi = blockIdx.y;
-    int pos = (kh - 1) + qi - ki;
+    int pos = int(qi*q_scale - ki*k_scale + (kh - 1)*k_scale);
 
     int s0 = C;
     int s1 = C * kh;

ggml/src/ggml.c

@@ -5479,7 +5479,8 @@ struct ggml_tensor * ggml_get_rel_pos(
         struct ggml_tensor  * a,
         int                   qh,
         int                   kh) {
-    GGML_ASSERT(qh + kh - 1 <= a->ne[1]);
+    GGML_ASSERT(qh >= 1 && kh >= 1);
+    GGML_ASSERT(2*MAX(qh, kh) - 1 == a->ne[1]);
 
     const int64_t ne[4] = { a->ne[0], kh, qh, 1, };
     struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 3, ne);