Add q4_0 x q8_1 matrix matrix multiplication support

Author: 0cc4m

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1897,6 +1897,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
         CREATE_MM(GGML_TYPE_IQ4_XS,  pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS].f16acc,  matmul_iq4_xs_f32,  _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
         CREATE_MM(GGML_TYPE_IQ4_NL,  pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f16acc,  matmul_iq4_nl_f32,  _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
 
+        CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0].f16acc, matmul_q4_0_q8_1, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
         CREATE_MMQ(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q8_0].f16acc, matmul_q8_0_q8_1, _f16acc, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
 
         CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 4, _id);
@@ -3161,6 +3162,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte
     // MMQ
     if (src1_type == GGML_TYPE_Q8_1) {
         switch (src0_type) {
+            case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q8_0:
                 break;
             default:
@@ -3893,8 +3895,8 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, int m, int
     return split_k;
 }
 
-static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned, ggml_type src0_type) {
-    VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ", " << ggml_type_name(src0_type) << ")");
+static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, bool aligned, ggml_type src0_type, ggml_type src1_type) {
+    VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")");
 
     if (ctx->device->coopmat2) {
         if ((ctx->device->mul_mat_l[src0_type] && (m % mmp->l->wg_denoms[0]) == 0 && (n % mmp->l->wg_denoms[1]) == 0) || (!ctx->device->mul_mat_m[src0_type] && !ctx->device->mul_mat_s[src0_type])) {
@@ -3906,7 +3908,7 @@ static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx,
         return aligned ? mmp->a_s : mmp->s;
     }
 
-    if ((ctx->device->mul_mat_s[src0_type] && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m[src0_type] && !ctx->device->mul_mat_l[src0_type])) {
+    if ((ctx->device->mul_mat_s[src0_type] && (m <= 32 || n <= 32)) || (!ctx->device->mul_mat_m[src0_type] && !ctx->device->mul_mat_l[src0_type]) || src1_type == GGML_TYPE_Q8_1) {
         return aligned ? mmp->a_s : mmp->s;
     }
     if ((ctx->device->mul_mat_m[src0_type] && (m <= 64 || n <= 64)) || !ctx->device->mul_mat_l[src0_type]) {
@@ -3915,9 +3917,9 @@ static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx,
     return aligned ? mmp->a_l : mmp->l;
 }
 
-static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, ggml_type src0_type) {
-    VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ", " << ggml_type_name(src0_type) << ")");
-    return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true, src0_type)->align;
+static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, ggml_type src0_type, ggml_type src1_type) {
+    VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")");
+    return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true, src0_type, src1_type)->align;
 }
 
 static void ggml_vk_matmul(
@@ -4177,10 +4179,10 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
     const int y_ne = ne11 * ne10;
     const int d_ne = ne11 * ne01;
 
-    const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11, qx_needs_dequant ? GGML_TYPE_F16 : src0->type));
+    const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11, qx_needs_dequant ? GGML_TYPE_F16 : src0->type, quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type)));
     const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && ne11 > 8;
 
-    vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned, qx_needs_dequant ? GGML_TYPE_F16 : src0->type);
+    vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned, qx_needs_dequant ? GGML_TYPE_F16 : src0->type, quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type));
 
     const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, pipeline);
 
@@ -7306,8 +7308,8 @@ static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m,
     ggml_vk_quantize_data(x, qx, x_ne, quant);
 
     for (size_t i = 0; i < y_ne; i++) {
-        y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f;
-        // y[i] = (i % k == i / k) ? 1.0f : 0.0f;
+        // y[i] = (rand() / (float)RAND_MAX) * 2.0f - 1.0f;
+        y[i] = (i % k == i / k) ? 1.0f : 0.0f;
     }
 
     ggml_pipeline_request_descriptor_sets(ctx->device, p, num_it);
@@ -7481,9 +7483,9 @@ static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx) {
     };
     const size_t num_it = 100;
 
-    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 0, GGML_TYPE_Q8_0, true);
-    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 1, GGML_TYPE_Q8_0, true);
-    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 2, GGML_TYPE_Q8_0, true);
+    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 0, GGML_TYPE_Q4_0, true);
+    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 1, GGML_TYPE_Q4_0, true);
+    ggml_vk_test_dequant_matmul(ctx, 4096, 512, 4096, 2, num_it, 1, 2, GGML_TYPE_Q4_0, true);
 
     abort();
 

ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_mmq.comp

