q8_1 quantization

Author: JohannesGaessler

ggml-cuda.cu

@@ -70,9 +70,11 @@ typedef void (*ggml_cuda_op_t)(
 
 // QK = number of values after dequantization
 // QR = QK / number of values before dequantization
+// QI = number of 32 bit integers before dequantization
 
 #define QK4_0 32
 #define QR4_0 2
+#define QI4_0 4
 typedef struct {
     half    d;              // delta
     uint8_t qs[QK4_0 / 2];  // nibbles / quants
@@ -81,6 +83,7 @@ static_assert(sizeof(block_q4_0) == sizeof(ggml_fp16_t) + QK4_0 / 2, "wrong q4_0
 
 #define QK4_1 32
 #define QR4_1 2
+#define QI4_1 4
 typedef struct {
     half    d;              // delta
     half    m;              // min
@@ -115,7 +118,16 @@ typedef struct {
 } block_q8_0;
 static_assert(sizeof(block_q8_0) == sizeof(ggml_fp16_t) + QK8_0, "wrong q8_0 block size/padding");
 
-typedef float (*vec_dot_q_cuda_t)(const void * vbq, const block_q8_0 * bq8_0, const int iqs);
+#define QK8_1 32
+#define QR8_1 1
+typedef struct {
+    half    d;              // delta
+    half    s;              // unquantized sum
+    int8_t  qs[QK8_0];      // quants
+} block_q8_1;
+static_assert(sizeof(block_q8_1) == 2*sizeof(ggml_fp16_t) + QK8_0, "wrong q8_1 block size/padding");
+
+typedef float (*vec_dot_q_cuda_t)(const void * vbq, const block_q8_1 * bq8_1, const int iqs);
 
 //================================= k-quants
 
@@ -1155,25 +1167,27 @@ static __device__ void convert_f16(const void * vx, const int ib, const int iqs,
     v.y = x[ib + iqs + 1];
 }
 
-static __global__ void quantize_q8_0(const float * x, void * vy, const int k) {
+static __global__ void quantize_q8_1(const float * x, void * vy, const int k) {
     const int i = blockDim.x*blockIdx.x + threadIdx.x;
 
     if (i >= k) {
         return;
     }
 
-    block_q8_0 * y = (block_q8_0 *) vy;
+    block_q8_1 * y = (block_q8_1 *) vy;
 
     const int ib = i / QK8_0; // block index
     const int iqs = i % QK8_0; // quant index
 
     const float xi = x[i];
     float amax = fabsf(xi);
+    float sum = xi;
 
     __syncwarp();
 #pragma unroll
     for (int mask = 16; mask > 0; mask >>= 1) {
         amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, mask, 32));
+        sum += __shfl_xor_sync(0xffffffff, sum, mask, 32);
     }
 
     const float d = amax / 127;
@@ -1186,51 +1200,47 @@ static __global__ void quantize_q8_0(const float * x, void * vy, const int k) {
     }
 
     y[ib].d = d;
+    y[ib].s = sum;
 }
 
-static __device__ float vec_dot_q4_0_q8_0(const void * vbq, const block_q8_0 * bq8_0, const int iqs) {
+static __device__ __forceinline__ float vec_dot_q4_0_q8_1(const void * vbq, const block_q8_1 * bq8_1, const int iqs) {
     const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq;
 
     int vi;
-    int ui0, ui1;
     memcpy(&vi,  &bq4_0->qs[sizeof(int) * (iqs + 0)], sizeof(int));
-    memcpy(&ui0, &bq8_0->qs[sizeof(int) * (iqs + 0)], sizeof(int));
-    memcpy(&ui1, &bq8_0->qs[sizeof(int) * (iqs + 4)], sizeof(int));
+    const int ui0 = *((int *) &bq8_1->qs[sizeof(int) * (iqs + 0)]);
+    const int ui1 = *((int *) &bq8_1->qs[sizeof(int) * (iqs + QI4_0)]);
 
-    const float d = bq4_0->d * bq8_0->d;
+    const float d = bq4_0->d * bq8_1->d;
 
     const int vi0 = __vsub4((vi >> 0) & 0x0F0F0F0F, 0x08080808);
     const int vi1 = __vsub4((vi >> 4) & 0x0F0F0F0F, 0x08080808);
 
-    const int sumi0 = __dp4a(vi0, ui0, 0);
-    const int sumi1 = __dp4a(vi1, ui1, 0);
+    int sumi = __dp4a(vi0, ui0, 0);
+    sumi     = __dp4a(vi1, ui1, sumi);
 
-    return (sumi0 + sumi1)*d;
+    return sumi*d;
 
 }
 
-static __device__ float vec_dot_q4_1_q8_0(const void * vbq, const block_q8_0 * bq8_0, const int iqs) {
+static __device__ __forceinline__ float vec_dot_q4_1_q8_1(const void * vbq, const block_q8_1 * bq8_1, const int iqs) {
     const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq;
 
-    int vi = *((int *) &bq4_1->qs[sizeof(int) * (iqs + 0)]);
-    int ui0, ui1;
-    memcpy(&ui0, &bq8_0->qs[sizeof(int) * (iqs + 0)], sizeof(int));
-    memcpy(&ui1, &bq8_0->qs[sizeof(int) * (iqs + 4)], sizeof(int));
+    const int vi  = *((int *) &bq4_1->qs[sizeof(int) * (iqs + 0)]);
+    const int ui0 = *((int *) &bq8_1->qs[sizeof(int) * (iqs + 0)]);
+    const int ui1 = *((int *) &bq8_1->qs[sizeof(int) * (iqs + QI4_1)]);
 
-    const float d4_1 = bq4_1->d;
-    const float m    = bq4_1->m;
-    const float d8_0 = bq8_0->d;
+    const float d = bq4_1->d * bq8_1->d;
+    const float m = bq4_1->m;
+    const float s = bq8_1->s;
 
     const int vi0 = (vi >> 0) & 0x0F0F0F0F;
     const int vi1 = (vi >> 4) & 0x0F0F0F0F;
 
-    const int sumi0 = __dp4a(vi0, ui0, 0);
-    const int sumi1 = __dp4a(vi1, ui1, 0);
-
-    const int sumi2 = __dp4a(0x01010101, ui0, 0);
-    const int sumi3 = __dp4a(0x01010101, ui1, 0);
+    int sumi = __dp4a(vi0, ui0, 0);
+    sumi     = __dp4a(vi1, ui1, sumi);
 
-    return (sumi0 + sumi1)*d4_1*d8_0 + (sumi2 + sumi3)*m*d8_0;
+    return sumi*d + m*s / QI4_1;
 
 }
 
@@ -1263,8 +1273,6 @@ static __global__ void mul_mat_vec_q(const void * vx, const void * vy, float * d
         return;
     }
 
-    const int tid = threadIdx.x;
-
     const int blocks_per_row = ncols / qk;
     const int blocks_per_warp = WARP_SIZE * sizeof(int)*2/qk;
     const int ints_per_block = qk / (2 * sizeof(int));
@@ -1273,14 +1281,14 @@ static __global__ void mul_mat_vec_q(const void * vx, const void * vy, float * d
     float tmp = 0.0f;
 
     const block_q_t  * x = (const block_q_t  *) vx;
-    const block_q8_0 * y = (const block_q8_0 *) vy;
+    const block_q8_1 * y = (const block_q8_1 *) vy;
 
     for (int i = 0; i < blocks_per_row; i += blocks_per_warp) {
-        const int ibx = row*blocks_per_row + i + tid/ints_per_block; // x block index
+        const int ibx = row*blocks_per_row + i + threadIdx.x/ints_per_block; // x block index
 
-        const int iby = i + tid/ints_per_block;
+        const int iby = i + threadIdx.x/ints_per_block;
 
-        const int iqs  = tid % ints_per_block;
+        const int iqs  = threadIdx.x % ints_per_block;
 
         tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);
     }
@@ -1292,7 +1300,7 @@ static __global__ void mul_mat_vec_q(const void * vx, const void * vy, float * d
         tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32);
     }
 
-    if (tid == 0) {
+    if (threadIdx.x == 0) {
         dst[row] = tmp;
     }
 }
@@ -1612,9 +1620,9 @@ static void rms_norm_f32_cuda(const float * x, float * dst, const int ncols, con
     rms_norm_f32<<<nrows, block_dims, 0, stream>>>(x, dst, ncols);
 }
 
-static void quantize_row_q8_0_cuda(const float * x, void * vy, const int k, cudaStream_t stream) {
+static void quantize_row_q8_1_cuda(const float * x, void * vy, const int k, cudaStream_t stream) {
     const int num_blocks = (k + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE;
-    quantize_q8_0<<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, vy, k);
+    quantize_q8_1<<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(x, vy, k);
 }
 
 static void dequantize_row_q4_0_cuda(const void * vx, float * y, const int k, cudaStream_t stream) {
@@ -1770,21 +1778,21 @@ static void dequantize_mul_mat_vec_q6_K_cuda(const void * vx, const float * y, f
     dequantize_mul_mat_vec_q6_k<<<block_nums, block_dims, 0, stream>>>(vx, y, dst, ncols, nrows);
 }
 
-static void mul_mat_vec_q4_0_q8_0_cuda(const void * vx, const void * vy, float * dst, const int ncols, const int nrows, cudaStream_t stream) {
+static void mul_mat_vec_q4_0_q8_1_cuda(const void * vx, const void * vy, float * dst, const int ncols, const int nrows, cudaStream_t stream) {
     GGML_ASSERT(ncols % GGML_CUDA_DMMV_X == 0);
     const int block_num_y = (nrows + GGML_CUDA_DMMV_Y - 1) / GGML_CUDA_DMMV_Y;
     const dim3 block_nums(1, block_num_y, 1);
     const dim3 block_dims(WARP_SIZE, GGML_CUDA_DMMV_Y, 1);
-    mul_mat_vec_q<QK4_0, block_q4_0, vec_dot_q4_0_q8_0>
+    mul_mat_vec_q<QK4_0, block_q4_0, vec_dot_q4_0_q8_1>
         <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
 }
 
-static void mul_mat_vec_q4_1_q8_0_cuda(const void * vx, const void * vy, float * dst, const int ncols, const int nrows, cudaStream_t stream) {
+static void mul_mat_vec_q4_1_q8_1_cuda(const void * vx, const void * vy, float * dst, const int ncols, const int nrows, cudaStream_t stream) {
     GGML_ASSERT(ncols % GGML_CUDA_DMMV_X == 0);
     const int block_num_y = (nrows + GGML_CUDA_DMMV_Y - 1) / GGML_CUDA_DMMV_Y;
     const dim3 block_nums(1, block_num_y, 1);
     const dim3 block_dims(WARP_SIZE, GGML_CUDA_DMMV_Y, 1);
-    mul_mat_vec_q<QK4_0, block_q4_1, vec_dot_q4_1_q8_0>
+    mul_mat_vec_q<QK4_0, block_q4_1, vec_dot_q4_1_q8_1>
         <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
 }
 
@@ -2302,14 +2310,14 @@ inline void ggml_cuda_op_mul_mat_vec_q(
 
     size_t as;
     void * src1_q8_0 = ggml_cuda_pool_malloc(ne00*sizeof(block_q8_0)/QK8_0, &as);
-    quantize_row_q8_0_cuda(src1_ddf_i, src1_q8_0, ne00, cudaStream_main);
+    quantize_row_q8_1_cuda(src1_ddf_i, src1_q8_0, ne00, cudaStream_main);
 
     switch (src0->type) {
         case GGML_TYPE_Q4_0:
-            mul_mat_vec_q4_0_q8_0_cuda(src0_ddq_i, src1_q8_0, dst_ddf_i, ne00, nrows, cudaStream_main);
+            mul_mat_vec_q4_0_q8_1_cuda(src0_ddq_i, src1_q8_0, dst_ddf_i, ne00, nrows, cudaStream_main);
             break;
         case GGML_TYPE_Q4_1:
-            mul_mat_vec_q4_1_q8_0_cuda(src0_ddq_i, src1_q8_0, dst_ddf_i, ne00, nrows, cudaStream_main);
+            mul_mat_vec_q4_1_q8_1_cuda(src0_ddq_i, src1_q8_0, dst_ddf_i, ne00, nrows, cudaStream_main);
             break;
         default:
             GGML_ASSERT(false);