(wip) Vulkan: Adreno Q4_1 fix

Author: Italo Nicola

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

@@ -47,8 +47,9 @@ vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     return vec2(vui & 0xF, vui >> 4);
 }
 vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
-    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
-    return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12);
+    const vec2 v01 = dequantize(ib, iqs, a_offset);
+    const vec2 v23 = dequantize(ib, iqs + 1, a_offset);
+    return vec4(v01.x, v01.y, v23.x, v23.y);
 }
 #endif
 

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

@@ -92,6 +92,37 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
 }
 #endif
 
+#if defined(DATA_A_Q4_1)
+#define BINDING_IDX_K 0
+#define BINDING_IDX_V 1
+layout (binding = 1) readonly buffer KV_PACKED16 {A_TYPE data[];} kv_packed[2];
+#define BLOCK_BYTE_SIZE 20
+
+vec4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
+    uint v00 =
+        uint(kv_packed[binding_idx].data[a_offset + ib].qs[(iqs & 0xF) + 0]);
+    uint v01 =
+        uint(kv_packed[binding_idx].data[a_offset + ib].qs[(iqs & 0xF) + 1]);
+    uint v10 =
+        uint(kv_packed[binding_idx].data[a_offset + ib].qs[(iqs & 0xF) + 2]);
+    uint v11 =
+        uint(kv_packed[binding_idx].data[a_offset + ib].qs[(iqs & 0xF) + 3]);
+
+    uint shift = (iqs & 0x10) >> 2;
+    v00 >>= shift;
+    v01 >>= shift;
+    v10 >>= shift;
+    v11 >>= shift;
+
+    v00 = v00 & 0xF;
+    v01 = v01 & 0xF;
+    v10 = v10 & 0xF;
+    v11 = v11 & 0xF;
+
+    return float(kv_packed[binding_idx].data[a_offset + ib].d) * (vec4(v00, v01, v10, v11) - 8.0f);
+}
+#endif
+
 #if defined(DATA_A_Q8_0)
 #define BINDING_IDX_K 0
 #define BINDING_IDX_V 1

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

@@ -10,7 +10,7 @@
 
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
 layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
-#if !defined(DATA_A_Q8_0) && !defined(DATA_A_Q4_0)
+#if !defined(DATA_A_Q8_0) && !defined(DATA_A_Q4_0) && !defined(DATA_A_Q4_1)
 layout (binding = 1) readonly buffer BV2 {B_TYPE_VEC2 data_b_v2[];};
 layout (binding = 1) readonly buffer BV4 {B_TYPE_VEC4 data_b_v4[];};
 #endif

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

@@ -140,6 +140,18 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
 }
 #endif
 
