Add q2_k and q3_k quantization

reeselevine · reeselevine · commit de4da8718370 · 2025-08-13T17:16:51.000-07:00
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -129,7 +129,7 @@ struct webgpu_context_struct {
     webgpu_buf_pool set_rows_error_buf_pool;
 
     wgpu::ComputePipeline memset_pipeline;
-    wgpu::ComputePipeline mul_mat_pipeline[10][2];
+    wgpu::ComputePipeline mul_mat_pipeline[12][2];
     wgpu::ComputePipeline set_rows_pipeline;
     wgpu::ComputePipeline cpy_pipeline;
 
@@ -910,7 +910,7 @@ static void ggml_webgpu_init_memset_pipeline(webgpu_context & webgpu_ctx) {
 }
 
 static void ggml_webgpu_init_mul_mat_pipeline(webgpu_context & webgpu_ctx) {
-    webgpu_pipeline_info pipeline_infos[16] = {
+    webgpu_pipeline_info pipeline_infos[10] = {
         { .name        = "mul_mat_f32_f32",
          .shader_code = wgsl_mul_mat_f32_f32,
          .src0_type   = GGML_TYPE_F32,
@@ -919,10 +919,6 @@ static void ggml_webgpu_init_mul_mat_pipeline(webgpu_context & webgpu_ctx) {
          .shader_code = wgsl_mul_mat_f16_f16,
          .src0_type   = GGML_TYPE_F16,
          .src1_type   = GGML_TYPE_F16 },
-        { .name        = "mul_mat_f32_f16",
-         .shader_code = wgsl_mul_mat_f32_f16,
-         .src0_type   = GGML_TYPE_F32,
-         .src1_type   = GGML_TYPE_F16 },
         { .name        = "mul_mat_f16_f32",
          .shader_code = wgsl_mul_mat_f16_f32,
          .src0_type   = GGML_TYPE_F16,
@@ -931,50 +927,30 @@ static void ggml_webgpu_init_mul_mat_pipeline(webgpu_context & webgpu_ctx) {
          .shader_code = wgsl_mul_mat_q4_0_f32,
          .src0_type   = GGML_TYPE_Q4_0,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q4_0_f16",
-         .shader_code = wgsl_mul_mat_q4_0_f16,
-         .src0_type   = GGML_TYPE_Q4_0,
-         .src1_type   = GGML_TYPE_F16 },
         { .name        = "mul_mat_q4_1_f32",
          .shader_code = wgsl_mul_mat_q4_1_f32,
          .src0_type   = GGML_TYPE_Q4_1,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q4_1_f16",
-         .shader_code = wgsl_mul_mat_q4_1_f16,
-         .src0_type   = GGML_TYPE_Q4_1,
-         .src1_type   = GGML_TYPE_F16 },
         { .name        = "mul_mat_q5_0_f32",
          .shader_code = wgsl_mul_mat_q5_0_f32,
          .src0_type   = GGML_TYPE_Q5_0,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q5_0_f16",
-         .shader_code = wgsl_mul_mat_q5_0_f16,
-         .src0_type   = GGML_TYPE_Q5_0,
-         .src1_type   = GGML_TYPE_F16 },
         { .name        = "mul_mat_q5_1_f32",
          .shader_code = wgsl_mul_mat_q5_1_f32,
          .src0_type   = GGML_TYPE_Q5_1,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q5_1_f16",
-         .shader_code = wgsl_mul_mat_q5_1_f16,
-         .src0_type   = GGML_TYPE_Q5_1,
-         .src1_type   = GGML_TYPE_F16 },
         { .name        = "mul_mat_q8_0_f32",
          .shader_code = wgsl_mul_mat_q8_0_f32,
          .src0_type   = GGML_TYPE_Q8_0,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q8_0_f16",
-         .shader_code = wgsl_mul_mat_q8_0_f16,
-         .src0_type   = GGML_TYPE_Q8_0,
-         .src1_type   = GGML_TYPE_F16 },
-        { .name        = "mul_mat_q8_1_f32",
-         .shader_code = wgsl_mul_mat_q8_1_f32,
-         .src0_type   = GGML_TYPE_Q8_1,
+        { .name        = "mul_mat_q2_k_f32",
+         .shader_code = wgsl_mul_mat_q2_k_f32,
+         .src0_type   = GGML_TYPE_Q2_K,
          .src1_type   = GGML_TYPE_F32 },
-        { .name        = "mul_mat_q8_1_f16",
-         .shader_code = wgsl_mul_mat_q8_1_f16,
-         .src0_type   = GGML_TYPE_Q8_1,
-         .src1_type   = GGML_TYPE_F16 }
+        { .name        = "mul_mat_q3_k_f32",
+         .shader_code = wgsl_mul_mat_q3_k_f32,
+         .src0_type   = GGML_TYPE_Q3_K,
+         .src1_type   = GGML_TYPE_F32 }
     };
 
     for (auto & pipeline_info : pipeline_infos) {
@@ -1058,28 +1034,29 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
         case GGML_OP_CPY:
         case GGML_OP_SET_ROWS:
             return op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F32;
-        case GGML_OP_MUL_MAT:
-            switch(op->src[0]->type) {
-                case GGML_TYPE_F32:
-                case GGML_TYPE_F16:
-                case GGML_TYPE_Q4_0:
-                case GGML_TYPE_Q4_1:
-                case GGML_TYPE_Q5_0:
-                case GGML_TYPE_Q5_1:
-                case GGML_TYPE_Q8_0:
-                case GGML_TYPE_Q8_1:
-                    break;
-                default:
-                    return false;
-            }
+        case GGML_OP_MUL_MAT: {
             switch(op->src[1]->type) {
-                case GGML_TYPE_F32:
                 case GGML_TYPE_F16:
-                    break;
+                    return op->src[0]->type == GGML_TYPE_F16;
+                case GGML_TYPE_F32:
+                    switch(op->src[0]->type) {
+                        case GGML_TYPE_F32:
+                        case GGML_TYPE_F16:
+                        case GGML_TYPE_Q4_0:
+                        case GGML_TYPE_Q4_1:
+                        case GGML_TYPE_Q5_0:
+                        case GGML_TYPE_Q5_1:
+                        case GGML_TYPE_Q8_0:
+                        case GGML_TYPE_Q2_K:
+                        case GGML_TYPE_Q3_K:
+                            return true;
+                        default:
+                            return false;
+                    }
                 default:
                     return false;
             }
-            return true;
+        }
         default:
             return false;
     }
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat.tmpl.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat.tmpl.wgsl
@@ -25,14 +25,6 @@
     },
     "DECLS" : "FLOAT"
   },
-  {
-    "REPLS": {
-      "SRC0_TYPE" : "f32",
-      "SRC1_TYPE" : "f16",
-      "BLOCK_SIZE" : 1
-    },
-    "DECLS" : "FLOAT"
-  },
   {
     "REPLS": {
       "SRC0_TYPE": "q4_0",
@@ -41,14 +33,6 @@
     },
     "DECLS": "Q4_0"
   },
-  {
-    "REPLS": {
-      "SRC0_TYPE": "q4_0",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
-    },
-    "DECLS": "Q4_0"
-  },
   {
     "REPLS": {
       "SRC0_TYPE": "q4_1",
@@ -57,14 +41,6 @@
     },
     "DECLS": "Q4_1"
   },
-  {
-    "REPLS": {
-      "SRC0_TYPE": "q4_1",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
-    },
-    "DECLS": "Q4_1"
-  },
   {
     "REPLS": {
       "SRC0_TYPE": "q5_0",
@@ -73,14 +49,6 @@
     },
     "DECLS": "Q5_0"
   },
-  {
-    "REPLS": {
-      "SRC0_TYPE": "q5_0",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
-    },
-    "DECLS": "Q5_0"
-  },
   {
     "REPLS": {
       "SRC0_TYPE": "q5_1",
@@ -89,14 +57,6 @@
     },
     "DECLS": "Q5_1"
   },
-  {
-    "REPLS": {
-      "SRC0_TYPE": "q5_1",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
-    },
-    "DECLS": "Q5_1"
-  },
   {
     "REPLS": {
       "SRC0_TYPE": "q8_0",
@@ -107,27 +67,19 @@
   },
   {
     "REPLS": {
-      "SRC0_TYPE": "q8_0",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
-    },
-    "DECLS": "Q8_0"
-  },
-  {
-    "REPLS": {
-      "SRC0_TYPE": "q8_1",
+      "SRC0_TYPE": "q2_k",
       "SRC1_TYPE": "f32",
-      "BLOCK_SIZE": 32
+      "BLOCK_SIZE": 256
     },
-    "DECLS": "Q8_1"
+    "DECLS": "Q2_K"
   },
   {
     "REPLS": {
-      "SRC0_TYPE": "q8_1",
-      "SRC1_TYPE": "f16",
-      "BLOCK_SIZE": 32
+      "SRC0_TYPE": "q3_k",
+      "SRC1_TYPE": "f32",
+      "BLOCK_SIZE": 256
     },
-    "DECLS": "Q8_1"
+    "DECLS": "Q3_K"
   }
 ]
 
@@ -300,6 +252,116 @@ fn multiply_add(src0_idx_base: u32, src1_idx_base: u32, offset: u32) -> f32 {
 }
 #enddecl(Q8_1)
 
+#decl(Q2_K)
+// 16 blocks of 16 elements each
+struct q2_k {
+    scales: array<u32, 4>,
+    qs: array<u32, 16>,
+    d: f16,
+    dmin: f16
+};
+
+fn multiply_add(src0_idx_base: u32, src1_idx_base: u32, offset: u32) -> f32 {
+    let block = src0[src0_idx_base + offset];
+    let d = f32(block.d);
+    let m = f32(block.dmin);
+    var sum = 0.0;
+    var src1_i = src1_idx_base + offset * 256;
+    var is: u32 = 0;
+    // 2 halves of the block (128 elements each)
+    for (var q_b_idx: u32 = 0; q_b_idx < 64; q_b_idx += 32) {
+        // 4 groups (each group has 2 blocks of 16 elements)
+        for (var shift: u32 = 0; shift < 8; shift += 2) {
+            // 2 blocks
+            for (var k: u32 = 0; k < 32; k += 16) {
+                let sc = (block.scales[is / 4] >> ((is % 4) * 8)) & 0xFF;
+                is++;
+                let dl = d * f32(sc & 0xF);
+                let ml = m * f32(sc >> 4);
+                for (var l: u32 = 0u; l < 16; l++) {
+                    let q_idx = q_b_idx + k + l;
+                    let q_byte = (block.qs[q_idx / 4] >> ((q_idx % 4) * 8)) & 0xFF;
+                    let qs_val = (q_byte >> shift) & 3;
+                    sum += (f32(qs_val) * dl - ml) * src1[src1_i];
+                    src1_i++;
+                }
+            }
+        }
+    }
+    return sum;
+}
+
+#enddecl(Q2_K)
+
+#decl(Q3_K)
+// 16 blocks of 16 elements each
+struct q3_k {
+    hmask: array<f16, 16>,
+    qs: array<f16, 32>,
+    scales: array<f16, 6>, // 6-bit quantized values
+    d: f16
+};
+
+fn multiply_add(src0_idx_base: u32, src1_idx_base: u32, offset: u32) -> f32 {
+    let block = src0[src0_idx_base + offset];
+    let d = f32(block.d);
+
+    // extract 6-bit scales, which consist of 4-bits from first 8 bytes of scale,
+    // and 2-bits from the last 4 bytes
+    let kmask1: u32 = 0x03030303;
+    let kmask2: u32 = 0x0f0f0f0f;
+    var scale_vals: array<u32, 4>;
+    for (var i: u32 = 0; i < 4; i++) {
+        scale_vals[i] = bitcast<u32>(vec2(block.scales[2 * i], block.scales[2 * i + 1]));
+    }
+    var tmp: u32 = scale_vals[2];
+    scale_vals[2] = ((scale_vals[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+    scale_vals[3] = ((scale_vals[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+    scale_vals[0] = (scale_vals[0] & kmask2) | ((tmp & kmask1) << 4);
+    scale_vals[1] = (scale_vals[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+
+    // convert half-precision floats to packed 32-bit integers
+    var hmask_vals: array<u32, 8>;
+    for (var i: u32 = 0; i < 8; i++) {
+        hmask_vals[i] = bitcast<u32>(vec2(block.hmask[2 * i], block.hmask[2 * i + 1]));
+    }
+    var qs_vals: array<u32, 16>;
+    for (var i: u32 = 0; i < 16; i++) {
+        qs_vals[i] = bitcast<u32>(vec2(block.qs[2 * i], block.qs[2 * i + 1]));
+    }
+
+    var sum = 0.0;
+    var src1_i = src1_idx_base + offset * 256;
+    var is: u32 = 0;
+    var m: u32 = 1;
+    // 2 halves of the block (128 elements each)
+    for (var q_b_idx: u32 = 0; q_b_idx < 64; q_b_idx += 32) {
+        // 4 groups (each group has 2 blocks of 16 elements)
+        for (var shift: u32 = 0; shift < 8; shift += 2) {
+            // 2 blocks
+            for (var k: u32 = 0; k < 32; k += 16) {
+                let sc = (scale_vals[is / 4] >> ((is % 4) * 8)) & 0xFF;
+                is++;
+                let dl = d * (f32(sc) - 32.0);
+                for (var l: u32 = 0u; l < 16u; l++) {
+                    let q_idx = q_b_idx + k + l;
+                    let hm_idx = k + l;
+                    let q_byte = (qs_vals[q_idx / 4] >> ((q_idx % 4) * 8)) & 0xFF;
+                    let hmask_byte = (hmask_vals[hm_idx / 4] >> ((hm_idx % 4) * 8)) & 0xFF;
+                    let hm = select(4.0, 0.0, (hmask_byte & m) != 0);
+                    let qs_val = (q_byte >> shift) & 3;
+                    sum += ((f32(qs_val) - hm) * dl) * src1[src1_i];
+                    src1_i++;
+                }
+            }
+            m <<= 1;
+        }
+    }
+    return sum;
+}
+
+#enddecl(Q3_K)
+
 #end(DECLS)
 
 #define(SHADER)