@@ -0,0 +1,37 @@
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+
+#include "types.comp"
+
+// Each iqs value maps to a 32-bit integer
+
+#if defined(DATA_A_Q4_0)
+i32vec2 repack(uint ib, uint iqs) {
+    const u16vec2 quants = u16vec2(data_a[ib].qs[iqs * 2 + 1],
+                                   data_a[ib].qs[iqs * 2    ]);
+    const uint32_t vui = pack32(quants);
+    return i32vec2(pack32(i8vec4(i16vec4(unpack8( vui       & 0x0F0F0F0F)) - int16_t(8))),
+                   pack32(i8vec4(i16vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - int16_t(8))));
+}
+#endif
+
+#if defined(DATA_A_Q8_0)
+int32_t repack(uint ib, uint iqs) {
+    const int16_t v0 = data_a[ib].qs[iqs * 2    ];
+    const int16_t v1 = data_a[ib].qs[iqs * 2 + 1];
+    return pack32(i16vec2(v1, v0));
+}
+#endif
+
+#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
+FLOAT_TYPE get_d(uint ib) {
+    return FLOAT_TYPE(data_a[ib].d);
+}
+#endif
+
+#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1)
+FLOAT_TYPE_VEC2 get_dm(uint ib) {
+    return FLOAT_TYPE_VEC2(data_a[ib].d, data_a[ib].m);
+}
+#endif

ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp

@@ -81,12 +81,12 @@ layout (constant_id = 10) const uint WARP = 32;
 
 // Assumption: BK == 32
 struct block_q8_0_internal {
-    FLOAT_TYPE ds;
+    FLOAT_TYPE d;
     int32_t qs[BK / 4];
 };
 
 struct block_q8_1_internal {
-    FLOAT_TYPE_VEC2 ds;
+    FLOAT_TYPE_VEC2 dm;
     int32_t qs[BK / 4];
 };
 
@@ -97,9 +97,10 @@ shared block_q8_0_internal buf_a[BM];
 shared block_q8_1_internal buf_a[BM];
 #endif
 
-shared block_q8_1_internal buf_b[BN];
+shared block_q8_0_internal buf_b[BN];
 
-#define LOAD_VEC 4
+#define LOAD_VEC_A (4 * QUANT_R)
+#define LOAD_VEC_B 4
 
 #ifdef MUL_MAT_ID
 shared u16vec2 row_ids[3072];
@@ -111,6 +112,8 @@ shared u16vec2 row_ids[3072];
 shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS];
 #endif
 
+#include "dequant_funcs_mmq.comp"
+
 void main() {
 #if defined(DATA_A_IQ4_NL)
     init_iq4nl_shmem();
@@ -162,10 +165,13 @@ void main() {
     const uint warp_r = warp_i % (BM / WM);
     const uint warp_c = warp_i / (BM / WM);
 
-    const uint loadr = gl_LocalInvocationID.x % (BK / LOAD_VEC);
-    const uint loadc = gl_LocalInvocationID.x / (BK / LOAD_VEC);
+    const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A);
+    const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A);
+    const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B);
+    const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B);
 
-    const uint loadstride = BLOCK_SIZE * LOAD_VEC / BK;
+    const uint loadstride_a = BLOCK_SIZE * LOAD_VEC_A / BK;
+    const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK;
 
 #ifdef MUL_MAT_ID
     uint _ne1 = 0;
