implement matrix kernel; wait for mlx bug fix

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

book/src/week2-overview.md

@@ -2,6 +2,8 @@ https://github.com/ml-explore/mlx/blob/main/mlx/backend/cpu/quantized.cpp
 https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py
 MLX uses INT4 W4A16
 https://ml-explore.github.io/mlx/build/html/dev/extensions.html
+https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-metal/ggml-metal.metal
+https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/kernels/quantized.h#L962
 
 pdm run ./build_ext.sh
 

build_ext.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+set -e
+pdm run build-ext-ref
+cp src/extensions_ref/build/lib/tiny_llm_ext_ref/tiny_llm_ext_ref.metallib .venv/lib/python3.12/site-packages/mlx/lib/
+pdm run test-week2-ref -k 'week_2_day_2'

src/extensions_ref/CMakeLists.txt

@@ -58,6 +58,7 @@ if(MLX_BUILD_METAL)
     tiny_llm_ext_ref
     SOURCES
     ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
+    ${CMAKE_CURRENT_LIST_DIR}/src/quantized_matmul.metal
     INCLUDE_DIRS
     ${PROJECT_SOURCE_DIR}
     ${MLX_INCLUDE_DIRS}

src/extensions_ref/bindings.cpp

@@ -12,6 +12,8 @@ using namespace nb::literals;
 NB_MODULE(_ext, m) {
     m.doc() = "tiny-llm extensions for MLX";
 
+    m.def("load_library", &tiny_llm_ext_ref::load_library, "device"_a, "path"_a);
+
     m.def("axpby", &tiny_llm_ext_ref::axpby, "x"_a, "y"_a, "alpha"_a, "beta"_a, nb::kw_only(), "stream"_a = nb::none(),
           R"(
         Scale and sum two vectors element-wise

src/extensions_ref/build.py

@@ -2,24 +2,23 @@
 import shutil
 from mlx import extension
 from setuptools import Distribution
+import inspect
+import mlx
+import os
 
 if __name__ == "__main__":
     src_dir = Path(__file__).parent
     distribution = Distribution(
         {
             "name": "tiny_llm_ext_ref",
             "ext_modules": [extension.CMakeExtension("tiny_llm_ext_ref._ext")],
+            "package_data": {"tiny_llm_ext_ref": ["*.so", "*.dylib", "*.metallib"]},
         }
     )
     cmd = extension.CMakeBuild(distribution)
     cmd.initialize_options()
     cmd.build_temp = Path("build")
     cmd.build_lib = Path("build") / "lib"
-    cmd.inplace = False  # we do the copy by ourselves
+    cmd.inplace = True
     cmd.ensure_finalized()
     cmd.run()
-    for output in cmd.get_outputs():
-        output = Path(output)
-        relative_extension = src_dir / output.relative_to(cmd.build_lib)
-        shutil.copyfile(output, relative_extension)
-        print(f"Copied {output} to {relative_extension}")

src/extensions_ref/src/quantized_matmul.cpp

@@ -54,6 +54,17 @@ mx::array quantized_matmul(const mx::array &scales,         // Input array scale
     if (!transpose_b) {
         throw std::runtime_error("quantized_matmul: b must be transposed");
     }
+
+    if (scales.shape() != biases.shape()) {
+        throw std::runtime_error("quantized_matmul: scales and biases must have the same shape");
+    }
+    if (b.shape()[0] != scales.shape()[0]) {
+        throw std::runtime_error("quantized_matmul: b must have the same number of rows as scales");
+    }
+    if (b.shape()[1] != scales.shape()[1] * group_size / 8) {
+        throw std::runtime_error("quantized_matmul: a must have the same number of columns as scales");
+    }
+    
     return mx::array(
         /* const mx::Shape& shape = */ out_shape,
         /* mx::Dtype dtype = */ mx::float16,
@@ -73,14 +84,11 @@ void quantized_matmul_impl(const mx::array &scales, const mx::array &biases, con
     encoder.set_input_array(b);
     encoder.set_output_array(out);
 
-    if (scales.shape() != biases.shape()) {
-        throw std::runtime_error("quantized_matmul: scales and biases must have the same shape");
+    if (!a.flags().row_contiguous) {
+        throw std::runtime_error("quantized_matmul: a must be contiguous");
     }
-    if (b.shape()[0] != scales.shape()[0]) {
-        throw std::runtime_error("quantized_matmul: b must have the same number of rows as scales");
-    }
-    if (b.shape()[1] != scales.shape()[1] * group_size / 8) {
-        throw std::runtime_error("quantized_matmul: a must have the same number of columns as scales");
+    if (!b.flags().row_contiguous) {
+        throw std::runtime_error("quantized_matmul: b must be contiguous");
     }
 
     // Launch the CPU kernel
@@ -100,32 +108,32 @@ void quantized_matmul_impl(const mx::array &scales, const mx::array &biases, con
         uint32_t item_mask = (1 << bits) - 1;
         for (int i = 0; i < M; i++) {
             for (int k = 0; k < K; k++) {
+                float sum = 0;
                 for (int group_idx = 0; group_idx < group_per_row; group_idx++) {
                     int64_t scales_loc =
-                        mx::elem_to_loc(k * N / group_size + group_idx, scales.shape(), scales.strides());
+                        mx::elem_to_loc(k * group_per_row + group_idx, scales.shape(), scales.strides());
                     int64_t biases_loc =
-                        mx::elem_to_loc(k * N / group_size + group_idx, biases.shape(), biases.strides());
-                    float16_t sum = 0;
+                        mx::elem_to_loc(k * group_per_row + group_idx, biases.shape(), biases.strides());
                     float16_t scale = scales_ptr[scales_loc];
                     float16_t bias = biases_ptr[biases_loc];
+                    int64_t b_loc = mx::elem_to_loc((k * N + group_idx * group_size) / 8, b.shape(), b.strides());
+                    int64_t a_loc = mx::elem_to_loc(i * N + group_idx * group_size, a.shape(), a.strides());
                     const int packs_per_item = 32 / bits;
                     for (int item_idx = 0; item_idx < group_size; item_idx += packs_per_item) {
-                        int64_t b_loc =
-                            mx::elem_to_loc((k * N + group_idx * group_size + item_idx) / 8, b.shape(), b.strides());
                         uint32_t b_val = b_ptr[b_loc];
                         uint8_t *b_bytes = reinterpret_cast<uint8_t *>(&b_val);
                         for (int pack_idx = 0; pack_idx < packs_per_item; pack_idx++) {
-                            int64_t a_loc = mx::elem_to_loc(i * N + group_idx * group_size + item_idx + pack_idx,
-                                                            a.shape(), a.strides());
                             uint8_t item_val = (b_bytes[pack_idx / 2] >> ((pack_idx % 2) * bits)) & item_mask;
-                            float16_t b = static_cast<float16_t>(item_val) * scale + bias;
-                            float16_t a = a_ptr[a_loc];
+                            float b = static_cast<float>(item_val) * scale + bias;
+                            float a = a_ptr[a_loc];
                             sum += a * b;
+                            a_loc += 1;
                         }
+                        b_loc += 1;
                     }
-                    int64_t out_loc = mx::elem_to_loc(i * K + k, out_shape, out_strides);
-                    out_ptr[out_loc] = sum;
                 }
+                int64_t out_loc = mx::elem_to_loc(i * K + k, out_shape, out_strides);
+                out_ptr[out_loc] = static_cast<float16_t>(sum);
             }
         }
     });
@@ -142,8 +150,65 @@ void QuantizedMatmul::eval_cpu(const std::vector<mx::array> &inputs, std::vector
     quantized_matmul_impl(scales, biases, a, b, out, group_size_, bits_, stream());
 }
 
-void QuantizedMatmul::eval_gpu(const std::vector<mx::array> &inputs, std::vector<mx::array> &out) {
-    throw std::runtime_error("QuantizedMatmul has no GPU implementation.");
+void load_library(mx::Device d, const char* path) {
+    auto &md = mx::metal::device(d);
+    md.register_library("tiny_llm_ext_ref", path);
+}
+
+void QuantizedMatmul::eval_gpu(const std::vector<mx::array> &inputs, std::vector<mx::array> &outputs) {
+    auto &scales = inputs[0];
+    auto &biases = inputs[1];
+    auto &a = inputs[2];
+    auto &b = inputs[3];
+    auto &out = outputs[0];
+
+    auto &s = stream();
+    auto &d = mx::metal::device(s.device);
+    out.set_data(mx::allocator::malloc(out.nbytes()));
+
+    // Make a kernel from this metal library
+    auto kernel = d.get_kernel("quantized_matmul_w4a16_g64", "tiny_llm_ext_ref");
+
+    // Prepare to encode kernel
+    auto &compute_encoder = d.get_command_encoder(s.index);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    // Kernel parameters are registered with buffer indices corresponding to
+    // those in the kernel declaration at axpby.metal
+    int ndim = out.ndim();
+
+    // Encode input arrays to kernel
+    compute_encoder.set_input_array(scales, 0);
+    compute_encoder.set_input_array(biases, 1);
+    compute_encoder.set_input_array(a, 2);
+    compute_encoder.set_input_array(b, 3);
+    // Encode output arrays to kernel
+    compute_encoder.set_output_array(out, 4);
+
+
+    if (!a.flags().row_contiguous) {
+        throw std::runtime_error("quantized_matmul: a must be contiguous");
+    }
+    if (!b.flags().row_contiguous) {
+        throw std::runtime_error("quantized_matmul: b must be contiguous");
+    }
+
+    int M = a.shape()[0];
+    int N = a.shape()[1];
+    int K = b.shape()[0];
+
+    // Encode matrix parameters
+    compute_encoder.set_bytes(M, 5);
+    compute_encoder.set_bytes(N, 6);
+    compute_encoder.set_bytes(K, 7);
+
+    size_t tgp_size = kernel->maxTotalThreadsPerThreadgroup();
+    MTL::Size num_threadgroups = MTL::Size((M * K + tgp_size - 1) / tgp_size, 1, 1);
+    MTL::Size num_threads_per_group = MTL::Size(tgp_size, 1, 1);
+
+    // Launch the grid with the given number of threads divided among
+    // the given threadgroups
+    compute_encoder.dispatch_threadgroups(num_threadgroups, num_threads_per_group);
 }
 
 bool QuantizedMatmul::is_equivalent(const Primitive &other) const {

src/extensions_ref/src/quantized_matmul.metal

@@ -0,0 +1,46 @@
+[[kernel]] void quantized_matmul_w4a16_g64(
+    device const half* scales [[buffer(0)]],
+    device const half* biases [[buffer(1)]],
+    device const half* a [[buffer(2)]],
+    device const uint32_t* b [[buffer(3)]],
+    device half* out [[buffer(4)]],
+    device const int &M [[buffer(5)]],
+    device const int &N [[buffer(6)]],
+    device const int &K [[buffer(7)]],
+    uint2 groupId [[threadgroup_position_in_grid]],
+    uint2 threadId [[thread_position_in_threadgroup]],
+    uint2 threads_per_threadgroup [[threads_per_threadgroup]]) {
+    const int group_size = 64;
+    const int bits = 4;
+    const int packs_per_item = 32 / bits;
+    const int item_mask = (1 << bits) - 1;
+    const int groups_per_row = N / group_size;
+    // Each threadgroup processes an element in the output matrix
+    const int64_t idx = groupId.x * threads_per_threadgroup.x + threadId.x;
+    const int64_t i = idx / K;
+    const int64_t k = idx % K;
+    float sum = 0;
+    for (int group_idx = 0; group_idx < groups_per_row; group_idx++) {
+        const int64_t scales_biases_loc = k * groups_per_row + group_idx;
+        const float scale = scales[scales_biases_loc];
+        const float bias = biases[scales_biases_loc];
+        int64_t b_loc = (k * N + group_idx * group_size) / 8;
+        int64_t a_loc = i * N + group_idx * group_size;
+        for (int item_idx = 0; item_idx < group_size; item_idx += packs_per_item) {
+            const uint32_t b_val = b[b_loc];
+            thread const uint32_t *b_val_ref = &b_val;
+            thread const uint8_t *b_bytes = reinterpret_cast<thread const uint8_t *>(b_val_ref);
+            for (int pack_idx = 0; pack_idx < packs_per_item; pack_idx++) {
+                const uint8_t item_val = (b_bytes[pack_idx / 2] >> ((pack_idx % 2) * bits)) & item_mask;
+                const float b_val = static_cast<float>(item_val) * scale + bias;
+                const float a_val = a[a_loc];
+                sum += a_val * b_val;
+                a_loc += 1;
+            }
+            b_loc += 1;
+        }
+    }
+    if (i < M && k < K) {
+        out[i * K + k] = sum;
+    }
+}

src/extensions_ref/src/tiny_llm_ext.h

@@ -7,6 +7,8 @@ namespace mx = mlx::core;
 
 namespace tiny_llm_ext_ref {
 
+void load_library(mx::Device d, const char* path);
+
 mx::array quantized_matmul(const mx::array &scales,   // Input array scales
                            const mx::array &biases,   // Input array biases
                            const int group_size,      // Group size

src/extensions_ref/test.py

@@ -0,0 +1,18 @@
+from tiny_llm_ext_ref import quantized_matmul
+import mlx.core as mx
+import numpy as np
+
+precision = np.float16
+input = mx.array(np.random.randn(3, 64).astype(precision))
+weight = mx.array(np.random.randn(5, 64).astype(precision))
+w_q, scales, biases = mx.quantize(weight)
+user_out = quantized_matmul(
+    scales=scales,
+    biases=biases,
+    group_size=64,
+    bits=4,
+    a=input,
+    b=w_q,
+    transpose_b=True,
+)
+print(user_out)

src/extensions_ref/tiny_llm_ext_ref/__init__.py

@@ -3,3 +3,7 @@
 import mlx.core as mx
 
 from ._ext import *
+from pathlib import Path
+
+current_path = Path(__file__).parent
+load_library(mx.gpu, str(current_path))

src/tiny_llm_week2_ref/quantize.py

@@ -26,6 +26,8 @@ def quantized_matmul(
 ) -> mx.array:
     *N, D = a.shape
     a = a.reshape(-1, D)
+    a = mx.contiguous(a)
+    b = mx.contiguous(b)
     return tiny_llm_ext_ref.quantized_matmul(
         scales, biases, group_size, bits, a, b, transpose_b
     ).reshape(*N, -1)

tests/test_week_2_day_2.py

@@ -42,3 +42,11 @@ def test_task_1_quantized_matmul_simple_f16_cpu():
 
 def test_task_1_quantized_matmul_complex_f16_cpu():
     quantized_matmul_helper(mx.cpu, False, np.float16)
+
+
+def test_task_2_quantized_matmul_simple_f16_gpu():
+    quantized_matmul_helper(mx.gpu, True, np.float16)
+
+
+def test_task_2_quantized_matmul_complex_f16_gpu():
+    quantized_matmul_helper(mx.gpu, False, np.float16)