metal : improve q4_K

28.3 -> 26.0 ms/token by avoiding a branch in the
calculation of the scales.

Author: Kawrakow

ggml-metal.metal

@@ -969,6 +969,10 @@ kernel void kernel_mul_mat_q4_k_f32(
         uint2 tpitg[[thread_position_in_threadgroup]],
         uint2  tptg[[threads_per_threadgroup]]) {
 
+    const uint16_t kmask1 = 0x3f3f;
+    const uint16_t kmask2 = 0x0f0f;
+    const uint16_t kmask3 = 0x0303;
+
     const int nb = ne00/QK_K;
 
     const int64_t r0 = tgpig.x;
@@ -983,29 +987,45 @@ kernel void kernel_mul_mat_q4_k_f32(
     const int tid = tpitg.y;   // 0...16
     const int il  = tid/4;     // 0...3
     const int ir  = tid%4;     // 0...3
-    const int n   = 8;
-    const int is  = 2*il;
+    const int n   = 4;
+
+    const int im = il/2;  // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224
+    const int in = il%2;
 
     sum[ith] = 0.0f;
 
+    //uchar2 sc1, sc2;
+
     float sumf = 0;
     for (int i = tpitg.x; i < nb; i += tptg.x) {
 
-        device const uint8_t * q = (x + i)->qs + 32*il + n*ir;
-        device const float   * y = yy + i*QK_K + 64*il + n*ir;
-        device const uint8_t * scales = (x + i)->scales;
+        device const uint8_t * q1 = (x + i)->qs + 32*im + n*(2*ir + in);
+        device const float   * y1 = yy + i*QK_K + 64*im + n*(2*ir + in);
+        device const uint8_t * q2 = q1 + 64;
+        device const float   * y2 = y1 + 128;
+
+        device const uint16_t * a = (device const uint16_t *)(x + i)->scales;
 
         const float dall = (float)((x + i)->d);
         const float dmin = (float)((x + i)->dmin);
 
-        const uchar4 sc = get_scale_min_k4(is, scales);
+        const uchar2 sc1 = as_type<uchar2>((uint16_t)(a[im+0] & kmask1));
+        const uchar2 sc2 = as_type<uchar2>((uint16_t)(a[im+2] & kmask1));
+        const uchar2 sc3 = as_type<uchar2>((uint16_t)(((a[im+4] >> 0) & kmask2) | (((a[im+0] >> 6) & kmask3) << 4)));
+        const uchar2 sc4 = as_type<uchar2>((uint16_t)(((a[im+4] >> 4) & kmask2) | (((a[im+2] >> 6) & kmask3) << 4)));
 
-        float4 s = {0.f, 0.f, 0.f, 0.f};
+        float4 s1 = {0.f, 0.f, 0.f, 0.f};
+        float4 s2 = {0.f, 0.f, 0.f, 0.f};
         for (int l = 0; l < n; ++l) {
-            s[0] += y[l+ 0] * (q[l] & 0xF); s[1] += y[l+ 0];
-            s[2] += y[l+32] * (q[l] >>  4); s[3] += y[l+32];
+            s1[0] += y1[l+ 0] * (q1[l] & 0xF); s1[1] += y1[l+ 0];
+            s1[2] += y1[l+32] * (q1[l] >>  4); s1[3] += y1[l+32];
+            s2[0] += y2[l+ 0] * (q2[l] & 0xF); s2[1] += y2[l+ 0];
+            s2[2] += y2[l+32] * (q2[l] >>  4); s2[3] += y2[l+32];
         }
-        sumf += dall * (s[0] * sc[0] + s[2] * sc[2]) - dmin * (s[1] * sc[1] + s[3] * sc[3]);
+        sumf += dall * (s1[0] * sc1[0] + s1[2] * sc1[1]
+                     +  s2[0] * sc3[0] + s2[2] * sc3[1])
+              - dmin * (s1[1] * sc2[0] + s1[3] * sc2[1]
+                     +  s2[1] * sc4[0] + s2[3] * sc4[1]);
 
     }
     sum[ith] = sumf;