diff --git a/ggml-metal.metal b/ggml-metal.metal index ccd36386b..0909c7b30 100644 --- a/ggml-metal.metal +++ b/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((uint16_t)(a[im+0] & kmask1)); + const uchar2 sc2 = as_type((uint16_t)(a[im+2] & kmask1)); + const uchar2 sc3 = as_type((uint16_t)(((a[im+4] >> 0) & kmask2) | (((a[im+0] >> 6) & kmask3) << 4))); + const uchar2 sc4 = as_type((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;