@@ -222,46 +228,52 @@ void main() {
     block_q8_1_internal cache_a[WMITER * TM];
 #endif
 
-    block_q8_1_internal cache_b[TN];
+    block_q8_0_internal cache_b[TN];
 
     [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) {
         sums[i] = ACC_TYPE(0.0f);
     }
 #endif
 
     for (uint block = start_k; block < end_k; block += BK) {
-        [[unroll]] for (uint l = 0; loadc + l < BM; l += loadstride) {
-#if defined(DATA_A_Q8_0)
-            const uint ib = pos_a_ib + (loadc + l) * p.stride_a / BK;
-            const uint iqs = loadr;
+        [[unroll]] for (uint l = 0; loadc_a + l < BM; l += loadstride_a) {
+            const uint ib = pos_a_ib + (loadc_a + l) * p.stride_a / BK;
+            const uint iqs = loadr_a;
 
-            const uint buf_ib = loadc + l;
+            const uint buf_ib = loadc_a + l;
 
             // Should ds be gated to a single thread?
             if (iqs == 0) {
-                buf_a[buf_ib].ds = FLOAT_TYPE(data_a[ib].d);
+#if QUANT_AUXF == 1
+                buf_a[buf_ib].d = get_d(ib);
+#else
+                buf_a[buf_ib].dm = get_dm(ib, 0);
+#endif
             }
-            const int16_t v0 = data_a[ib].qs[iqs * 2    ];
-            const int16_t v1 = data_a[ib].qs[iqs * 2 + 1];
-            buf_a[buf_ib].qs[iqs] = pack32(i16vec2(v1, v0));
+#if QUANT_R == 1
+            buf_a[buf_ib].qs[iqs] = repack(ib, iqs);
+#else
+            const i32vec2 vals = repack(ib, iqs);
+            buf_a[buf_ib].qs[iqs] = vals.x;
+            buf_a[buf_ib].qs[iqs + 4] = vals.y;
 #endif
         }
-        [[unroll]] for (uint l = 0; loadc + l < BN; l += loadstride) {
+        [[unroll]] for (uint l = 0; loadc_b + l < BN; l += loadstride_b) {
 #ifdef MUL_MAT_ID
-            const u16vec2 row_idx = row_ids[ic * BN + loadc + l];
-            const uint idx = pos_b_ib + row_idx.y * p.batch_stride_b / LOAD_VEC + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC + loadr;
+            const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l];
+            const uint idx = pos_b_ib + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + loadr_b;
             const uint ib = idx / 8;
             const uint iqs = idx & 0x7;
 #else
-            const uint ib = pos_b_ib + (loadc + l) * p.stride_b / BK;
-            const uint iqs = loadr;
+            const uint ib = pos_b_ib + (loadc_b + l) * p.stride_b / BK;
+            const uint iqs = loadr_b;
 #endif
 
-            const uint buf_ib = loadc + l;
+            const uint buf_ib = loadc_b + l;
 
             // Should ds be gated to a single thread?
             if (iqs == 0) {
-                buf_b[buf_ib].ds = FLOAT_TYPE_VEC2(data_b[ib].ds);
+                buf_b[buf_ib].d = FLOAT_TYPE(data_b[ib].ds.x);
             }
             const int32_t v0 = int32_t(data_b[ib].qs[iqs * 2    ]);
             const int32_t v1 = int32_t(data_b[ib].qs[iqs * 2 + 1]);
@@ -311,7 +323,12 @@ void main() {
                             q_sum = dotPacked4x8AccSatEXT(cache_a[cache_a_idx].qs[idx_k], cache_b[cc].qs[idx_k], q_sum);
                         }
 
-                        const float factor = float(cache_a[cache_a_idx].ds) * float(cache_b[cc].ds.x);
+#if QUANT_AUXF == 1
+                        const float factor = float(cache_a[cache_a_idx].d) * float(cache_b[cc].d);
+#else
+                        // TODO
+                        // const float factor = float(cache_a[cache_a_idx].d) * float(cache_b[cc].d);
+#endif
 
                         sums[sums_idx] = ACC_TYPE(fma(float(q_sum), factor, float(sums[sums_idx])));
                     }

ggml/src/ggml-vulkan/vulkan-shaders/types.comp

@@ -49,6 +49,7 @@ struct block_q4_0_packed16
 #if defined(DATA_A_Q4_0)
 #define QUANT_K QUANT_K_Q4_0
 #define QUANT_R QUANT_R_Q4_0
+#define QUANT_AUXF 1
 #define A_TYPE block_q4_0
 #define A_TYPE_PACKED16 block_q4_0_packed16
 #endif
@@ -73,6 +74,7 @@ struct block_q4_1_packed16
 #if defined(DATA_A_Q4_1)
 #define QUANT_K QUANT_K_Q4_1
 #define QUANT_R QUANT_R_Q4_1
+#define QUANT_AUXF 2
 #define A_TYPE block_q4_1
 #define A_TYPE_PACKED16 block_q4_1_packed16
 #endif
@@ -97,6 +99,7 @@ struct block_q5_0_packed16
 #if defined(DATA_A_Q5_0)
 #define QUANT_K QUANT_K_Q5_0
 #define QUANT_R QUANT_R_Q5_0
+#define QUANT_AUXF 1
 #define A_TYPE block_q5_0
 #define A_TYPE_PACKED16 block_q5_0_packed16
 #endif
@@ -123,6 +126,7 @@ struct block_q5_1_packed16
 #if defined(DATA_A_Q5_1)
 #define QUANT_K QUANT_K_Q5_1
 #define QUANT_R QUANT_R_Q5_1
+#define QUANT_AUXF 2
 #define A_TYPE block_q5_1
 #define A_TYPE_PACKED16 block_q5_1_packed16
 #endif

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -349,7 +349,7 @@ void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool
             string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name,  merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a},           {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc);
         }
 
-        if (!coopmat2 && !coopmat && tname == "q8_0") {
+        if (!coopmat2 && !coopmat && (tname == "q4_0" || tname == "q8_0")) {
             string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc);
         }
     }