tuning q1_0 metal kernels

Author: khosravipasha

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -8,7 +8,7 @@
 //
 // TODO: for optimal performance, become function of the device and work size
 
-#define N_R0_Q1_0 4
+#define N_R0_Q1_0 8
 #define N_SG_Q1_0 2
 
 #define N_R0_Q4_0 4

ggml/src/ggml-metal/ggml-metal.metal

@@ -124,34 +124,29 @@ void dequantize_q1_0(device const block_q1_0 * xb, short il, thread type4x4 & re
     const float d = xb->d;
     const float neg_d = -d;
 
-    // Process 16 bits starting at offset il*16
-    // Optimization: process 2 bytes (16 bits) at once for better memory access
     const int byte_offset = il * 2;  // il*16 bits = il*2 bytes
     const uint8_t b0 = qs[byte_offset];
     const uint8_t b1 = qs[byte_offset + 1];
 
     float4x4 reg_f;
 
-    // Unroll completely for better ILP
-    // First byte (bits 0-7)
-    reg_f[0][0] = (b0 & 0x01) ? d : neg_d;
-    reg_f[0][1] = (b0 & 0x02) ? d : neg_d;
-    reg_f[0][2] = (b0 & 0x04) ? d : neg_d;
-    reg_f[0][3] = (b0 & 0x08) ? d : neg_d;
-    reg_f[1][0] = (b0 & 0x10) ? d : neg_d;
-    reg_f[1][1] = (b0 & 0x20) ? d : neg_d;
-    reg_f[1][2] = (b0 & 0x40) ? d : neg_d;
-    reg_f[1][3] = (b0 & 0x80) ? d : neg_d;
-
-    // Second byte (bits 8-15)
-    reg_f[2][0] = (b1 & 0x01) ? d : neg_d;
-    reg_f[2][1] = (b1 & 0x02) ? d : neg_d;
-    reg_f[2][2] = (b1 & 0x04) ? d : neg_d;
-    reg_f[2][3] = (b1 & 0x08) ? d : neg_d;
-    reg_f[3][0] = (b1 & 0x10) ? d : neg_d;
-    reg_f[3][1] = (b1 & 0x20) ? d : neg_d;
-    reg_f[3][2] = (b1 & 0x40) ? d : neg_d;
-    reg_f[3][3] = (b1 & 0x80) ? d : neg_d;
+    reg_f[0][0] = select(neg_d, d, bool(b0 & 0x01));
+    reg_f[0][1] = select(neg_d, d, bool(b0 & 0x02));
+    reg_f[0][2] = select(neg_d, d, bool(b0 & 0x04));
+    reg_f[0][3] = select(neg_d, d, bool(b0 & 0x08));
+    reg_f[1][0] = select(neg_d, d, bool(b0 & 0x10));
+    reg_f[1][1] = select(neg_d, d, bool(b0 & 0x20));
+    reg_f[1][2] = select(neg_d, d, bool(b0 & 0x40));
+    reg_f[1][3] = select(neg_d, d, bool(b0 & 0x80));
+
+    reg_f[2][0] = select(neg_d, d, bool(b1 & 0x01));
+    reg_f[2][1] = select(neg_d, d, bool(b1 & 0x02));
+    reg_f[2][2] = select(neg_d, d, bool(b1 & 0x04));
+    reg_f[2][3] = select(neg_d, d, bool(b1 & 0x08));
+    reg_f[3][0] = select(neg_d, d, bool(b1 & 0x10));
+    reg_f[3][1] = select(neg_d, d, bool(b1 & 0x20));
+    reg_f[3][2] = select(neg_d, d, bool(b1 & 0x40));
+    reg_f[3][3] = select(neg_d, d, bool(b1 & 0x80));
 
     reg = (type4x4) reg_f;
 }
