q6_K works

Author: JohannesGaessler

ggml-cuda.cu

@@ -207,6 +207,8 @@ typedef struct {
 static_assert(sizeof(block_q5_K) == 2*sizeof(ggml_fp16_t) + K_SCALE_SIZE + QK_K/2 + QK_K/8, "wrong q5_K block size/padding");
 #endif
 
+#define QR6_K 2
+#define QI6_K (QK_K / (4*QR6_K))
 typedef struct {
     uint8_t ql[QK_K/2];   // quants, lower 4 bits
     uint8_t qh[QK_K/4];   // quants, upper 2 bits
@@ -1577,6 +1579,47 @@ static __device__ __forceinline__ float vec_dot_q5_K_q8_1(
 #endif // __CUDA_ARCH__ >= 610
 }
 
+static __device__ __forceinline__ float vec_dot_q6_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_q6_K * bq6_K = (const block_q6_K *) vbq;
+
+    const int bq8_offset = 4 * (iqs / 16) + (iqs % 16) / 8;
+
+    float sumf = 0.0f;
+
+    const float d = bq6_K->d;
+
+    int vil;
+    memcpy(&vil, &bq6_K->ql[sizeof(int) * iqs], sizeof(int));
+
+    int vih;
+    memcpy(&vih, &bq6_K->qh[sizeof(int) * (8 * (iqs / 16) + iqs % 8)], sizeof(int));
+
+    for (int i = 0; i < 2; ++i) {
+        const int isc = 8 * (iqs / 16) + (iqs % 16) / 4 + 4*i;
+        const int sc = bq6_K->scales[isc];
+
+        const block_q8_1 * bq8i = bq8_1 + bq8_offset + 2*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 >> (2*((iqs%16)/8) + 4*i)) << 4) & 0x30303030;
+
+        const int vii = __vsubss4((viil | viih), 0x20202020);
+
+        sumf += d8i * (__dp4a(vii, uii, 0) * sc);
+    }
+
+    return d*sumf;
+#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;
@@ -2171,6 +2214,15 @@ static void mul_mat_vec_q5_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_q6_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, QI6_K, block_q6_K, vec_dot_q6_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);
@@ -2637,8 +2689,8 @@ inline void ggml_cuda_op_mul_mat_vec(
         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_Q5_K;
+        src0->type == GGML_TYPE_Q5_K ||
+        src0->type == GGML_TYPE_Q6_K;
 
     const bool use_mul_mat_vec_q = g_compute_capabilities[id] >= 610 && mul_mat_vec_q_implemented;
 #endif
@@ -2678,6 +2730,9 @@ inline void ggml_cuda_op_mul_mat_vec(
             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;
+            case GGML_TYPE_Q6_K:
+                mul_mat_vec_q6_K_q8_1_cuda(src0_ddq_i, src1_q8_1, dst_ddf_i, ne00, nrows, cudaStream_main);
+                break;
             default:
                 GGML_ASSERT(false);
                 break;