+#if defined(DATA_A_Q4_1)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+    return vec2(vui & 0xF, vui >> 4);
+}
+vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+    const vec2 v01 = dequantize(ib, iqs, a_offset);
+    const vec2 v23 = dequantize(ib, iqs + 1, a_offset);
+    return vec4(v01.x, v01.y, v23.x, v23.y);
+}
+#endif
+
 void main() {
 #ifdef NEEDS_INIT_IQ_SHMEM
     init_iq_shmem(gl_WorkGroupSize);
@@ -382,6 +394,7 @@ void main() {
             const uint ib = idx / 4;
             const uint iqs = idx & 0x03;
 
+#if 0
             const float d = float(data_a_packed16[ib].d);
             const float m = float(data_a_packed16[ib].m);
             const uint vui = uint(data_a_packed16[ib].qs[2*iqs]) | (uint(data_a_packed16[ib].qs[2*iqs + 1]) << 16);
@@ -396,6 +409,21 @@ void main() {
             buf_a[buf_idx + 17] = FLOAT_TYPE(v1.y);
             buf_a[buf_idx + 18] = FLOAT_TYPE(v1.z);
             buf_a[buf_idx + 19] = FLOAT_TYPE(v1.w);
+#else
+            const float d = float(data_a[ib].d);
+            const float m = float(data_a[ib].m);
+            const vec4 vxy = dequantize4(ib, 4*iqs, 0) * d + m;
+            const vec4 vzw = dequantize4(ib, 4*iqs + 2, 0) * d + m;
+
+            buf_a[buf_idx     ] = FLOAT_TYPE(vxy.x);
+            buf_a[buf_idx + 1 ] = FLOAT_TYPE(vxy.z);
+            buf_a[buf_idx + 2 ] = FLOAT_TYPE(vzw.x);
+            buf_a[buf_idx + 3 ] = FLOAT_TYPE(vzw.z);
+            buf_a[buf_idx + 16] = FLOAT_TYPE(vxy.y);
+            buf_a[buf_idx + 17] = FLOAT_TYPE(vxy.w);
+            buf_a[buf_idx + 18] = FLOAT_TYPE(vzw.y);
+            buf_a[buf_idx + 19] = FLOAT_TYPE(vzw.w);
+#endif
 #elif defined(DATA_A_Q5_0)
             const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
             const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + 2 * loadr_a;

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

@@ -24,7 +24,7 @@
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-#if defined(DATA_A_Q8_0) || defined(DATA_A_Q4_0)
+#if defined(DATA_A_Q8_0) || defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
 #else
 layout (binding = 0) readonly buffer A {A_TYPE_PACKED16 data_a[];};

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

@@ -43,10 +43,30 @@ ACC_TYPE mul_q8_1(int32_t q_sum, float da, vec2 dsb) {
 
 #if defined(DATA_A_Q4_1)
 i32vec2 repack(uint ib, uint iqs) {
+#if 0
     // Use 4-byte loads since a q4_1 block (20 bytes) is divisible by 4
     const uint32_t vui = data_a_packed32[ib].qs[iqs];
     return i32vec2( vui       & 0x0F0F0F0F,
                    (vui >> 4) & 0x0F0F0F0F);
+#else
+    int32_t u0 = int32_t(uint(data_a[ib].qs[iqs * 4]));
+    int32_t u1 = int32_t(uint(data_a[ib].qs[iqs * 4 + 1]));
+    int32_t u2 = int32_t(uint(data_a[ib].qs[iqs * 4 + 2]));
+    int32_t u3 = int32_t(uint(data_a[ib].qs[iqs * 4 + 3]));
+
+    int32_t v0 = int32_t(
+            (u0 & 0xF) |
+            ((u1 & 0xF) << 8) |
+            ((u2 & 0xF) << 16) |
+            ((u3 & 0xF) << 24));
+    int32_t v1 = int32_t(
+            ((u0 >> 4) & 0xF) |
+            (((u1 >> 4) & 0xF) << 8) |
+            (((u2 >> 4) & 0xF) << 16) |
+            (((u3 >> 4) & 0xF) << 24));
+    
+    return i32vec2(v0, v1);
+#endif
 }
 
 ACC_TYPE mul_q8_1(int32_t q_sum, vec2 dma, vec2 dsb) {
@@ -123,8 +143,15 @@ FLOAT_TYPE get_d(uint ib) {
 }
 #endif
 
-#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1)
+#if defined(DATA_A_Q4_1)
+FLOAT_TYPE_VEC2 get_dm(uint ib) {
+    return FLOAT_TYPE_VEC2(data_a[ib].d, data_a[ib].m);
+}
+#endif
+
+#if defined(DATA_A_Q5_1)
 FLOAT_TYPE_VEC2 get_dm(uint ib) {
     return FLOAT_TYPE_VEC2(data_a_packed32[ib].dm);
 }
 #endif
+

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

@@ -96,8 +96,8 @@ struct block_q4_1_packed32
 #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
-#define A_TYPE_PACKED32 block_q4_1_packed32
+//#define A_TYPE_PACKED16 block_q4_1_packed16
+//#define A_TYPE_PACKED32 block_q4_1_packed32
 #endif
 
 #define QUANT_K_Q5_0 32

tests/test-backend-ops.cpp

@@ -4922,7 +4922,7 @@ struct test_falcon : public test_llm {
 static const ggml_type all_types[] = {
     GGML_TYPE_F32, GGML_TYPE_F16, // GGML_TYPE_BF16,
     GGML_TYPE_Q4_0,
-//    GGML_TYPE_Q4_1,
+    GGML_TYPE_Q4_1,
 //    GGML_TYPE_Q5_0, GGML_TYPE_Q5_1,
     GGML_TYPE_Q8_0,
 //    GGML_TYPE_Q2_K, GGML_TYPE_Q3_K,
@@ -4938,14 +4938,14 @@ static const ggml_type base_types[] = {
     GGML_TYPE_F32, GGML_TYPE_F16,
     GGML_TYPE_Q8_0, // for I8MM tests
     GGML_TYPE_Q4_0,
-//    GGML_TYPE_Q4_1, // for I8MM tests
+    GGML_TYPE_Q4_1, // for I8MM tests
 //    GGML_TYPE_Q4_K,
 //    GGML_TYPE_IQ2_XXS
 };
 
 static const ggml_type other_types[] = {
     GGML_TYPE_Q4_0,
-//    GGML_TYPE_Q4_1,
+    GGML_TYPE_Q4_1,
 //    GGML_TYPE_Q5_0, GGML_TYPE_Q5_1,
     GGML_TYPE_Q8_0,
 //    GGML_TYPE_Q2_K, GGML_TYPE_Q3_K,