mlx-bun 0.5.0 → 1.0.0

package/src/train/steel-qmm-header.ts

@@ -1,1461 +0,0 @@
-/* eslint-disable */
-// AUTO-PORTED VERBATIM from MLX quantized.h + steel/* (github.com/ml-explore/mlx).
-export const STEEL_QMM_HEADER = String.raw`// mlx metal_kernel preamble provides METAL_FUNC + using namespace metal; we add
-// the steel macros, type_traits + integral_constant (each self-manages its own
-// #pragma METAL internals enable/disable), then the GEMM-specific structs.
-#define STEEL_CONST static constant constexpr const
-#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
-#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)")
-#define MLX_MTL_CONST static constant constexpr const
-MLX_MTL_CONST int SIMD_SIZE = 32;
-// ===== type_traits.h (self-managed internals pragma) =====
-
-
-
-#pragma METAL internals : enable
-
-namespace metal {
-
-template <typename T>
-struct is_empty : metal::bool_constant<__is_empty(T)> {};
-
-#ifdef __cpp_variable_templates
-template <typename T>
-constexpr constant bool is_empty_v = is_empty<T>::value;
-#endif
-
-template <typename... Ts>
-struct make_void {
-  typedef void type;
-};
-
-template <typename... Ts>
-using void_t = typename make_void<Ts...>::type;
-
-template <class T>
-struct is_static : metal::bool_constant<is_empty<remove_cv_t<T>>::value> {};
-
-template <typename T>
-struct pointer_element {};
-
-template <typename T>
-struct pointer_element<thread T*> {
-  using type = remove_cv_t<T>;
-};
-template <typename T>
-struct pointer_element<device T*> {
-  using type = remove_cv_t<T>;
-};
-template <typename T>
-struct pointer_element<constant T*> {
-  using type = remove_cv_t<T>;
-};
-template <typename T>
-struct pointer_element<threadgroup T*> {
-  using type = remove_cv_t<T>;
-};
-
-template <typename T>
-using pointer_element_t = typename pointer_element<remove_cv_t<T>>::type;
-
-} // namespace metal
-
-#pragma METAL internals : disable// ===== integral_constant.h (self-managed internals pragma) =====
-
-
-
-#pragma METAL internals : enable
-
-namespace mlx {
-namespace steel {
-
-///////////////////////////////////////////////////////////////////////////////
-// Integral constant with casting
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, T v>
-struct integral_constant {
-  static constexpr constant T value = v;
-  using value_type = T;
-  using type = integral_constant;
-
-  METAL_FUNC constexpr operator value_type() const noexcept {
-    return value;
-  }
-};
-
-template <bool B>
-using bool_constant = integral_constant<bool, B>;
-using true_type = bool_constant<true>;
-using false_type = bool_constant<false>;
-
-template <class T>
-struct is_integral : bool_constant<metal::is_integral<T>::value> {};
-
-template <class T, T v>
-struct is_integral<integral_constant<T, v>>
-    : bool_constant<metal::is_integral<T>::value> {};
-
-template <typename T>
-constexpr constant bool is_integral_v = is_integral<T>::value;
-
-template <int val>
-using Int = integral_constant<int, val>;
-
-///////////////////////////////////////////////////////////////////////////////
-// Binary Operators on Integral constants
-///////////////////////////////////////////////////////////////////////////////
-
-#define integral_const_binop(__op__, __operator__)          \
-  template <typename T, T tv, typename U, U uv>             \
-  METAL_FUNC constexpr auto __operator__(                   \
-      integral_constant<T, tv>, integral_constant<U, uv>) { \
-    constexpr auto res = tv __op__ uv;                      \
-    return integral_constant<decltype(res), res>{};         \
-  }
-
-integral_const_binop(+, operator+);
-integral_const_binop(-, operator-);
-integral_const_binop(*, operator*);
-integral_const_binop(/, operator/);
-
-integral_const_binop(==, operator==);
-integral_const_binop(!=, operator!=);
-integral_const_binop(<, operator<);
-integral_const_binop(>, operator>);
-integral_const_binop(<=, operator<=);
-integral_const_binop(>=, operator>=);
-
-integral_const_binop(&&, operator&&);
-integral_const_binop(||, operator||);
-
-template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
-METAL_FUNC constexpr auto operator||(true_type, T) {
-  return true_type{};
-}
-template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
-METAL_FUNC constexpr auto operator||(T, true_type) {
-  return true_type{};
-}
-
-template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
-METAL_FUNC constexpr auto operator&&(false_type, T) {
-  return false_type{};
-}
-
-template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
-METAL_FUNC constexpr auto operator&&(T, false_type) {
-  return false_type{};
-}
-
-// Dispatch utilities
-template <typename F>
-void dispatch_bool(bool v, F f) {
-  if (v) {
-    f(true_type{});
-  } else {
-    f(false_type{});
-  }
-}
-
-template <int start, int stop, int step, typename F>
-constexpr void const_for_loop(F f) {
-  if constexpr (start < stop) {
-    constexpr auto idx = Int<start>{};
-    f(idx);
-    const_for_loop<start + step, stop, step, F>(f);
-  }
-}
-
-#undef integral_const_binop
-
-///////////////////////////////////////////////////////////////////////////////
-// Reduction operators
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T>
-METAL_FUNC constexpr T sum(T x) {
-  return x;
-}
-
-template <typename T, typename... Us>
-METAL_FUNC constexpr auto sum(T x, Us... us) {
-  return x + sum(us...);
-}
-
-} // namespace steel
-} // namespace mlx
-
-#pragma METAL internals : disable
-// ===== transforms.h =====
-
-
-
-///////////////////////////////////////////////////////////////////////////////
-// Transforms and Epilogues
-///////////////////////////////////////////////////////////////////////////////
-
-namespace mlx {
-namespace steel {
-
-template <typename OutT, typename InT>
-struct TransformNone {
-  static METAL_FUNC OutT apply(InT x) {
-    return static_cast<OutT>(x);
-  }
-
-  static METAL_FUNC OutT apply(InT x, OutT) {
-    return static_cast<OutT>(x);
-  }
-};
-
-template <typename OutT, typename InT>
-struct TransformAdd {
-  TransformAdd(const float, const float) {}
-
-  static METAL_FUNC OutT apply(InT x) {
-    return static_cast<OutT>(x);
-  }
-
-  static METAL_FUNC OutT apply(InT x, OutT c) {
-    return static_cast<OutT>(x) + c;
-  }
-};
-
-template <typename OutT, typename InT>
-struct TransformAxpby {
-  const float alpha;
-  const float beta;
-
-  TransformAxpby(const float alpha_, const float beta_)
-      : alpha(alpha_), beta(beta_) {}
-
-  static METAL_FUNC OutT apply(InT x) {
-    return static_cast<OutT>(x);
-  }
-
-  METAL_FUNC OutT apply(InT x, OutT c) const {
-    return static_cast<OutT>(
-        x * static_cast<InT>(alpha) + (static_cast<OutT>(beta) * c));
-  }
-};
-
-template <typename T>
-struct AccumHelper {
-  typedef float accum_type;
-};
-
-struct BlockSwizzle {
-  static METAL_FUNC int2
-  swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) {
-    const int tid_x = (tid.x) >> swizzle_log;
-    const int tid_y =
-        ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1));
-    return int2(tid_x, tid_y);
-  }
-};
-
-} // namespace steel
-} // namespace mlx// ===== quantized helpers =====
-template <int bits, int wsize = 8>
-inline constexpr short get_pack_factor() {
-  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
-}
-
-template <int bits, int wsize = 8>
-inline constexpr short get_bytes_per_pack() {
-  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
-  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
-}
-
-// ===== dequantize =====
-template <typename U, int N, int bits>
-inline void
-dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) {
-  static_assert(
-      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
-          bits == 8,
-      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
-
-  if (bits == 2) {
-    U s[4] = {
-        scale,
-        scale / static_cast<U>(4.0f),
-        scale / static_cast<U>(16.0f),
-        scale / static_cast<U>(64.0f)};
-    for (int i = 0; i < (N / 4); i++) {
-      w_local[4 * i] = s[0] * (w[i] & 0x03) + bias;
-      w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias;
-      w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias;
-      w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias;
-    }
-  }
-
-  else if (bits == 3) {
-    for (int i = 0; i < (N / 8); i++) {
-      w_local += 8 * i;
-      w += 3 * i;
-
-      w_local[0] = (w[0] & 0x7) * scale + bias;
-      w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias;
-      w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias;
-      w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias;
-      w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias;
-      w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias;
-      w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias;
-      w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias;
-    }
-  }
-
-  else if (bits == 4) {
-    U s[2] = {scale, scale / static_cast<U>(16.0f)};
-    for (int i = 0; i < (N / 2); i++) {
-      w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias;
-      w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias;
-    }
-  }
-
-  else if (bits == 5) {
-    for (int i = 0; i < (N / 8); i++) {
-      w_local += 8 * i;
-      w += 5 * i;
-
-      w_local[0] = (w[0] & 0x1f) * scale + bias;
-      w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias;
-      w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias;
-      w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias;
-      w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias;
-      w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias;
-      w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias;
-      w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias;
-    }
-  }
-
-  else if (bits == 6) {
-    for (int i = 0; i < (N / 4); i++) {
-      w_local += 4 * i;
-      w += 3 * i;
-      w_local[0] = (w[0] & 0x3f) * scale + bias;
-      w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias;
-      w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias;
-      w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias;
-    }
-  }
-
-  else if (bits == 8) {
-    for (int i = 0; i < N; i++) {
-      w_local[i] = scale * w[i] + bias;
-    }
-  }
-}
-
-// ===== QuantizedBlockLoader =====
-template <
-    typename T,
-    short BROWS,
-    short BCOLS,
-    short dst_ld,
-    short reduction_dim,
-    short tgp_size,
-    short group_size,
-    short bits>
-struct QuantizedBlockLoader {
-  static_assert(
-      BCOLS <= group_size,
-      "The group size should be larger than the columns");
-  static_assert(
-      group_size % BCOLS == 0,
-      "The group size should be divisible by the columns");
-  static_assert(
-      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
-          bits == 8,
-      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
-
-  MLX_MTL_CONST short pack_factor = get_pack_factor<bits, 8>();
-  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack<bits>();
-  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
-  MLX_MTL_CONST short n_reads =
-      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
-  MLX_MTL_CONST short group_steps = group_size / BCOLS;
-
-  const int src_ld;
-  const int tile_stride;
-  short group_step_cnt;
-  const int group_stride;
-
-  const short thread_idx;
-  const short bi;
-  const short bj;
-
-  threadgroup T* dst;
-  const device uint8_t* src;
-  const device T* scales;
-  const device T* biases;
-
-  QuantizedBlockLoader(
-      const device uint8_t* src_,
-      const device T* scales_,
-      const device T* biases_,
-      const int src_ld_,
-      threadgroup T* dst_,
-      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
-      ushort simd_lane_id [[thread_index_in_simdgroup]])
-      : src_ld(src_ld_),
-        tile_stride(
-            reduction_dim ? BCOLS_PACKED * bytes_per_pack
-                          : BROWS * src_ld * bytes_per_pack / pack_factor),
-        group_step_cnt(0),
-        group_stride(BROWS * src_ld / group_size),
-        thread_idx(simd_group_id * 32 + simd_lane_id),
-        bi(n_reads * thread_idx / BCOLS_PACKED),
-        bj((n_reads * thread_idx) % BCOLS_PACKED),
-        dst(dst_ + bi * dst_ld + bj * pack_factor),
-        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
-            bj * bytes_per_pack),
-        scales(scales_ + bi * src_ld / group_size),
-        biases(biases_ + bi * src_ld / group_size) {}
-
-  void load_unsafe() const {
-    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
-      return;
-    }
-
-    T scale = *scales;
-    T bias = *biases;
-    for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor, bits>(
-          src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
-    }
-  }
-
-  void load_safe(short2 src_tile_dim) const {
-    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
-      return;
-    }
-
-    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
-      for (int i = 0; i < n_reads * pack_factor; i++) {
-        dst[i] = T(0);
-      }
-      return;
-    }
-
-    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
-      for (int i = 0; i < n_reads * pack_factor; i++) {
-        dst[i] = T(0);
-      }
-      return;
-    }
-
-    T scale = *scales;
-    T bias = *biases;
-    for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor, bits>(
-          (device uint8_t*)(src + i * bytes_per_pack),
-          scale,
-          bias,
-          dst + i * pack_factor);
-    }
-  }
-
-  void next() {
-    src += tile_stride;
-    if (reduction_dim == 1) {
-      if (group_steps > 1) {
-        group_step_cnt++;
-        if (group_step_cnt == group_steps) {
-          group_step_cnt = 0;
-          scales++;
-          biases++;
-        }
-      } else {
-        scales++;
-        biases++;
-      }
-    } else {
-      scales += group_stride;
-      biases += group_stride;
-    }
-  }
-};
-// ===== BlockLoader (loader.h) =====
-
-
-
-///////////////////////////////////////////////////////////////////////////////
-// Loading helper
-///////////////////////////////////////////////////////////////////////////////
-
-namespace mlx {
-namespace steel {
-
-template <
-    typename T,
-    short BROWS,
-    short BCOLS,
-    short dst_ld,
-    short reduction_dim,
-    short tgp_size,
-    short alignment = 1,
-    short n_reads = (BCOLS * BROWS) / (tgp_size),
-    short TCOLS = BCOLS / n_reads,
-    short TROWS = tgp_size / TCOLS>
-struct BlockLoader {
-  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
-  STEEL_CONST short vec_size = n_reads;
-
-  // Leading dimension for src
-  const int src_ld;
-  const int tile_stride;
-
-  // Thread location indices
-  const short thread_idx;
-  const short bi;
-  const short bj;
-
-  // threadgroup and device memory
-  threadgroup T* dst;
-  const device T* src;
-
-  struct alignas(alignment * sizeof(T)) ReadVector {
-    uint8_t v[sizeof(T) * vec_size];
-  };
-
-  /* Constructor */
-  METAL_FUNC BlockLoader(
-      const device T* src_,
-      const int src_ld_,
-      threadgroup T* dst_,
-      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
-      ushort simd_lane_id [[thread_index_in_simdgroup]])
-      : src_ld(src_ld_),
-        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
-        thread_idx(simd_group_id * 32 + simd_lane_id),
-        bi(thread_idx / TCOLS),
-        bj(vec_size * (thread_idx % TCOLS)),
-        dst(dst_ + bi * dst_ld + bj),
-        src(src_ + bi * src_ld + bj) {}
-
-  /* Apply operation to threadgroup without bound checking */
-  template <typename UnaryOp>
-  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < BROWS; i += TROWS) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < vec_size; j++) {
-        dst[i * dst_ld + j] = op.apply(dst[i * dst_ld + j]);
-      }
-    }
-  }
-
-  /* Load from device memory into threadgroup memory - without bound checking */
-  METAL_FUNC void load_unsafe() const {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < BROWS; i += TROWS) {
-      *((threadgroup ReadVector*)(&dst[i * dst_ld])) =
-          *((const device ReadVector*)(&src[i * src_ld]));
-    }
-  }
-
-  /* Load from device memory into threadgroup memory - with bound checking */
-  METAL_FUNC void load_safe(short2 src_tile_dim) const {
-    src_tile_dim = src_tile_dim - short2(bj, bi);
-
-    // Skip loading if thread has no valid reads
-    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
-      STEEL_PRAGMA_UNROLL
-      for (short i = 0; i < BROWS; i += TROWS) {
-        STEEL_PRAGMA_UNROLL
-        for (short j = 0; j < vec_size; j++) {
-          dst[i * dst_ld + j] = T(0);
-        }
-      }
-      return;
-    }
-
-    // Use fast thread memory for bound checks
-    bool tmp_idx[vec_size];
-    T tmp_val[vec_size];
-
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < BROWS; i += TROWS) {
-      // Make sure tmp_idx only contains valid indices
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < vec_size; j++) {
-        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
-      }
-
-      // Read valid indices into tmp_val
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < vec_size; j++) {
-        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
-      }
-
-      // Zero out unneeded values
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < vec_size; j++) {
-        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
-      }
-
-      // Copy values to threadgroup memory
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < vec_size; j++) {
-        dst[i * dst_ld + j] = tmp_val[j];
-      }
-    }
-  }
-
-  /* Iteration helper */
-  METAL_FUNC void next() {
-    src += tile_stride;
-  }
-};
-
-} // namespace steel
-} // namespace mlx
-// ===== mma.h (BaseMMAFrag/MMATile/tile_matmad) =====
-namespace mlx {
-namespace steel {
-
-template <typename T, int kFragRows_, int kFragCols_>
-struct BaseMMAFrag {
-  static_assert(
-      kFragRows_ == 8,
-      "Only 8 x 8 fragment matrices are currently supported");
-  static_assert(
-      kFragCols_ == 8,
-      "Only 8 x 8 fragment matrices are currently supported");
-};
-
-template <typename T>
-struct BaseMMAFrag<T, 8, 8> {
-  STEEL_CONST int kFragRows = 8;
-  STEEL_CONST int kFragCols = 8;
-
-  STEEL_CONST int kElemsPerFrag = (kFragRows * kFragCols) / 32;
-
-  STEEL_CONST int kElemRows = 1;
-  STEEL_CONST int kElemCols = 2;
-
-  static_assert(
-      kElemRows * kElemCols == kElemsPerFrag,
-      "MMAFrag shape is not consistent with MMAFrag size");
-
-  typedef metal::simdgroup_matrix<T, kFragRows, kFragCols> mat_type;
-  typedef metal::vec<T, kElemsPerFrag> frag_type;
-
-  METAL_FUNC static constexpr short2 get_coord(
-      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
-    const short qid = simd_lane_id / 4;
-    const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
-    const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
-    return short2{fn, fm};
-  }
-
-  template <typename SrcPtrType, typename StrX, typename StrY>
-  METAL_FUNC static constexpr void
-  load(thread frag_type& dst, SrcPtrType src, StrX str_x, StrY str_y) {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kElemRows; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kElemCols; j++) {
-        dst[i * kElemCols + j] = static_cast<T>(src[i * str_x + j * str_y]);
-      }
-    }
-  }
-
-  template <
-      typename SrcPtrType,
-      typename StrX,
-      typename StrY,
-      typename LimX,
-      typename LimY,
-      typename OffX,
-      typename OffY>
-  METAL_FUNC static constexpr void load_safe(
-      thread frag_type& dst,
-      SrcPtrType src,
-      StrX str_x,
-      StrY str_y,
-      LimX lim_x,
-      LimY lim_y,
-      OffX off_x = Int<0>{},
-      OffY off_y = Int<0>{}) {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kElemRows; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kElemCols; j++) {
-        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
-          dst[i * kElemCols + j] =
-              static_cast<T>(src[(off_x + i) * str_x + (off_y + j) * str_y]);
-        } else {
-          dst[i * kElemCols + j] = T(0);
-        }
-      }
-    }
-  }
-
-  template <typename DstPtrType, typename StrX, typename StrY>
-  METAL_FUNC static constexpr void
-  store(const thread frag_type& src, DstPtrType dst, StrX str_x, StrY str_y) {
-    using U = pointer_element_t<DstPtrType>;
-
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kElemRows; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kElemCols; j++) {
-        dst[i * str_x + j * str_y] = static_cast<U>(src[i * kElemCols + j]);
-      }
-    }
-  }
-
-  template <
-      typename DstPtrType,
-      typename StrX,
-      typename StrY,
-      typename LimX,
-      typename LimY,
-      typename OffX,
-      typename OffY>
-  METAL_FUNC static constexpr void store_safe(
-      const thread frag_type& src,
-      DstPtrType dst,
-      StrX str_x,
-      StrY str_y,
-      LimX lim_x,
-      LimY lim_y,
-      OffX off_x = Int<0>{},
-      OffY off_y = Int<0>{}) {
-    using U = pointer_element_t<DstPtrType>;
-
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kElemRows; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kElemCols; j++) {
-        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
-          dst[(off_x + i) * str_x + (off_y + j) * str_y] =
-              static_cast<U>(src[i * kElemCols + j]);
-        }
-      }
-    }
-  }
-
-  template <
-      typename DstPtrType,
-      typename StrX,
-      typename StrY,
-      typename StartX,
-      typename StopX,
-      typename StartY,
-      typename StopY,
-      typename OffX,
-      typename OffY>
-  METAL_FUNC static constexpr void store_slice(
-      const thread frag_type& src,
-      DstPtrType dst,
-      StrX str_x,
-      StrY str_y,
-      StartX start_x,
-      StopX stop_x,
-      StartY start_y,
-      StopY stop_y,
-      OffX off_x = Int<0>{},
-      OffY off_y = Int<0>{}) {
-    using U = pointer_element_t<DstPtrType>;
-
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kElemRows; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kElemCols; j++) {
-        if ((off_x + i) < stop_x && (off_x + i) >= start_x &&
-            (off_y + j) < stop_y && (off_y + j) >= start_y) {
-          dst[(off_x + i) * str_x + (off_y + j) * str_y] =
-              static_cast<U>(src[i * kElemCols + j]);
-        }
-      }
-    }
-  }
-
-  METAL_FUNC static constexpr void mma(
-      thread frag_type& D,
-      thread frag_type& A,
-      thread frag_type& B,
-      thread frag_type& C) {
-    mat_type D_mat;
-    mat_type A_mat;
-    mat_type B_mat;
-    mat_type C_mat;
-
-    reinterpret_cast<thread frag_type&>(A_mat.thread_elements()) = A;
-    reinterpret_cast<thread frag_type&>(B_mat.thread_elements()) = B;
-    reinterpret_cast<thread frag_type&>(C_mat.thread_elements()) = C;
-
-    mma(D_mat, A_mat, B_mat, C_mat);
-
-    D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
-  }
-
-  METAL_FUNC static constexpr void mma(
-      thread mat_type& D,
-      thread mat_type& A,
-      thread mat_type& B,
-      thread mat_type& C) {
-    simdgroup_multiply_accumulate(D, A, B, C);
-  }
-};
-
-template <
-    typename T,
-    int kTileRows_,
-    int kTileCols_,
-    class MMAFrag_ = BaseMMAFrag<T, 8, 8>>
-struct MMATile {
-  using MMAFrag_t = MMAFrag_;
-  using elem_type = T;
-  STEEL_CONST int kFragRows = MMAFrag_t::kFragRows;
-  STEEL_CONST int kFragCols = MMAFrag_t::kFragCols;
-  STEEL_CONST int kElemsPerFrag = MMAFrag_t::kElemsPerFrag;
-
-  STEEL_CONST int kTileRows = kTileRows_;
-  STEEL_CONST int kTileCols = kTileCols_;
-
-  STEEL_CONST int kRows = kTileRows * kFragRows;
-  STEEL_CONST int kCols = kTileCols * kFragCols;
-
-  STEEL_CONST int kNumFrags = kTileRows * kTileCols;
-  STEEL_CONST int kElemsPerTile = kNumFrags * kElemsPerFrag;
-
-  typedef typename MMAFrag_t::mat_type mat_type;
-  typedef typename MMAFrag_t::frag_type frag_type;
-
-  frag_type val_frags[kNumFrags] = {frag_type(0)};
-
-  METAL_FUNC MMATile() thread {}
-
-  METAL_FUNC constexpr void clear() {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kNumFrags; ++i) {
-      val_frags[i] = frag_type(0);
-    }
-  }
-
-  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
-    return val_frags[i * kTileCols + j];
-  }
-
-  METAL_FUNC constexpr const thread frag_type& frag_at(
-      const short i,
-      const short j) const {
-    return val_frags[i * kTileCols + j];
-  }
-
-  METAL_FUNC mat_type mat_at(const short i, const short j) {
-    mat_type val_mat;
-    STEEL_PRAGMA_UNROLL
-    for (short ii = 0; ii < kElemsPerFrag; ++ii) {
-      val_mat.thread_elements()[ii] = frag_at(i, j)[ii];
-    }
-    return val_mat;
-  }
-
-  METAL_FUNC thread elem_type* elems() {
-    return reinterpret_cast<thread elem_type*>(val_frags);
-  }
-
-  METAL_FUNC const thread elem_type* elems() const {
-    return reinterpret_cast<const thread elem_type*>(val_frags);
-  }
-
-  template <typename U, int w_x, int w_y, int str_x, int str_y>
-  METAL_FUNC void load(const threadgroup U* src) {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::load(
-            frag_at(i, j),
-            &(
-                src[(i * kFragRows) * w_x * str_x +
-                    (j * kFragCols) * w_y * str_y]),
-            Int<str_x>{},
-            Int<str_y>{});
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y, int str_x, int str_y>
-  METAL_FUNC void store(threadgroup U* dst) const {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::store(
-            frag_at(i, j),
-            &(
-                dst[(i * kFragRows) * w_x * str_x +
-                    (j * kFragCols) * w_y * str_y]),
-            Int<str_x>{},
-            Int<str_y>{});
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y>
-  METAL_FUNC void load(const device U* src, const int ld) {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::load(
-            frag_at(i, j),
-            &(src[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
-            ld,
-            Int<1>{});
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y>
-  METAL_FUNC void store(device U* dst, const int ld) const {
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::store(
-            frag_at(i, j),
-            &(dst[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
-            ld,
-            Int<1>{});
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y>
-  METAL_FUNC void
-  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
-    STEEL_PRAGMA_UNROLL
-    for (int i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (int j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::load_safe(
-            frag_at(i, j),
-            src,
-            ld,
-            Int<1>{},
-            src_tile_dims.y,
-            src_tile_dims.x,
-            (i * kFragRows) * w_x,
-            (j * kFragCols) * w_y);
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y>
-  METAL_FUNC void
-  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
-    STEEL_PRAGMA_UNROLL
-    for (int i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (int j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::store_safe(
-            frag_at(i, j),
-            dst,
-            ld,
-            Int<1>{},
-            dst_tile_dims.y,
-            dst_tile_dims.x,
-            (i * kFragRows) * w_x,
-            (j * kFragCols) * w_y);
-      }
-    }
-  }
-
-  template <typename U, int w_x, int w_y>
-  METAL_FUNC void store_slice(
-      device U* dst,
-      const int ld,
-      const short2 start,
-      const short2 stop) const {
-    STEEL_PRAGMA_UNROLL
-    for (int i = 0; i < kTileRows; ++i) {
-      STEEL_PRAGMA_UNROLL
-      for (int j = 0; j < kTileCols; ++j) {
-        MMAFrag_t::store_slice(
-            frag_at(i, j),
-            dst,
-            ld,
-            Int<1>{},
-            start.y,
-            stop.y,
-            start.x,
-            stop.x,
-            (i * kFragRows) * w_x,
-            (j * kFragCols) * w_y);
-      }
-    }
-  }
-};
-
-template <typename T, typename U, int M, int N, int K>
-METAL_FUNC void tile_matmad(
-    thread MMATile<T, M, N>& D,
-    thread MMATile<U, M, K>& A,
-    thread MMATile<U, K, N>& B,
-    thread MMATile<T, M, N>& C) {
-  STEEL_PRAGMA_UNROLL
-  for (short m = 0; m < M; ++m) {
-    STEEL_PRAGMA_UNROLL
-    for (short n = 0; n < N; ++n) {
-      short n_serp = (m % 2) ? (N - 1 - n) : n;
-      STEEL_PRAGMA_UNROLL
-      for (short k = 0; k < K; ++k) {
-        MMATile<T, M, N>::MMAFrag_t::mma(
-            D.frag_at(m, n_serp),
-            A.frag_at(m, k),
-            B.frag_at(k, n_serp),
-            C.frag_at(m, n_serp));
-      }
-    }
-  }
-}
-// ===== mma.h (BlockMMA non-complex) =====
-template <
-    typename T,
-    typename U,
-    int BM,
-    int BN,
-    int BK,
-    int WM,
-    int WN,
-    bool transpose_a,
-    bool transpose_b,
-    short lda_tgp,
-    short ldb_tgp,
-    typename AccumType = float,
-    typename Epilogue = TransformNone<U, AccumType>>
-struct BlockMMA {
-  // MMAFrag size
-  STEEL_CONST short kFragSize = 8;
-  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
-
-  // Warp tile simdgroup matrix strides along M
-  STEEL_CONST short TM_stride = kFragSize * WM;
-  // Warp tile simdgroup matrix strides along M
-  STEEL_CONST short TN_stride = kFragSize * WN;
-
-  // Warp tile size along M
-  STEEL_CONST short TM = BM / (kFragSize * WM);
-  // Warp tile size along N
-  STEEL_CONST short TN = BN / (kFragSize * WN);
-
-  // Threadgroup A strides
-  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
-  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
-
-  // Threadgroup B strides
-  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
-  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
-
-  // Threadgroup strides along K
-  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
-  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
-
-  // Simdgroup matrices
-  MMATile<AccumType, TM, 1, MMAFrag_acc_t> Atile;
-  MMATile<AccumType, 1, TN, MMAFrag_acc_t> Btile;
-  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile;
-
-  // Offsets within threadgroup
-  short sm;
-  short sn;
-
-  short As_offset;
-  short Bs_offset;
-
-  /* Constructor */
-  METAL_FUNC BlockMMA(
-      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
-      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
-    // Determine thread position in simdgroup matrix
-    short tm = kFragSize * (simd_group_id / WN);
-    short tn = kFragSize * (simd_group_id % WN);
-
-    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
-    sm = simd_coord.y;
-    sn = simd_coord.x;
-
-    // Determine thread and simdgroup offset
-    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // M, K
-    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // K, N
-
-    sm += tm;
-    sn += tn;
-  }
-
-  /* (BM, BK) X (BK, BN) multiply accumulate function */
-  METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) {
-    // Adjust for simdgroup and thread location
-    As += As_offset;
-    Bs += Bs_offset;
-
-    // Iterate over BK in blocks of kFragSize
-    STEEL_PRAGMA_UNROLL
-    for (short kk = 0; kk < BK; kk += kFragSize) {
-      simdgroup_barrier(mem_flags::mem_none);
-
-      Atile.template load<T, WM, 1, A_str_m, A_str_k>(As);
-
-      simdgroup_barrier(mem_flags::mem_none);
-
-      Btile.template load<T, 1, WN, B_str_k, B_str_n>(Bs);
-
-      simdgroup_barrier(mem_flags::mem_none);
-
-      tile_matmad(Ctile, Atile, Btile, Ctile);
-
-      // Progress to next simdgroup tile
-      As += tile_stride_a;
-      Bs += tile_stride_b;
-    }
-  }
-
-  /* Store results from simdgroup_matrix results into device memory */
-  METAL_FUNC void store_result(device U* D, const int ldd) {
-    // Apply epilogue
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
-      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
-    }
-
-    // Adjust for simdgroup and thread location
-    D += sm * ldd + sn;
-
-    Ctile.template store<U, WM, WN>(D, ldd);
-  }
-
-  METAL_FUNC void
-  store_result_slice(device U* D, const int ldd, short2 start, short2 stop) {
-    // Apply epilogue
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
-      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
-    }
-
-    D += sm * ldd + sn;
-    start -= short2(sn, sm);
-    stop -= short2(sn, sm);
-
-    // TODO: Check the start as well
-    if (stop.y <= 0 || stop.x <= 0) {
-      return;
-    }
-
-    Ctile.template store_slice<U, WM, WN>(D, ldd, start, stop);
-  }
-
-  METAL_FUNC void
-  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
-    // Apply epilogue
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
-      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
-    }
-
-    // Adjust for simdgroup and thread location
-    D += sm * ldd + sn;
-    dst_tile_dims -= short2(sn, sm);
-
-    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
-      return;
-
-    Ctile.template store_safe<U, WM, WN>(D, ldd, dst_tile_dims);
-  }
-
-  /* Apply epilogue */
-  template <typename UnaryEpilogue>
-  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
-    // Loop over all simdgroup tiles
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
-      Ctile.elems()[i] = epilogue_op.apply(Ctile.elems()[i]);
-    }
-  }
-
-  /* Apply epilogue */
-  template <typename BinaryEpilogue>
-  METAL_FUNC void apply_epilogue(
-      const device U* C,
-      const int ldc,
-      const int fdc,
-      thread const BinaryEpilogue& epilogue_op) {
-    // Adjust for simdgroup and thread location
-    C += (sm)*ldc + (sn)*fdc;
-
-    // Loop over all simdgroup tiles
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < TM; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < TN; j++) {
-        // Get accumulated result and associated offset in C
-        thread auto& accum = Ctile.frag_at(i, j);
-        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
-
-        // Apply epilogue
-        STEEL_PRAGMA_UNROLL
-        for (short k = 0; k < decltype(Ctile)::kElemsPerFrag; k++) {
-          accum[k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
-        }
-      }
-    }
-  }
-
-  /* Apply epilogue */
-  template <typename BinaryEpilogue>
-  METAL_FUNC void apply_epilogue_safe(
-      const device U* C,
-      const int ldc,
-      const int fdc,
-      short2 dst_tile_dims,
-      thread const BinaryEpilogue& epilogue_op) {
-    // Adjust for simdgroup and thread location
-    C += (sm)*ldc + (sn)*fdc;
-    dst_tile_dims -= short2(sn, sm);
-
-    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
-      return;
-
-    // Loop over all simdgroup tiles
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < TM; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < TN; j++) {
-        // Get accumulated result and associated offset in C
-        thread auto& accum = Ctile.frag_at(i, j);
-        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
-
-        constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
-
-        // Read C
-        U c_elems[kelems] = {0};
-
-        STEEL_PRAGMA_UNROLL
-        for (short k = 0; k < kelems; k++) {
-          if ((j * TN_stride + k) < dst_tile_dims.x) {
-            c_elems[k] = C[offset_c + k * fdc];
-          }
-        }
-
-        // Apply epilogue
-        STEEL_PRAGMA_UNROLL
-        for (short k = 0; k < kelems; k++) {
-          accum[k] = epilogue_op.apply(accum[k], c_elems[k]);
-        }
-      }
-    }
-  }
-
-  /* Store results from simdgroup_matrix results into device memory */
-  METAL_FUNC void store_result(
-      device U* D,
-      const int ldd,
-      const device U* C,
-      const int ldc,
-      const int fdc,
-      thread const Epilogue& epilogue_op) const {
-    // Adjust for simdgroup and thread location
-    C += (sm)*ldc + (sn)*fdc;
-    D += (sm)*ldd + sn;
-
-    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
-
-    // Loop over all simdgroup tiles
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < TM; i++) {
-      STEEL_PRAGMA_UNROLL
-      for (short j = 0; j < TN; j++) {
-        // Get accumulated result and associated offset in C
-        thread const auto& accum = Ctile.frag_at(i, j);
-        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
-        int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
-
-        // Apply epilogue
-        STEEL_PRAGMA_UNROLL
-        for (short k = 0; k < kelems; k++) {
-          D[offset_d + k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
-        }
-      }
-    }
-  }
-
-  METAL_FUNC void store_result_safe(
-      device U* D,
-      const int ldd,
-      const device U* C,
-      const int ldc,
-      const int fdc,
-      short2 dst_tile_dims,
-      thread const Epilogue& epilogue_op) const {
-    // Adjust for simdgroup and thread location
-    C += (sm)*ldc + (sn)*fdc;
-    D += (sm)*ldd + sn;
-    dst_tile_dims -= short2(sn, sm);
-
-    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
-      return;
-
-    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
-
-    STEEL_PRAGMA_UNROLL
-    for (int i = 0; i < TM; i++) {
-      if (i * TM_stride < dst_tile_dims.y) {
-        STEEL_PRAGMA_UNROLL
-        for (int j = 0; j < TN; j++) {
-          // Get accumulated result and associated offset in C
-          thread const auto& accum = Ctile.frag_at(i, j);
-          int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
-          int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
-
-          // Apply epilogue
-          STEEL_PRAGMA_UNROLL
-          for (short k = 0; k < kelems; k++) {
-            if ((j * TN_stride + k) < dst_tile_dims.x) {
-              D[offset_d + k] =
-                  epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
-            }
-          }
-        }
-      }
-    }
-  }
-};
-} // namespace steel
-} // namespace mlx
-// ===== qmm_t_impl =====
-template <
-    typename T,
-    const int group_size,
-    const int bits,
-    const bool aligned_N,
-    const int BM = 32,
-    const int BK = 32,
-    const int BN = 32>
-METAL_FUNC void qmm_t_impl(
-    const device uint32_t* w,
-    const device T* scales,
-    const device T* biases,
-    const device T* x,
-    device T* y,
-    threadgroup T* Xs,
-    threadgroup T* Ws,
-    const constant int& K,
-    const constant int& N,
-    const constant int& M,
-    const constant int& K_eff,
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint lid [[thread_index_in_threadgroup]],
-    uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]]) {
-  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
-  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
-
-  (void)lid;
-
-  constexpr int WM = 2;
-  constexpr int WN = 2;
-  constexpr int pack_factor = get_pack_factor<bits, 8>();
-  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
-
-  constexpr int BK_padded = (BK + 16 / sizeof(T));
-
-  // Instantiate the appropriate BlockMMA and Loader
-  using mma_t = mlx::steel::
-      BlockMMA<T, T, BM, BN, BK, WM, WN, false, true, BK_padded, BK_padded>;
-  using loader_x_t =
-      mlx::steel::BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE>;
-  using loader_w_t = QuantizedBlockLoader<
-      T,
-      BN,
-      BK,
-      BK_padded,
-      1,
-      WM * WN * SIMD_SIZE,
-      group_size,
-      bits>;
-
-  // Set the block
-  const int K_w = K * bytes_per_pack / pack_factor;
-  const int K_g = K / group_size;
-  const int y_row = tid.y * BM;
-  const int y_col = tid.x * BN;
-
-  auto wl = (const device uint8_t*)w;
-
-  x += y_row * static_cast<int64_t>(K);
-  wl += y_col * K_w;
-  scales += y_col * K_g;
-  biases += y_col * K_g;
-  y += y_row * static_cast<int64_t>(N) + y_col;
-
-  // Make the x loader and mma operation
-  const short num_els = min(BM, M - y_row);
-  const short num_outs = min(BN, N - y_col);
-  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
-  loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid);
-  mma_t mma_op(simd_gid, simd_lid);
-
-  if (num_els < BM) {
-    if (!aligned_N && num_outs < BN) {
-      for (int k = 0; k < K_eff; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_safe(short2(BK, num_els));
-        loader_w.load_safe(short2(BK, num_outs));
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-    } else {
-      for (int k = 0; k < K_eff; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_safe(short2(BK, num_els));
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-    }
-  } else {
-    if (!aligned_N && num_outs < BN) {
-      for (int k = 0; k < K_eff; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_unsafe();
-        loader_w.load_safe(short2(BK, num_outs));
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-    } else {
-      for (int k = 0; k < K_eff; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_unsafe();
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-    }
-  }
-
-  // Store results to device memory
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  if (num_els < BM || num_outs < BN) {
-    mma_op.store_result_safe(y, N, short2(num_outs, num_els));
-  } else {
-    mma_op.store_result(y, N);
-  }
-}
-`;