q5_K works

Author: JohannesGaessler

ggml-cuda.cu

@@ -186,6 +186,8 @@ typedef struct {
 static_assert(sizeof(block_q4_K) == 2*sizeof(ggml_fp16_t) + 3*QK_K/64 + QK_K/2, "wrong q4_K block size/padding");
 #endif
 
+#define QR5_K 2
+#define QI5_K (QK_K / (4*QR5_K))
 #ifdef GGML_QKK_64
 typedef struct {
     half d;                  // super-block scale
@@ -1531,6 +1533,50 @@ static __device__ __forceinline__ float vec_dot_q4_K_q8_1(
 #endif // __CUDA_ARCH__ >= 610
 }
 
+static __device__ __forceinline__ float vec_dot_q5_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int iqs) {
+
+#if __CUDA_ARCH__ >= 610 // lowest compute capability for integer intrinsics
+    const block_q5_K * bq5_K = (const block_q5_K *) vbq;
+
+    const int bq8_offset = 2 * (iqs / 8);
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0;
+
+    const float d = bq5_K->d;
+    const float dmin = bq5_K->dmin;
+
+    const int vil = *((int *) &bq5_K->qs[sizeof(int) * iqs]);
+
+    const int vih = (*((int *) &bq5_K->qh[sizeof(int) * (iqs % 8)])) >> bq8_offset;
+
+    for (int i = 0; i < 2; ++i) {
+        const int isc = bq8_offset + i;
+
+        uint8_t sc, m;
+        get_scale_min_k4(isc, bq5_K->scales, sc, m);
+
+        const block_q8_1 * bq8i = bq8_1 + bq8_offset + i;
+        const int uii = *((int*) &bq8i->qs[sizeof(int) * (iqs%8)]);
+        const float d8i = bq8i->d;
+
+        const int viil = (vil >> (4*i)) & 0x0F0F0F0F;
+
+        const int viih = ((vih >> i) << 4) & 0x10101010;
+
+        const int vii = viil | viih;
+
+        sumf_d += d8i * (__dp4a(vii,        uii, 0) * sc);
+        sumf_m += d8i * (__dp4a(0x01010101, uii, 0) * m);
+    }
+
+    return d*sumf_d - dmin*sumf_m;
+#else
+    return 0.0f; // only to satisfy the compiler
+#endif // __CUDA_ARCH__ >= 610
+}
+
 template <int qk, int qi, typename block_q_t, vec_dot_q_cuda_t vec_dot_q_cuda>
 static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, const int ncols, const int nrows) {
     const int row = blockIdx.y*blockDim.y + threadIdx.y;
@@ -2116,6 +2162,15 @@ static void mul_mat_vec_q4_K_q8_1_cuda(const void * vx, const void * vy, float *
         <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
 }
 
+static void mul_mat_vec_q5_K_q8_1_cuda(const void * vx, const void * vy, float * dst, const int ncols, const int nrows, cudaStream_t stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
+    const dim3 block_nums(1, block_num_y, 1);
+    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
+    mul_mat_vec_q<QK_K, QI5_K, block_q5_K, vec_dot_q5_K_q8_1>
+        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
+}
+
 static void convert_fp16_to_fp32_cuda(const void * vx, float * y, const int k, cudaStream_t stream) {
     const int num_blocks = (k + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE;
     dequantize_block<1, 1, convert_f16><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(vx, y, k);
@@ -2581,8 +2636,8 @@ inline void ggml_cuda_op_mul_mat_vec(
         src0->type == GGML_TYPE_Q8_0 ||
         src0->type == GGML_TYPE_Q2_K ||
         src0->type == GGML_TYPE_Q3_K ||
-        src0->type == GGML_TYPE_Q4_K;
-        // src0->type == GGML_TYPE_Q5_K ||
+        src0->type == GGML_TYPE_Q4_K ||
+        src0->type == GGML_TYPE_Q5_K;
         // src0->type == GGML_TYPE_Q5_K;
 
     const bool use_mul_mat_vec_q = g_compute_capabilities[id] >= 610 && mul_mat_vec_q_implemented;
@@ -2620,6 +2675,9 @@ inline void ggml_cuda_op_mul_mat_vec(
             case GGML_TYPE_Q4_K:
                 mul_mat_vec_q4_K_q8_1_cuda(src0_ddq_i, src1_q8_1, dst_ddf_i, ne00, nrows, cudaStream_main);
                 break;
+            case GGML_TYPE_Q5_K:
+                mul_mat_vec_q5_K_q8_1_cuda(src0_ddq_i, src1_q8_1, dst_ddf_i, ne00, nrows, cudaStream_main);
+                break;
             default:
                 GGML_ASSERT(false);
                 break;