mlx 1.0.0

data/mlx/mlx/backend/metal/kernels/rms_norm.metal

@@ -0,0 +1,391 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+constant bool has_w [[function_constant(20)]];
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_single_row(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup float local_inv_mean[1];
+  threadgroup float local_sums[SIMD_SIZE];
+
+  float acc = 0;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      float xi = x[i];
+      acc += xi * xi;
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        float xi = x[i];
+        acc += xi * xi;
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_looped(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_inv_mean[1];
+  threadgroup float local_sums[SIMD_SIZE];
+
+  float acc = 0;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        acc += xi * xi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          acc += xi * xi;
+        }
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        out[r + i] = w[w_stride * (i + r)] *
+            static_cast<T>(x[r + i] * local_inv_mean[0]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          out[r + i] = w[w_stride * (i + r)] *
+              static_cast<T>(x[r + i] * local_inv_mean[0]);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void vjp_rms_single_row(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  // Allocate registers for the computation and accumulators
+  float thread_x[N_READS];
+  float thread_w[N_READS];
+  float thread_g[N_READS];
+  float sumx2 = 0;
+  float sumgwx = 0;
+
+  // Allocate shared memory to implement the reduction
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_sumgwx[SIMD_SIZE];
+  threadgroup float local_normalizer[1];
+  threadgroup float local_meangwx[1];
+
+  // Read and accumulate locally
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = x[i];
+      thread_w[i] = w[w_stride * i];
+      thread_g[i] = g[i];
+
+      sumx2 += thread_x[i] * thread_x[i];
+      sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        thread_x[i] = x[i];
+        thread_w[i] = w[w_stride * i];
+        thread_g[i] = g[i];
+
+        sumx2 += thread_x[i] * thread_x[i];
+        sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
+      }
+    }
+  }
+
+  // Accumulate across threads
+  sumx2 = simd_sum(sumx2);
+  sumgwx = simd_sum(sumgwx);
+  if (simd_group_id == 0) {
+    local_sumx2[simd_lane_id] = 0;
+    local_sumgwx[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_lane_id == 0) {
+    local_sumx2[simd_group_id] = sumx2;
+    local_sumgwx[simd_group_id] = sumgwx;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_meangwx[0] = sumgwx / axis_size;
+      local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  float meangwx = local_meangwx[0];
+  float normalizer = local_normalizer[0];
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      gx[i] = static_cast<T>(
+          thread_g[i] * thread_w[i] * normalizer -
+          thread_x[i] * meangwx * normalizer3);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      }
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        gx[i] = static_cast<T>(
+            thread_g[i] * thread_w[i] * normalizer -
+            thread_x[i] * meangwx * normalizer3);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void vjp_rms_looped(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  // Allocate registers for the accumulators
+  float sumx2 = 0;
+  float sumgwx = 0;
+
+  // Allocate shared memory to implement the reduction
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_sumgwx[SIMD_SIZE];
+  threadgroup float local_normalizer[1];
+  threadgroup float local_meangwx[1];
+
+  // Read and accumulate locally
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        float wi = w[w_stride * (i + r)];
+        float gi = g[i + r];
+
+        sumx2 += xi * xi;
+        sumgwx += xi * wi * gi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          float wi = w[w_stride * (i + r)];
+          float gi = g[i + r];
+
+          sumx2 += xi * xi;
+          sumgwx += xi * wi * gi;
+        }
+      }
+    }
+  }
+
+  // Accumulate across threads
+  sumx2 = simd_sum(sumx2);
+  sumgwx = simd_sum(sumgwx);
+  if (simd_group_id == 0) {
+    local_sumx2[simd_lane_id] = 0;
+    local_sumgwx[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_lane_id == 0) {
+    local_sumx2[simd_group_id] = sumx2;
+    local_sumgwx[simd_group_id] = sumgwx;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_meangwx[0] = sumgwx / axis_size;
+      local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  float meangwx = local_meangwx[0];
+  float normalizer = local_normalizer[0];
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        float wi = w[w_stride * (i + r)];
+        float gi = g[i + r];
+
+        gx[i + r] =
+            static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        }
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          float wi = w[w_stride * (i + r)];
+          float gi = g[i + r];
+
+          gx[i + r] =
+              static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          }
+        }
+      }
+    }
+  }
+}
+
+// clang-format off
+#define instantiate_rms(name, itype)                                \
+  instantiate_kernel("rms" #name, rms_single_row, itype)            \
+  instantiate_kernel("vjp_rms" #name, vjp_rms_single_row, itype)    \
+  instantiate_kernel("rms_looped" #name, rms_looped, itype)         \
+  instantiate_kernel("vjp_rms_looped" #name, vjp_rms_looped, itype)
+
+instantiate_rms(float32, float)
+instantiate_rms(float16, half)
+instantiate_rms(bfloat16, bfloat16_t) // clang-format on

data/mlx/mlx/backend/metal/kernels/rope.metal

@@ -0,0 +1,229 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_math>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+constant bool forward [[function_constant(1)]];
+constant bool traditional [[function_constant(2)]];
+constant bool hs_transpose [[function_constant(3)]];
+
+template <typename T>
+void rope_single_impl(
+    const device T* in,
+    device T* out,
+    constant const int& offset,
+    const float inv_freq,
+    constant const float& scale,
+    constant const int64_t& stride,
+    uint2 pos,
+    uint2 grid) {
+  float L = scale * static_cast<float>(offset);
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = metal::fast::cos(theta);
+  float sintheta = metal::fast::sin(theta);
+
+  // Compute the input and output indices
+  uint index_1, index_2;
+  if (traditional) {
+    index_1 = 2 * pos.x + pos.y * stride;
+    index_2 = index_1 + 1;
+  } else {
+    index_1 = pos.x + pos.y * stride;
+    index_2 = index_1 + grid.x;
+  }
+
+  // Read and write the output
+  float x1 = static_cast<float>(in[index_1]);
+  float x2 = static_cast<float>(in[index_2]);
+  float rx1;
+  float rx2;
+  if (forward) {
+    rx1 = x1 * costheta - x2 * sintheta;
+    rx2 = x1 * sintheta + x2 * costheta;
+  } else {
+    rx1 = x2 * sintheta + x1 * costheta;
+    rx2 = x2 * costheta - x1 * sintheta;
+  }
+  out[index_1] = static_cast<T>(rx1);
+  out[index_2] = static_cast<T>(rx2);
+}
+
+template <typename T>
+[[kernel]] void rope_single(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const int& offset,
+    constant const float& scale,
+    constant const int64_t& stride,
+    constant const float& base [[buffer(10)]],
+    uint2 pos [[thread_position_in_grid]],
+    uint2 grid [[threads_per_grid]]) {
+  float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
+  float inv_freq = metal::exp2(-d * base);
+  rope_single_impl<T>(in, out, offset, inv_freq, scale, stride, pos, grid);
+}
+
+template <typename T>
+[[kernel]] void rope_single_freqs(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const int& offset,
+    constant const float& scale,
+    constant const int64_t& stride,
+    const device float* freqs [[buffer(10)]],
+    constant const int64_t& freq_stride [[buffer(11)]],
+    uint2 pos [[thread_position_in_grid]],
+    uint2 grid [[threads_per_grid]]) {
+  float inv_freq = 1.0 / (freqs[freq_stride * pos.x]);
+  rope_single_impl<T>(in, out, offset, inv_freq, scale, stride, pos, grid);
+}
+
+template <typename T, typename IdxT, int N = 4>
+void rope_impl(
+    const device T* in,
+    device T* out,
+    const device int* offset,
+    const float inv_freq,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    uint3 pos,
+    uint3 grid) {
+  auto n_head_up = N * ((n_head + N - 1) / N);
+  auto head_idx = static_cast<int>((pos.z * N) % n_head_up);
+  auto batch_idx = (pos.z * N) / n_head_up;
+  auto batch_offset = offset[batch_idx * offset_stride];
+  float L = scale * static_cast<float>(pos.y + batch_offset);
+  auto mat_idx = batch_idx * n_head + head_idx;
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = metal::fast::cos(theta);
+  float sintheta = metal::fast::sin(theta);
+  // Compute the input and output indices
+  IdxT in_index_1;
+  if (hs_transpose) {
+    IdxT batch_stride = grid.y * IdxT(strides[1]);
+    in_index_1 =
+        batch_idx * batch_stride + pos.y * strides[1] + head_idx * strides[0];
+  } else {
+    in_index_1 = pos.y * IdxT(strides[1]) + mat_idx * IdxT(strides[0]);
+  }
+  IdxT in_index_2;
+  IdxT out_index_1 =
+      pos.y * IdxT(out_strides[1]) + mat_idx * IdxT(out_strides[0]);
+  IdxT out_index_2;
+  if (traditional) {
+    out_index_1 += 2 * pos.x * IdxT(out_strides[2]);
+    out_index_2 = out_index_1 + 1;
+    in_index_1 += 2 * pos.x * IdxT(strides[2]);
+    in_index_2 = in_index_1 + IdxT(strides[2]);
+  } else {
+    out_index_1 += pos.x * IdxT(out_strides[2]);
+    out_index_2 = out_index_1 + grid.x * IdxT(out_strides[2]);
+    in_index_1 += pos.x * IdxT(strides[2]);
+    in_index_2 = in_index_1 + grid.x * IdxT(strides[2]);
+  }
+  for (int i = 0; i < N && head_idx + i < n_head; ++i) {
+    // Read and write the output
+    float x1 = static_cast<float>(in[in_index_1]);
+    float x2 = static_cast<float>(in[in_index_2]);
+    float rx1;
+    float rx2;
+    if (forward) {
+      rx1 = x1 * costheta - x2 * sintheta;
+      rx2 = x1 * sintheta + x2 * costheta;
+    } else {
+      rx1 = x2 * sintheta + x1 * costheta;
+      rx2 = x2 * costheta - x1 * sintheta;
+    }
+    out[out_index_1] = static_cast<T>(rx1);
+    out[out_index_2] = static_cast<T>(rx2);
+    in_index_1 += IdxT(strides[0]);
+    in_index_2 += IdxT(strides[0]);
+    out_index_1 += IdxT(out_strides[0]);
+    out_index_2 += IdxT(out_strides[0]);
+  }
+}
+
+template <typename T, typename IdxT, int N = 4>
+[[kernel]] void rope(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const device int* offset,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    constant const float& base [[buffer(10)]],
+    uint3 pos [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
+  float inv_freq = metal::exp2(-d * base);
+  rope_impl<T, IdxT, N>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      grid);
+}
+
+template <typename T, typename IdxT, int N = 4>
+[[kernel]] void rope_freqs(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const device int* offset,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    const device float* freqs [[buffer(10)]],
+    constant const int64_t& freq_stride [[buffer(11)]],
+    uint3 pos [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  float inv_freq = 1.0 / (freqs[freq_stride * pos.x]);
+  rope_impl<T, IdxT, N>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      grid);
+}
+
+// clang-format off
+#define instantiate_rope_g(name, type) \
+  instantiate_kernel("rope_" #name, rope, type, int32_t) \
+  instantiate_kernel("rope_freqs_" #name, rope_freqs, type, int32_t) \
+  instantiate_kernel("rope_large_" #name, rope, type, int64_t) \
+  instantiate_kernel("rope_freqs_large_" #name, rope_freqs, type, int64_t)
+
+#define instantiate_rope_s(name, type) \
+  instantiate_kernel("rope_single_" #name, rope_single, type) \
+  instantiate_kernel("rope_single_freqs_" #name, rope_single_freqs, type)
+
+#define instantiate_rope(name, type) \
+  instantiate_rope_s(name, type)     \
+  instantiate_rope_g(name, type)
+
+instantiate_rope(float16, half)
+instantiate_rope(bfloat16, bfloat16_t)
+instantiate_rope(float32, float) // clang-format on

data/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -0,0 +1,44 @@
+#include <metal_stdlib>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/sdpa_vector.h"
+
+using namespace metal;
+
+// SDPA vector instantiations
+#define instantiate_sdpa_vector_aggregation(type, value_dim) \
+  instantiate_kernel(                                        \
+      "sdpa_vector_2pass_2_" #type "_" #value_dim,           \
+      sdpa_vector_2pass_2,                                   \
+      type,                                                  \
+      value_dim)
+
+#define instantiate_sdpa_vector(type, qk_dim, value_dim)       \
+  instantiate_kernel(                                          \
+      "sdpa_vector_" #type "_" #qk_dim "_" #value_dim,         \
+      sdpa_vector,                                             \
+      type,                                                    \
+      qk_dim,                                                  \
+      value_dim)                                               \
+  instantiate_kernel(                                          \
+      "sdpa_vector_2pass_1_" #type "_" #qk_dim "_" #value_dim, \
+      sdpa_vector_2pass_1,                                     \
+      type,                                                    \
+      qk_dim,                                                  \
+      value_dim)
+
+#define instantiate_sdpa_vector_heads(type)      \
+  instantiate_sdpa_vector(type, 64, 64)          \
+  instantiate_sdpa_vector(type, 96, 96)          \
+  instantiate_sdpa_vector(type, 128, 128)        \
+  instantiate_sdpa_vector(type, 256, 256)        \
+  instantiate_sdpa_vector_aggregation(type, 64)  \
+  instantiate_sdpa_vector_aggregation(type, 96)  \
+  instantiate_sdpa_vector_aggregation(type, 128) \
+  instantiate_sdpa_vector_aggregation(type, 256)
+
+instantiate_sdpa_vector_heads(float)
+instantiate_sdpa_vector_heads(bfloat16_t)
+instantiate_sdpa_vector_heads(float16_t)
+    // clang-format on