@@ -160,20 +155,18 @@ template <typename type4>
 void dequantize_q1_0_t4(device const block_q1_0 * xb, short il, thread type4 & reg) {
     device const uint8_t * qs = xb->qs;
     const float d = xb->d;
+    const float neg_d = -d;
 
-    float4 reg_f;
-
-    // Process 4 bits for each call
-    const int offset = il * 4;
-
-    for (int i = 0; i < 4; i++) {
-        const int bit_idx = offset + i;
-        const int byte_idx = bit_idx / 8;
-        const int bit_offset = bit_idx % 8;
+    // 4 consecutive bits starting at il*4
+    const int base = il * 4;
+    const uint8_t byte = qs[base / 8];
+    const int bit_base = base % 8;
 
-        const bool bit_val = (qs[byte_idx] >> bit_offset) & 1;
-        reg_f[i] = bit_val ? d : -d;
-    }
+    float4 reg_f;
+    reg_f[0] = select(neg_d, d, bool((byte >> (bit_base    )) & 1));
+    reg_f[1] = select(neg_d, d, bool((byte >> (bit_base + 1)) & 1));
+    reg_f[2] = select(neg_d, d, bool((byte >> (bit_base + 2)) & 1));
+    reg_f[3] = select(neg_d, d, bool((byte >> (bit_base + 3)) & 1));
 
     reg = (type4) reg_f;
 }
@@ -3178,25 +3171,36 @@ kernel void kernel_group_norm_f32(
 
 // function for calculate inner product between part of a q1_0 block and 16 floats (yl), sumy is SUM(yl[i])
 // il indicates where the q1 quants begin (0, 16, 32, ..., 112 for 128-element block)
-// we assume that the yl's have been multiplied with the appropriate scale factor
+// Q1_0 encodes weights as {+1, -1}: dot = d * (2 * Σ(yl[i] where bit=1) - sumy)
 inline float block_q_n_dot_y(device const block_q1_0 * qb_curr, float sumy, thread float * yl, int il) {
-    float d = qb_curr->d;
+    // 16 elements = 2 bytes starting at byte offset il/8
+    device const uint8_t * qs = qb_curr->qs + il / 8;
+    const uint8_t b0 = qs[0];
+    const uint8_t b1 = qs[1];
 
+    // Accumulate yl[i] only where bit is 1, using select (branchless)
     float acc = 0.0f;
 
-    // il represents which 16-element chunk of the 128-element block (0, 16, 32, ..., 112)
-    // 16 weights = 16 bits = 2 bytes
-    const int byte_offset = il / 8;
-    device const uint8_t * qs = qb_curr->qs + byte_offset;
+    acc += select(0.0f, yl[ 0], bool(b0 & 0x01));
+    acc += select(0.0f, yl[ 1], bool(b0 & 0x02));
+    acc += select(0.0f, yl[ 2], bool(b0 & 0x04));
+    acc += select(0.0f, yl[ 3], bool(b0 & 0x08));
+    acc += select(0.0f, yl[ 4], bool(b0 & 0x10));
+    acc += select(0.0f, yl[ 5], bool(b0 & 0x20));
+    acc += select(0.0f, yl[ 6], bool(b0 & 0x40));
+    acc += select(0.0f, yl[ 7], bool(b0 & 0x80));
 
-    for (int i = 0; i < 16; i++) {
-        const uint8_t byte_idx = i / 8;
-        const uint8_t bit_idx = i % 8;
-        const int8_t qval = ((qs[byte_idx] >> bit_idx) & 1) ? 1 : -1;
-        acc += yl[i] * qval;
-    }
+    acc += select(0.0f, yl[ 8], bool(b1 & 0x01));
+    acc += select(0.0f, yl[ 9], bool(b1 & 0x02));
+    acc += select(0.0f, yl[10], bool(b1 & 0x04));
+    acc += select(0.0f, yl[11], bool(b1 & 0x08));
+    acc += select(0.0f, yl[12], bool(b1 & 0x10));
+    acc += select(0.0f, yl[13], bool(b1 & 0x20));
+    acc += select(0.0f, yl[14], bool(b1 & 0x40));
+    acc += select(0.0f, yl[15], bool(b1 & 0x80));
 
-    return d * acc;
+    // dot = d * (2 * sum_where_bit1 - sumy)
+    return qb_curr->d * (2.0f * acc - sumy);
 }
 
 // function for calculate inner product between half a q4_0 block and 16 floats (yl), sumy is SUM(yl[i])
@@ -10022,6 +10026,7 @@ template [[host_name("kernel_mul_mm_iq4_xs_f32")]]  kernel mul_mm_t kernel_mul_m
 
 template [[host_name("kernel_mul_mm_f32_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_f16_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
+template [[host_name("kernel_mul_mm_q1_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q4_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q4_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q5_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;