Adding metal kernels for the gated delta nets.

benchmarks/python/gated_delta_bench.py

@@ -0,0 +1,145 @@
+import argparse
+import csv
+import itertools
+import os
+import time
+from datetime import datetime
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import numpy as np
+
+RED_BOLD = "\033[1;31m"
+GREEN = "\033[0;32m"
+RESET = "\033[0m"
+
+
+N_warmup = 8
+N_iter_bench = 80
+N_iter_func = 5
+
+
+# similar to ./blas/bench_gemm.py
+def bench(f, *args):
+    for _ in range(N_warmup):
+        f(*args)
+    mx.synchronize()
+
+    s = time.perf_counter_ns()
+    for _ in range(N_iter_bench):
+        f(*args)
+    mx.synchronize()
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9  # total seconds for N_iter_bench * N_iter_func calls
+
+
+def do_kernel_bench(f, *args):
+    ys = []
+    for _ in range(N_iter_func):
+        out, hf = f(*args)
+        ys.append(out)
+        ys.append(hf)
+    mx.eval(ys)
+    return ys
+
+
+def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes):
+    mx.random.seed(42)
+    q = mx.random.normal(shape=(B, T, Hk, Dk))
+    k = mx.random.normal(shape=(B, T, Hk, Dk))
+    k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6)
+    v = mx.random.normal(shape=(B, T, Hv, Dv))
+    g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0
+    b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv)))
+
+    shape_str = f"B={B} T={T} Hk={Hk} Hv={Hv} Dk={Dk} Dv={Dv}"
+    denom = N_iter_bench * N_iter_func
+
+    os.environ["GATED_DELTA_CHUNK"] = "0"
+    h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+    mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0))
+    ms_seq = (
+        bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0)
+        / denom
+        * 1e3
+    )
+
+    speedups = []
+    for C in (c for c in chunk_sizes if c != 0):
+        try:
+            os.environ["GATED_DELTA_CHUNK"] = str(C)
+            h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+            mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0))
+            ms_c = (
+                bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0)
+                / denom
+                * 1e3
+            )
+            speedups.append(ms_seq / ms_c if ms_c > 0 else float("nan"))
+        except Exception as ex:
+            print(f"  chunk {C} failed: {ex}")
+            speedups.append(float("nan"))
+
+    return shape_str, f"{ms_seq:.3f}", speedups, ms_seq
+
+
+def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"):
+    if run_full:
+        Bs = [1, 4, 8, 16]
+        Ts = [8, 64, 256, 512, 1024, 2048, 4096]
+        Hks = [16]
+        Hvs = [32]
+        Dks = [128]
+        Dvs = [128]
+    else:
+        Bs = [1, 8, 16]
+        Ts = [8, 512, 1024, 2048]
+        Hks = [16]
+        Hvs = [32]
+        Dks = [128]
+        Dvs = [128]
+
+    chunk_sizes = [0, 8, 16]
+    non_zero_Cs = [C for C in chunk_sizes if C != 0]
+
+    headers = ["B", "T", "Hk", "Hv", "Dk", "Dv", "time_seq (ms)"] + [
+        f"C={C} (speedup)" for C in non_zero_Cs
+    ]
+
+    col_widths = [6, 6, 6, 6, 6, 6, 15] + [25] * (len(non_zero_Cs))
+    fmt = "".join(f"{{:<{w}}}" for w in col_widths)
+
+    rows = []
+
+    print(fmt.format(*headers))
+    print("-" * (sum(col_widths)))
+
+    for B, T, Hk, Hv, Dk, Dv in itertools.product(Bs, Ts, Hks, Hvs, Dks, Dvs):
+        shapes_s, base_time_s, speedups, base_time = benchmark_shape(
+            B, T, Hk, Hv, Dk, Dv, chunk_sizes
+        )
+        row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}", base_time_s]
+        for speed in speedups:
+            row.append(f"{(base_time / speed):<8.2f} ({speed:<5.2f}x)")
+
+        print(fmt.format(*row), end="")
+        print(f"{RESET}")
+
+        rows.append(row)
+
+    if to_csv:
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow(headers)
+            writer.writerows(rows)
+        print(f"\nResults also written to {csv_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Gated delta benchmark")
+    parser.add_argument("--full", "-f", action="store_true")
+    parser.add_argument("--csv", "-c", action="store_true")
+    parser.add_argument("--csv_out", "-co", default="benchmark_results.csv")
+    args = parser.parse_args()
+
+    run_benchmark(args.full, to_csv=args.csv, csv_path=args.csv_out)

