Method for checking if can dequantize mul mat

Author: JohannesGaessler

ggml-cuda.cu

@@ -618,6 +618,8 @@ static void ggml_cuda_mul_mat_q_f32(const ggml_tensor * src0, const ggml_tensor
     const int nb3  = dst->nb[3];
     const ggml_type type = src0->type;
 
+    const bool can_dmmv = ggml_cuda_can_dequantize_mul_mat_vec(src0, src1, dst);
+
     const float alpha = 1.0f;
     const float beta = 0.0f;
     const int x_ne = ne01 * ne00;
@@ -628,7 +630,7 @@ static void ggml_cuda_mul_mat_q_f32(const ggml_tensor * src0, const ggml_tensor
 
     size_t x_size, y_size, d_size, q_size;
     float * d_X;
-    if (ne11 > 1) {
+    if (!can_dmmv) {
         d_X = (float *) ggml_cuda_pool_malloc(n_mm * sizeof(float) * x_ne, &x_size);
     }
     float * d_Y = (float *) ggml_cuda_pool_malloc(n_mm * sizeof(float) * y_ne, &y_size);
@@ -658,7 +660,7 @@ static void ggml_cuda_mul_mat_q_f32(const ggml_tensor * src0, const ggml_tensor
             } else {
                 GGML_ASSERT(false);
             }
-            if (ne11 == 1) {
+            if (can_dmmv) { // specialized dequantize_mul_mat_vec kernel
                 CUDA_CHECK(cudaEventRecord(cudaEvent, cudaStream2));
 
                 // copy src1 to device
@@ -671,7 +673,7 @@ static void ggml_cuda_mul_mat_q_f32(const ggml_tensor * src0, const ggml_tensor
                 dmmv(c_Q, c_Y, c_D, ne00, ne01, cudaStream);
                 CUDA_CHECK(cudaGetLastError());
 
-            } else {
+            } else { // general matrix matrix multiplication
                 float * c_X = d_X + i * x_ne;
 
                 // convert src0 to fp32 on device
@@ -702,7 +704,7 @@ static void ggml_cuda_mul_mat_q_f32(const ggml_tensor * src0, const ggml_tensor
     }
 
     CUDA_CHECK(cudaDeviceSynchronize());
-    if (ne11 > 1) {
+    if (!can_dmmv) {
         ggml_cuda_pool_free(d_X, x_size);
     }
     ggml_cuda_pool_free(d_Y, y_size);
@@ -720,13 +722,26 @@ bool ggml_cuda_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_te
     if ((src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) &&
         src1->type == GGML_TYPE_F32 &&
         dst->type == GGML_TYPE_F32 &&
-        ((ne0 >= 32 && ne1 >= 32 && ne10 >= 32) || src0->backend == GGML_BACKEND_CUDA)) {
+        (ne0 >= 32 && ne1 >= 32 && ne10 >= 32)) {
         return true;
     }
 
     return false;
 }
 
+bool ggml_cuda_can_dequantize_mul_mat_vec(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst) {
+    const int64_t ne11 = src1->ne[1];
+    const int64_t ne12 = src1->ne[2];
+    const int64_t ne13 = src1->ne[3];
+    if (ggml_is_quantized(src0->type) &&
+        src1->type == GGML_TYPE_F32 &&
+        dst->type == GGML_TYPE_F32 &&
+        (ne11 == 1 && ne12 == 1 && ne13 == 1)) {
+        return true;
+    }
+    return false;
+}
+
 bool ggml_cuda_mul_mat_use_f16(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * /* dst */) {
     size_t src0_sz = ggml_nbytes(src0);
     size_t src1_sz = ggml_nbytes(src1);
@@ -743,7 +758,9 @@ bool ggml_cuda_mul_mat_use_f16(const struct ggml_tensor * src0, const struct ggm
 }
 
 void ggml_cuda_mul_mat(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, void * wdata, size_t wsize) {
-    GGML_ASSERT(ggml_cuda_can_mul_mat(src0, src1, dst));
+    GGML_ASSERT(ggml_cuda_can_mul_mat(src0, src1, dst) ||
+                ggml_cuda_can_dequantize_mul_mat_vec(src0, src1, dst) ||
+                src0->backend == GGML_BACKEND_CUDA);
 
     if (src0->type == GGML_TYPE_F32) {
         ggml_cuda_mul_mat_f32(src0, src1, dst);

ggml-cuda.h

@@ -7,6 +7,7 @@ extern "C" {
 void   ggml_init_cublas(void);
 
 bool   ggml_cuda_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst);
+bool   ggml_cuda_can_dequantize_mul_mat_vec(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst);
 size_t ggml_cuda_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst);
 void   ggml_cuda_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst, void * wdata, size_t wsize);
 

ggml.c

@@ -7799,7 +7799,9 @@ static void ggml_compute_forward_mul_mat_q_f32(
     //   compute by src0 rows
 
 #if defined(GGML_USE_CUBLAS)
-    if (ggml_cuda_can_mul_mat(src0, src1, dst)) {
+    if (ggml_cuda_can_mul_mat(src0, src1, dst) ||
+        ggml_cuda_can_dequantize_mul_mat_vec(src0, src1, dst) ||
+        src0->backend == GGML_BACKEND_CUDA) {
         if (params->ith == 0 && params->type == GGML_TASK_COMPUTE) {
             ggml_cuda_mul_mat(src0, src1, dst, params->wdata, params->wsize);
         }