Make Q8_0 KV cache work with mla=2,fa on CUDA (#264)

Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>

1 files changed, 19 insertions, 5 deletions

diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu
index b40617f6..f89ce974 100644
--- a/ggml/src/ggml-cuda/concat.cu
+++ b/ggml/src/ggml-cuda/concat.cu
@@ -209,6 +209,10 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 
     if (dim == 0 && src0->nb[0] == ggml_type_size(src0->type) && src1->nb[0] == ggml_type_size(src1->type) &&
             src0->nb[1] % sizeof(float) == 0 && src1->nb[1] % sizeof(float) == 0) {
+        auto bs = ggml_blck_size(dst->type);
+        auto ts = ggml_type_size(dst->type);
+        auto ne00_eff = (src0->ne[0]/bs)*ts/sizeof(float);
+        auto ne0_eff  = (dst->ne[0]/bs)*ts/sizeof(float);
         if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
             //if (dst->ne[1] >= 65536 || dst->ne[2] >= 65536) {
             //    fprintf(stderr, "%s: ne1 = %ld, ne2 = %ld exceed max. blocks when computing %s\n", __func__, dst->ne[1], dst->ne[2], dst->name);
@@ -217,25 +221,35 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
             const float * src0_d = (const float *)src0->data;
             const float * src1_d = (const float *)src1->data;
             float * dst_d = (float *)dst->data;
+            //printf("%s(%s, %s): %ld %zu %zu  %ld %zu %zu\n", __func__, src0->name, src1->name, src0->ne[0], src0->nb[0], src0->nb[1], dst->ne[0], dst->nb[0], dst->nb[1]);
             for (int i3 = 0; i3 < dst->ne[3]; i3++) {
                 concat_f32_cuda(
                         src0_d + i3 * (src0->nb[3] / 4),
                         src1_d + i3 * (src1->nb[3] / 4),
                         dst_d + i3 * ( dst->nb[3] / 4),
-                        src0->ne[0]*src0->nb[0]/sizeof(float), src0->ne[1], src0->ne[2],
-                        dst->ne[0]*dst->nb[0]/sizeof(float),  dst->ne[1],  dst->ne[2], dim, stream);
+                        ne00_eff, src0->ne[1], src0->ne[2],
+                        ne0_eff, dst->ne[1], dst->ne[2], dim, stream);
+                        //src0->nb[1]/sizeof(float), src0->ne[1], src0->ne[2],
+                        //dst->nb[1]/sizeof(float), dst->ne[1], dst->ne[2], dim, stream);
+                        //src0->ne[0]*src0->nb[0]/sizeof(float), src0->ne[1], src0->ne[2],
+                        //dst->ne[0]*dst->nb[0]/sizeof(float),  dst->ne[1],  dst->ne[2], dim, stream);
             }
         } else {
+            //printf("%s(not contiguous): %s(%s) and %s(%s)\n", __func__, src0->name, ggml_type_name(src0->type), src1->name, ggml_type_name(src1->type));
+            auto ne10_eff = (src1->ne[0]/bs)*ts/sizeof(float);
             dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]);
             concat_f32_non_cont<<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
                     (const char *)src0->data,
                     (const char *)src1->data,
                     (      char *)dst->data,
-                    src0->ne[0]*src0->nb[0]/sizeof(float), src0->ne[1], src0->ne[2], src0->ne[3],
+                    ne00_eff, src0->ne[1], src0->ne[2], src0->ne[3],
+                    //src0->ne[0]*src0->nb[0]/sizeof(float), src0->ne[1], src0->ne[2], src0->ne[3],
                     sizeof(float), src0->nb[1], src0->nb[2], src0->nb[3],
-                    src1->ne[0]*src1->nb[0]/sizeof(float), src1->ne[1], src1->ne[2], src1->ne[3],
+                    ne10_eff, src1->ne[1], src1->ne[2], src1->ne[3],
+                    //src1->ne[0]*src1->nb[0]/sizeof(float), src1->ne[1], src1->ne[2], src1->ne[3],
                     sizeof(float), src1->nb[1], src1->nb[2], src1->nb[3],
-                    dst->ne[0]*dst->nb[0]/sizeof(float),  dst->ne[1],  dst->ne[2],  dst->ne[3],
+                    ne0_eff,  dst->ne[1],  dst->ne[2],  dst->ne[3],
+                    //dst->ne[0]*dst->nb[0]/sizeof(float),  dst->ne[1],  dst->ne[2],  dst->ne[3],
                     sizeof(float),  dst->nb[1],  dst->nb[2],  dst->nb[3], dim);
         }
         return;