mlx/backend/cuda/primitives.cpp

@@ -24,6 +24,16 @@ namespace mlx::core {
     throw std::runtime_error(#func " has no CUDA implementation.");   \
   }
 
+bool fast::GatedDeltaUpdate::use_fallback(
+    const int Hk,
+    const int Dk,
+    const int Hv,
+    const int Dv,
+    const bool has_mask,
+    Stream s) {
+  return true;
+}
+
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
 NO_GPU_MULTI(SVD)
@@ -32,6 +42,10 @@ NO_GPU(Cholesky)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
 
+namespace fast {
+NO_GPU_MULTI(GatedDeltaUpdate)
+}
+
 namespace distributed {
 NO_GPU_MULTI(Send)
 NO_GPU_MULTI(Recv)

mlx/backend/metal/CMakeLists.txt

@@ -102,6 +102,7 @@ if(MLX_METAL_JIT)
                     kernels/fp4.h)
 
     make_jit_source(steel/attn/kernels/steel_attention_nax)
+    make_jit_source(gated_delta_update_nax)
 
   else()
     message(
@@ -136,6 +137,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gated_delta_update.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp

mlx/backend/metal/gated_delta_update.cpp

@@ -0,0 +1,197 @@
+// Copyright © 2024 Apple Inc.
+#include <sstream>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+bool GatedDeltaUpdate::use_fallback(
+    const int Hk,
+    const int Dk,
+    const int Hv,
+    const int Dv,
+    const bool has_mask,
+    Stream s) {
+  if (s.device == Device::cpu) {
+    return true;
+  }
+
+  if (has_mask) {
+    return true;
+  }
+
+  if (Dk != 128 || Dv != 128) {
+    return true;
+  }
+
+  const bool supported_heads = (Hk == 24 && Hv == 24) ||
+      (Hk == 32 && Hv == 32) || (Hk == 16 && Hv == 32) ||
+      (Hk == 16 && Hv == 48) || (Hk == 16 && Hv == 16) ||
+      (Hk == 16 && Hv == 64);
+  if (!supported_heads) {
+    return true;
+  }
+
+  return false;
+}
+
+inline array
+ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    metal::get_command_encoder(s).add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+void GatedDeltaUpdate::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto q = ensure_row_contiguous(inputs[0], d, s);
+  auto k = ensure_row_contiguous(inputs[1], d, s);
+  auto v = ensure_row_contiguous(inputs[2], d, s);
+  auto g = ensure_row_contiguous(inputs[3], d, s);
+  auto beta = ensure_row_contiguous(inputs[4], d, s);
+  auto h0 = ensure_row_contiguous(inputs[5], d, s);
+
+  auto& out = outputs[0];
+  auto& hf = outputs[1];
+
+  int B = q.shape(0);
+  int T = q.shape(1);
+  int Hk = q.shape(2);
+  int Dk = q.shape(3);
+  int Hv = v.shape(2);
+  int Dv = v.shape(3);
+
+  int C = 1;
+  const char* threashold_env = std::getenv("GATED_DELTA_THRESH");
+  int threshold = threashold_env ? std::stoi(threashold_env) : 16;
+  if (T > threshold) {
+    if (metal::is_nax_available())
+      C = 16;
+    else
+      C = 8;
+  }
+  const char* chunk_env = std::getenv("GATED_DELTA_CHUNK");
+  C = chunk_env ? std::stoi(chunk_env) : C;
+
+  if (!metal::is_nax_available())
+    C = std::min(C, 8); // override in case nax is not available.
+
+  std::string suffix = get_type_string(q.dtype()) + "_" + std::to_string(Dk) +
+      "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" +
+      std::to_string(Hv);
+
+  auto& compute_encoder = metal::get_command_encoder(s);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+  hf.set_data(allocator::malloc(hf.nbytes()));
+
+  switch (C) {
+    case 16: {
+      std::string kernel_name = "gated_delta_fused_nax_";
+      std::string base_name = kernel_name + suffix;
+
+      base_name += "_" + std::to_string(C);
+
+      std::string hash_name = base_name;
+
+      metal::MTLFCList func_consts = {};
+
+      auto delta_kernel =
+          get_gated_delta_nax_kernel(d, base_name, hash_name, func_consts);
+
+      compute_encoder.set_compute_pipeline_state(delta_kernel);
+      compute_encoder.set_input_array(q, 0);
+      compute_encoder.set_input_array(k, 1);
+      compute_encoder.set_input_array(v, 2);
+      compute_encoder.set_input_array(h0, 3); // initial state in
+      compute_encoder.set_input_array(g, 4);
+      compute_encoder.set_input_array(beta, 5);
+      compute_encoder.set_output_array(out, 6);
+      compute_encoder.set_output_array(hf, 7); // final state out
+      compute_encoder.set_bytes(T, 8);
+
+      auto grid = MTL::Size(32, Dv / 16, B * Hv);
+      auto threads = MTL::Size(32, 4, 1);
+      compute_encoder.dispatch_threads(grid, threads);
+      break;
+    }
+    case 8: {
+      std::string kernel_name = "gated_delta_fused_chunk_";
+      std::string base_name = kernel_name + suffix;
+
+      base_name += "_" + std::to_string(C);
+
+      std::string hash_name = base_name;
+
+      metal::MTLFCList func_consts = {};
+
+      auto delta_kernel =
+          get_gated_delta_kernel(d, base_name, hash_name, func_consts);
+
+      compute_encoder.set_compute_pipeline_state(delta_kernel);
+      compute_encoder.set_input_array(q, 0);
+      compute_encoder.set_input_array(k, 1);
+      compute_encoder.set_input_array(v, 2);
+      compute_encoder.set_input_array(h0, 3); // initial state in
+      compute_encoder.set_input_array(g, 4);
+      compute_encoder.set_input_array(beta, 5);
+      compute_encoder.set_output_array(out, 6);
+      compute_encoder.set_output_array(hf, 7); // final state out
+      compute_encoder.set_bytes(T, 8);
+
+      auto grid = MTL::Size(32, Dv / 8, B * Hv);
+      auto threads = MTL::Size(32, 4, 1);
+      compute_encoder.dispatch_threads(grid, threads);
+      break;
+    }
+    case 1:
+    case 0: {
+      std::string kernel_name = "seq_gated_delta_";
+      std::string base_name = kernel_name + suffix;
+      std::string hash_name = base_name;
+
+      metal::MTLFCList func_consts = {};
+
+      auto delta_kernel =
+          get_gated_delta_kernel(d, base_name, hash_name, func_consts);
+
+      compute_encoder.set_compute_pipeline_state(delta_kernel);
+
+      compute_encoder.set_input_array(q, 0);
+      compute_encoder.set_input_array(k, 1);
+      compute_encoder.set_input_array(v, 2);
+      compute_encoder.set_input_array(g, 3);
+      compute_encoder.set_input_array(beta, 4);
+      compute_encoder.set_input_array(h0, 5);
+      compute_encoder.set_bytes(T, 6);
+      compute_encoder.set_output_array(out, 7);
+      compute_encoder.set_output_array(hf, 8);
+
+      auto grid = MTL::Size(32, Dv, B * Hv);
+      auto threads = MTL::Size(32, 4, 1);
+      compute_encoder.dispatch_threads(grid, threads);
+      break;
+    }
+    default: {
+      throw std::runtime_error(
+          "NYI: Only sequential and chunk size 8,16 are supported");
+    }
+  }
+}
+
+} // namespace mlx::core::fast

mlx/backend/metal/jit/includes.h

@@ -59,4 +59,7 @@ const char* fp_quantized_nax();
 
 const char* steel_attention_nax();
 
+const char* gated_delta_update();
+const char* gated_delta_update_nax();
+
 } // namespace mlx::core::metal