[MLX] GGUF vectorized loads (pytorch#20675)

Small improvement to do vectorized loads in GGUF kernels, matching
llama.cpp.

Author: metascroy

backends/mlx/custom_kernel_ops/gguf/q4k/common.py

@@ -64,6 +64,14 @@
 #include <metal_simdgroup_matrix>
 using namespace metal;
 
+#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
+
+typedef matrix<bfloat, 2, 4> bfloat2x4;
+template<typename T> struct vec2x4;
+template<> struct vec2x4<float>   { using type = float2x4;  };
+template<> struct vec2x4<half>    { using type = half2x4;   };
+template<> struct vec2x4<bfloat>  { using type = bfloat2x4; };
+
 #define QK_K 256
 #define K_SCALE_SIZE 12
 

backends/mlx/custom_kernel_ops/gguf/q4k/linear.py

@@ -105,7 +105,7 @@ def _q4k_matvec_source(has_bias: bool) -> str:
 
             float4 acc1 = {{0.f, 0.f, 0.f, 0.f}};
             float4 acc2 = {{0.f, 0.f, 0.f, 0.f}};
-            for (short i = 0; i < 4; ++i) {{
+            FOR_UNROLL (short i = 0; i < 4; ++i) {{
                 acc1[0] += yl[2*i + 0] * (float)(q1[i] & 0x000F);
                 acc1[1] += yl[2*i + 1] * (float)(q1[i] & 0x0F00);
                 acc1[2] += yl[2*i + 8] * (float)(q1[i] & 0x00F0);
@@ -218,9 +218,9 @@ def _q4k_matmul_source(has_bias: bool) -> str:
         const short ly_b = (tid / NL1) % 8;
         const short ib_b = 4 * sx_b + sy_b;
 
-        for (short i = 0; i < 8; ++i) {{
-            *(sb + 64 * ib_b + 8 * ly_b + i) = (half) *(yp + i);
-        }}
+        using InT2x4 = typename vec2x4<InT>::type;
+        *(threadgroup half2x4 *)(sb + 64 * ib_b + 8 * ly_b) =
+            (half2x4)(*(device const InT2x4 *)(yp));
 
         // Advance weight pointer through Q4_K sub-blocks.
         il = (il + 2 < NL) ? il + 2 : il % 2;

backends/mlx/custom_kernel_ops/gguf/q5k/common.py

@@ -70,6 +70,14 @@
 #include <metal_simdgroup_matrix>
 using namespace metal;
 
+#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
+
+typedef matrix<bfloat, 2, 4> bfloat2x4;
+template<typename T> struct vec2x4;
+template<> struct vec2x4<float>   { using type = float2x4;  };
+template<> struct vec2x4<half>    { using type = half2x4;   };
+template<> struct vec2x4<bfloat>  { using type = bfloat2x4; };
+
 #define QK_K 256
 #define K_SCALE_SIZE 12
 

backends/mlx/custom_kernel_ops/gguf/q5k/linear.py

@@ -135,7 +135,7 @@ def _q5k_matvec_source(has_bias: bool) -> str:
 
             float4 acc1 = {{0.f, 0.f, 0.f, 0.f}};
             float4 acc2 = {{0.f, 0.f, 0.f, 0.f}};
-            for (short l = 0; l < 8; ++l) {{
+            FOR_UNROLL (short l = 0; l < 8; ++l) {{
                 const uint8_t h = qh[l];
                 acc1[0] += yl[l + 0] * (float)(q1[l] & 0x0F);
                 acc1[1] += yl[l + 8] * (float)(q1[l] & 0xF0);
@@ -253,9 +253,9 @@ def _q5k_matmul_source(has_bias: bool) -> str:
         const short ly_b = (tid / NL1) % 8;
         const short ib_b = 4 * sx_b + sy_b;
 
-        for (short i = 0; i < 8; ++i) {{
-            *(sb + 64 * ib_b + 8 * ly_b + i) = (half) *(yp + i);
-        }}
+        using InT2x4 = typename vec2x4<InT>::type;
+        *(threadgroup half2x4 *)(sb + 64 * ib_b + 8 * ly_b) =
+            (half2x4)(*(device const InT2x4 *)(yp));
 
         // Advance weight pointer through Q5_K sub-blocks.
         il = (il + 2 < NL) ? il + 2 : il % 2;

backends/mlx/custom_kernel_ops/gguf/q6k/common.py

@@ -65,6 +65,14 @@
 #include <metal_simdgroup_matrix>
 using namespace metal;
 
+#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
+
+typedef matrix<bfloat, 2, 4> bfloat2x4;
+template<typename T> struct vec2x4;
+template<> struct vec2x4<float>   { using type = float2x4;  };
+template<> struct vec2x4<half>    { using type = half2x4;   };
+template<> struct vec2x4<bfloat>  { using type = bfloat2x4; };
+
 #define QK_K 256
 
 typedef struct {

backends/mlx/custom_kernel_ops/gguf/q6k/linear.py

@@ -126,7 +126,7 @@ def _q6k_matvec_source(has_bias: bool) -> str:
             const float d = (float) blk->d;
 
             float4 sums = {{0.f, 0.f, 0.f, 0.f}};
-            for (short l = 0; l < 4; ++l) {{
+            FOR_UNROLL (short l = 0; l < 4; ++l) {{
                 sums[0] += yl[4*l + 0] * (float)((int8_t)((q1[l] & 0xF) | ((qh[l] & 0x03) << 4)) - 32);
                 sums[1] += yl[4*l + 1] * (float)((int8_t)((q2[l] & 0xF) | ((qh[l] & 0x0C) << 2)) - 32);
                 sums[2] += yl[4*l + 2] * (float)((int8_t)((q1[l] >> 4)  | ((qh[l] & 0x30) << 0)) - 32);
@@ -228,9 +228,9 @@ def _q6k_matmul_source(has_bias: bool) -> str:
         const short ly_b = (tid / NL1) % 8;
         const short ib_b = 4 * sx_b + sy_b;
 
-        for (short i = 0; i < 8; ++i) {{
-            *(sb + 64 * ib_b + 8 * ly_b + i) = (half) *(yp + i);
-        }}
+        using InT2x4 = typename vec2x4<InT>::type;
+        *(threadgroup half2x4 *)(sb + 64 * ib_b + 8 * ly_b) =
+            (half2x4)(*(device const InT2x4 *)(yp));
 
         // Advance weight pointer through Q6_K sub-blocks.
         il = (il + 2 < NL) ? il + 2 : il % 2;