vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

metal : improve q4_K

Browse source 
28.3 -> 26.0 ms/token by avoiding a branch in the
calculation of the scales.
This commit is contained in:
Iwan Kawrakow 2023-06-10 11:31:12 +03:00
parent fa84c4b3e8
commit a75c12932d
1 changed files with 30 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

40
ggml-metal.metal
View file

@ -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;