Merge branch 'ikawrakow:main' into main

Author: Thireus

common/common.cpp

@@ -1020,6 +1020,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
         params.fused_up_gate = false;
         return true;
     }
+    if (arg == "-no-mmad" || arg == "--no-fused-mul-multiadd") {
+        params.fused_mmad = false;
+        return true;
+    }
     if (arg == "-ser" || arg == "--smart-expert-reduction") {
         CHECK_ARG
         auto values = string_split_pairs<int,float>(argv[i], ',');
@@ -1806,6 +1810,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
     options.push_back({ "*",           "-fmoe, --fused-moe",            "enable fused MoE (default: %s)", params.fused_moe_up_gate ? "enabled" : "disabled" });
     options.push_back({ "*",           "-ger,  --grouped-expert-routing", "enable grouped expert routing (default: %s)", params.grouped_expert_routing ? "enabled" : "disabled" });
     options.push_back({ "*",           "-no-fug, --no-fused-up-gate",   "disaable fused up-gate (default: %s)", params.fused_up_gate ? "enabled" : "disabled" });
+    options.push_back({ "*",           "-no-mmad, --no-fused-mul-multiadd", "disaable fused mul-multi_add (default: %s)", params.fused_mmad? "enabled" : "disabled" });
     options.push_back({ "*",         "-ser,  --smart-expert-reduction,","experts reduction (default: %d,%g)", params.min_experts, params.thresh_experts});
     options.push_back({ "*",           "-p,    --prompt PROMPT",        "prompt to start generation with\n"
                                                                         "in conversation mode, this will be used as system prompt\n"
@@ -2762,6 +2767,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     cparams.fused_moe_up_gate = params.fused_moe_up_gate;
     cparams.grouped_expert_routing = params.grouped_expert_routing;
     cparams.fused_up_gate     = params.fused_up_gate;
+    cparams.fused_mmad        = params.fused_mmad;
     cparams.min_experts       = params.min_experts;
     cparams.thresh_experts    = params.thresh_experts;
     cparams.only_active_experts = params.only_active_exps;
@@ -3879,6 +3885,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l
     fprintf(stream, "fused_moe: %s # default: false\n", params.fused_moe_up_gate ? "true" : "false");
     fprintf(stream, "grouped_expert_routing: %s # default: false\n", params.grouped_expert_routing ? "true" : "false");
     fprintf(stream, "fused_up_gate: %s # default: true\n", params.fused_up_gate ? "true" : "false");
+    fprintf(stream, "fused_mmad: %s # default: true\n", params.fused_mmad? "true" : "false");
     fprintf(stream, "ser: %d,%g # defaulr: -1,0\n", params.min_experts, params.thresh_experts);
     fprintf(stream, "temp: %f # default: 0.8\n", sparams.temp);
 

common/common.h

@@ -235,6 +235,7 @@ struct gpt_params {
     int  attn_max_batch    = 0;     // Max batch size to use when computing attention (only applicable if flash_attn = false)
     bool fused_moe_up_gate = false; // fused up*unary(gate) op for MoE models
     bool fused_up_gate     = true;  // fused up*unary(gate) op
+    bool fused_mmad        = true;  // fused mul+multi_add op
     bool grouped_expert_routing = false; // if to use grouped expert routing (BailingMoeV2 arch)
     int  min_experts       = -1;
     float thresh_experts   = 0;

ggml/include/ggml.h

@@ -619,6 +619,7 @@ extern "C" {
         GGML_OP_OUT_PROD,
         GGML_OP_FUSED_UP_GATE,
         GGML_OP_MOE_FUSED_UP_GATE,
+        GGML_OP_MUL_MULTI_ADD,
 
         GGML_OP_SCALE,
         GGML_OP_SET,
@@ -1083,6 +1084,11 @@ extern "C" {
             struct ggml_tensor  * a,
             int n_experts);
 
+    GGML_API struct ggml_tensor * ggml_mul_multi_add(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b);
+
     // dst = a
     // view(dst, nb1, nb2, nb3, offset) += b
     // return dst

ggml/src/ggml-cuda.cu

@@ -46,6 +46,7 @@
 #include "ggml-cuda/conv2d-dw.cuh"
 #include "ggml-cuda/set-rows.cuh"
 #include "ggml-cuda/argmax.cuh"
+#include "ggml-cuda/multiadd.cuh"
 
 #include <algorithm>
 #include <array>
@@ -3178,6 +3179,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_MULTI_ADD:
             ggml_cuda_op_multi_add(ctx, dst);
             break;
+        case GGML_OP_MUL_MULTI_ADD:
+            ggml_cuda_op_mul_multi_add(ctx, dst);
+            break;
         case GGML_OP_ACC:
             ggml_cuda_op_acc(ctx, dst);
             break;
@@ -4408,6 +4412,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
         case GGML_OP_ADD:
         case GGML_OP_ADD_ID:
         case GGML_OP_MULTI_ADD:
+        case GGML_OP_MUL_MULTI_ADD:
         case GGML_OP_MUL:
         case GGML_OP_DIV:
         case GGML_OP_FUSED_RMS_NORM:

ggml/src/ggml-cuda/multiadd.cu

@@ -0,0 +1,87 @@
+#include "multiadd.cuh"
+
+static __global__ void multi_add_f32(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, const char * src0, char * dst) {
+    const int64_t i = blockDim.x*blockIdx.x + threadIdx.x;
+    int64_t k = ne0*ne1;
+    if (i >= k) {
+        return;
+    }
+    int i1 = i / ne0;
+    int i0 = i % ne0;
+    float * result = (float *)(dst + i1*nb1);
+    const float * s = (const float *)(src0 + i1*nb01) + i0;
+    if (nused == 1) {
+        result[i0] = s[0];
+    } else {
+        float sum = s[0] + s[ne0];
+        for (int j = 2; j < nused; ++j) sum += s[j*ne0];
+        result[i0] = sum;
+    }
+}
+
+static void multi_add_f32_cuda(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, const char * src0, char * dst, cudaStream_t stream) {
+    int64_t k = ne0 * ne1;
+    const int num_blocks = (k + CUDA_MULTI_ADD_BLOCK_SIZE - 1) / CUDA_MULTI_ADD_BLOCK_SIZE;
+    multi_add_f32<<<num_blocks, CUDA_MULTI_ADD_BLOCK_SIZE, 0, stream>>>(nused, ne0, ne1, nb1, nb01, src0, dst);
+}
+
+void ggml_cuda_op_multi_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->ne[2] == 1 && dst->ne[3] == 1);
+    GGML_ASSERT(dst->nb[0] == sizeof(float));
+    int nused = dst->op_params[0];
+    GGML_ASSERT(nused >= 1);
+    const char * src0 = (const char *)dst->src[0]->data;
+    cudaStream_t stream = ctx.stream();
+    multi_add_f32_cuda(nused, dst->ne[0], dst->ne[1], dst->nb[1], dst->src[0]->nb[1], src0, (char *)dst->data, stream);
+}
+
+
+static __global__ void mul_multi_add_f32(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, int64_t nb02, int64_t nb11, int64_t nb12, const char * src0, const char * src1, char * dst) {
+    const int64_t i = blockDim.x*blockIdx.x + threadIdx.x;
+    int64_t k = ne0*ne1;
+    if (i >= k) {
+        return;
+    }
+    int i1 = i / ne0;
+    int i0 = i % ne0;
+    float * result = (float *)(dst + i1*nb1);
+
+    auto c0 = src0 + i1*nb02;
+    auto c1 = src1 + i1*nb12;
+
+    float sum = 0;
+    for (int j = 0; j < nused; ++j) {
+        auto x0 = (const float *)c0;
+        auto x1 = (const float *)c1;
+        sum += x0[i0] * x1[0];
+        c0 += nb01;
+        c1 += nb11;
+    }
+    result[i0] = sum;
+}
+
+static void mul_multi_add_f32_cuda(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, int64_t nb02, int64_t nb11, int64_t nb12,
+        const char * src0, const char * src1, char * dst, cudaStream_t stream) {
+    int64_t k = ne0 * ne1;
+    const int num_blocks = (k + CUDA_MULTI_ADD_BLOCK_SIZE - 1) / CUDA_MULTI_ADD_BLOCK_SIZE;
+    mul_multi_add_f32<<<num_blocks, CUDA_MULTI_ADD_BLOCK_SIZE, 0, stream>>>(nused, ne0, ne1, nb1, nb01, nb02, nb11, nb12, src0, src1, dst);
+}
+
+void ggml_cuda_op_mul_multi_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    auto src0 = dst->src[0];
+    auto src1 = dst->src[1];
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->ne[0] ==  dst->ne[0]);
+    GGML_ASSERT(src0->ne[2] ==  dst->ne[1]);
+    GGML_ASSERT(src0->ne[1] == src1->ne[1]);
+    GGML_ASSERT(src0->ne[2] == src1->ne[2]);
+    GGML_ASSERT(src0->ne[3] == src1->ne[3]);
+    GGML_ASSERT(src0->ne[3] == 1);
+    GGML_ASSERT(src1->ne[0] == 1);
+
+    mul_multi_add_f32_cuda(src0->ne[1], dst->ne[0], dst->ne[1], dst->nb[1], src0->nb[1], src0->nb[2], src1->nb[1], src1->nb[2],
+            (const char *)src0->data, (const char *)src1->data, (char *)dst->data, ctx.stream());
+}

ggml/src/ggml-cuda/multiadd.cuh

@@ -0,0 +1,14 @@
+//
+// Copyright (C) 2023-2024 The ggml authors
+// Copyright (C) 2024 Iwan Kawrakow
+// MIT license
+// SPDX-License-Identifier: MIT
+//
+
+#include "common.cuh"
+
+#define CUDA_MULTI_ADD_BLOCK_SIZE 256
+
+void ggml_cuda_op_multi_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_mul_multi_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-cuda/unary.cu

@@ -59,25 +59,6 @@ static __global__ void fused_mul_silu_f32(const float * x, const float * y, floa
     dst[i] = x[i] * y[i] / (1.0f + expf(-x[i]));
 }
 
-static __global__ void multi_add_f32(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, const char * src0, char * dst) {
-    const int64_t i = blockDim.x*blockIdx.x + threadIdx.x;
-    int64_t k = ne0*ne1;
-    if (i >= k) {
-        return;
-    }
-    int i1 = i / ne0;
-    int i0 = i % ne0;
-    float * result = (float *)(dst + i1*nb1);
-    const float * s = (const float *)(src0 + i1*nb01) + i0;
-    if (nused == 1) {
-        result[i0] = s[0];
-    } else {
-        float sum = s[0] + s[ne0];
-        for (int j = 2; j < nused; ++j) sum += s[j*ne0];
-        result[i0] = sum;
-    }
-}
-
 static __global__ void fused_mul_relu_f32(const float * x, const float * y, float * dst, const int k) {
     const int i = blockDim.x*blockIdx.x + threadIdx.x;
 
@@ -261,23 +242,6 @@ static void sqrt_f32_cuda(const float * x, float * dst, const int k, cudaStream_
     sqrt_f32<<<num_blocks, CUDA_SQRT_BLOCK_SIZE, 0, stream>>>(x, dst, k);
 }
 
-static void multi_add_f32_cuda(int nused, int64_t ne0, int64_t ne1, int64_t nb1, int64_t nb01, const char * src0, char * dst, cudaStream_t stream) {
-    int64_t k = ne0 * ne1;
-    const int num_blocks = (k + CUDA_MULTI_ADD_BLOCK_SIZE - 1) / CUDA_MULTI_ADD_BLOCK_SIZE;
-    multi_add_f32<<<num_blocks, CUDA_MULTI_ADD_BLOCK_SIZE, 0, stream>>>(nused, ne0, ne1, nb1, nb01, src0, dst);
-}
-
-void ggml_cuda_op_multi_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->ne[2] == 1 && dst->ne[3] == 1);
-    GGML_ASSERT(dst->nb[0] == sizeof(float));
-    int nused = dst->op_params[0];
-    GGML_ASSERT(nused >= 1);
-    const char * src0 = (const char *)dst->src[0]->data;
-    cudaStream_t stream = ctx.stream();
-    multi_add_f32_cuda(nused, dst->ne[0], dst->ne[1], dst->nb[1], dst->src[0]->nb[1], src0, (char *)dst->data, stream);
-}
-
 void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
     const float * src0_d = (const float *)src0->data;