vulkan: integer-dot mat-vec kernel for PTQ1_0 (8x faster decode) and PQ2_0 support

ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl

@@ -168,6 +168,17 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
 }
 #endif
 
+#if defined(DATA_A_PQ2_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]) >> (2u * (iqs % 4u));
+    return vec2(bits & 3u, (bits >> 2u) & 3u) - 1.0f;
+}
+vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]);
+    return vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f;
+}
+#endif
+
 #if defined(DATA_A_IQ1_S)
 vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     const uint ib32 = iqs / 32;
@@ -572,7 +583,7 @@ vec2 get_dm(uint ib, uint a_offset) {
 }
 #endif
 
-#if defined(DATA_A_Q2_0) || 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)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_PQ2_0) || 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)
 vec2 get_dm(uint ib, uint a_offset) {
     return vec2(float(data_a[a_offset + ib].d), 0);
 }

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

@@ -0,0 +1,30 @@
+#version 450
+
+#include "dequant_head.glsl"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_pq2_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+// 64 invocations cover 256 elements = two 128-weight blocks, 4 elements (one byte) each.
+void main() {
+    const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+    const uint tid = gl_LocalInvocationID.x % 64;
+    const uint il  = tid % 32;
+    const uint ir  = tid / 32;
+    const uint ib = 2*i + ir;
+    if (ib >= p.nel / QUANT_K_PQ2_0) {
+        return;
+    }
+
+    const uint b_idx = 256*i + QUANT_K_PQ2_0*ir + 4*il;
+    const uint bits = uint(data_a[ib].qs[il]);
+    const float d = float(data_a[ib].d);
+
+    data_b[b_idx    ] = D_TYPE(d * (float(bits & 3u) - 1.0f));
+    data_b[b_idx + 1] = D_TYPE(d * (float((bits >> 2u) & 3u) - 1.0f));
+    data_b[b_idx + 2] = D_TYPE(d * (float((bits >> 4u) & 3u) - 1.0f));
+    data_b[b_idx + 3] = D_TYPE(d * (float(bits >> 6u) - 1.0f));
+}

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

@@ -11,7 +11,7 @@
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-#if defined(DATA_A_Q2_0) || defined(DATA_A_QUANT_K)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_PQ2_0) || defined(DATA_A_QUANT_K)
 #define K_PER_ITER 16
 #elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
 #define K_PER_ITER 8

ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl

@@ -8,6 +8,10 @@
 FLOAT_TYPE get_dm(uint ib) {
     return FLOAT_TYPE(data_a[ib / 2].d);
 }
+#elif defined(DATA_A_PQ2_0)
+FLOAT_TYPE get_dm(uint ib) {
+    return FLOAT_TYPE(data_a[ib / 4].d);
+}
 #elif 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_dm(uint ib) {
     return FLOAT_TYPE(data_a[ib].d);
@@ -50,6 +54,27 @@ i32vec4 repack4(uint ib, uint iqs) {
                    unpack_q2_0(bits >> 16u), unpack_q2_0(bits >> 24u));
 }
 
+FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const int32_t sum_divisor) {
+    return FLOAT_TYPE(da * (float(q_sum) * dsb.x - dsb.y / float(sum_divisor)));
+}
+#endif
+// PQ2_0: group-128 variant of the Q2_0 path (4 x 32-element chunks per block)
+#if defined(DATA_A_PQ2_0)
+uint unpack_pq2_0(uint bits) {
+    // Move bit pairs [1:0], [3:2], [5:4], [7:6] to [1:0], [9:8], [17:16], [25:24].
+    bits &= 0xffu;
+    bits = (bits | (bits << 12u)) & 0x000f000fu;
+    return (bits | (bits << 6u)) & 0x03030303u;
+}
+
+i32vec4 repack4(uint ib, uint iqs) {
+    const uint qs_idx = (ib & 3u) * 4u + iqs * 2u;
+    const uint bits = pack32(u16vec2(data_a_packed16[ib / 4].qs[qs_idx],
+                                     data_a_packed16[ib / 4].qs[qs_idx + 1]));
+    return i32vec4(unpack_pq2_0(bits), unpack_pq2_0(bits >> 8u),
+                   unpack_pq2_0(bits >> 16u), unpack_pq2_0(bits >> 24u));
+}
+
 FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const int32_t sum_divisor) {
     return FLOAT_TYPE(da * (float(q_sum) * dsb.x - dsb.y / float(sum_divisor)));
 }
@@ -157,7 +182,7 @@ FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const i
 }
 #endif
 
-#if defined(DATA_A_Q2_0)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_PQ2_0)
 FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) {
     int32_t q_sum = 0;
     const i32vec4 qs_a = repack4(ib_a, iqs);