Faster contiguous gather for indices in the first axis (#2552)

* faster contiguous gather for indices in the first axis

* work per thread > 1

* angelos suggestion for scales / biases

mlx/backend/metal/CMakeLists.txt

@@ -33,10 +33,11 @@ make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h)
 make_jit_source(binary_ops)
 make_jit_source(ternary_ops)
 make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h)
-make_jit_source(scatter kernels/indexing.h)
+make_jit_source(indexing/scatter kernels/indexing/indexing.h)
-make_jit_source(gather kernels/indexing.h)
+make_jit_source(indexing/gather kernels/indexing/indexing.h)
-make_jit_source(gather_axis)
+make_jit_source(indexing/gather_front kernels/indexing/indexing.h)
-make_jit_source(scatter_axis)
+make_jit_source(indexing/gather_axis)
+make_jit_source(indexing/scatter_axis)
 make_jit_source(hadamard)
 
 if(MLX_METAL_JIT)

mlx/backend/metal/indexing.cpp

@@ -52,8 +52,10 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  size_t slice_size = 1;
-  size_t ndim = src.ndim();
+  for (auto s : slice_sizes_) {
+    slice_size *= s;
+  }
 
   bool large_index = nidx && inputs[1].size() > INT32_MAX;
   bool large_src = src.size() > INT32_MAX;
@@ -61,6 +63,55 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   bool large = large_index || large_src || large_out;
 
   std::string idx_type_name = nidx ? type_to_name(inputs[1]) : "";
+
+  if (src.flags().row_contiguous && nidx == 1 && axes_[0] == 0 &&
+      inputs[1].flags().row_contiguous && slice_size == src.strides()[0]) {
+    int work_per_thread = (slice_size > 8 && src.dtype().size() < 4) ? 2 : 1;
+    auto& indices = inputs[1];
+    std::string kernel_name = fmt::format(
+        "gather_front{0}_{1}_{2}_{3}",
+        type_to_name(out),
+        idx_type_name,
+        large ? "int64_t" : "int",
+        work_per_thread);
+    std::string lib_name = kernel_name;
+
+    auto lib = d.get_library(lib_name, [&]() {
+      std::string kernel_source = metal::utils();
+      kernel_source += metal::gather_front();
+      kernel_source += get_template_definition(
+          kernel_name,
+          "gather_front",
+          get_type_string(out.dtype()),
+          get_type_string(indices.dtype()),
+          large ? "int64_t" : "int",
+          work_per_thread);
+
+      return kernel_source;
+    });
+
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kernel_name, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    size_t dim_x = (slice_size + work_per_thread - 1) / work_per_thread;
+    size_t dim_y = indices.size();
+    auto group_dims = get_block_dims(dim_x, dim_y, 1);
+    MTL::Size grid_dims = MTL::Size(dim_x, dim_y, 1);
+
+    compute_encoder.set_input_array(src, 0);
+    compute_encoder.set_input_array(indices, 1);
+    compute_encoder.set_output_array(out, 2);
+    compute_encoder.set_bytes(slice_size, 3);
+    compute_encoder.set_bytes(src.shape(0), 4);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+    return;
+  }
+
+  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  size_t ndim = src.ndim();
+
   std::string kernel_name = fmt::format(
       "gather{0}{1}_{2}_{3}_{4}",
       type_to_name(out),
@@ -96,11 +147,6 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto kernel = d.get_kernel(kernel_name, lib);
   compute_encoder.set_compute_pipeline_state(kernel);
 
-  size_t slice_size = 1;
-  for (auto s : slice_sizes_) {
-    slice_size *= s;
-  }
-
   // Launch 3D grid of threads
   // First two dimensions for the indices, the last one for the slice
   size_t dim0 = 1;

mlx/backend/metal/jit/includes.h

@@ -19,6 +19,7 @@ const char* binary_two();
 const char* copy();
 const char* fft();
 const char* gather_axis();
+const char* gather_front();
 const char* hadamard();
 const char* logsumexp();
 const char* quantized_utils();

mlx/backend/metal/kernels/indexing/gather.h

@@ -2,7 +2,7 @@
 
 #pragma once
 
-#include "mlx/backend/metal/kernels/indexing.h"
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
 
 template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
 METAL_FUNC void gather_impl(

mlx/backend/metal/kernels/indexing/gather_front.h

@@ -0,0 +1,24 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, typename LocT, int N>
+[[kernel]] void gather_front(
+    const device T* src,
+    const device IdxT* indices,
+    device T* out,
+    const constant int64_t& stride,
+    const constant int& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto idx = offset_neg_idx(indices[index.y], size);
+  LocT src_idx = static_cast<LocT>(stride) * idx;
+  LocT out_idx = static_cast<LocT>(stride) * index.y;
+
+  int s_idx = N * index.x;
+  for (int i = 0; i < N && s_idx < stride; ++i, ++s_idx) {
+    out[out_idx + s_idx] = src[src_idx + s_idx];
+  }
+}

mlx/backend/metal/kernels/indexing/scatter.h

@@ -2,7 +2,7 @@
 
 #pragma once
 
-#include "mlx/backend/metal/kernels/indexing.h"
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
 
 template <
     typename T,

python/mlx/nn/layers/quantized.py

@@ -98,11 +98,10 @@ class QuantizedEmbedding(Module):
         # Initialize the quantized weight
         scale = math.sqrt(1 / dims)
         weight = mx.random.normal(shape=(num_embeddings, dims), scale=scale)
-        self.weight, *scales_biases = mx.quantize(weight, group_size, bits, mode=mode)
+        self.weight, self.scales, *biases = mx.quantize(
-        if mode == "affine":
+            weight, group_size, bits, mode=mode
-            self.scales, self.biases = scales_biases
+        )
-        else:
+        self.biases = biases[0] if biases else None
-            (self.scales,) = scales_biases
         self.num_embeddings = num_embeddings
         self.dims = dims
 
@@ -155,16 +154,13 @@ class QuantizedEmbedding(Module):
         """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
         embedding_dims, dims = embedding_layer.weight.shape
         ql = cls(embedding_dims, dims, group_size, bits, mode=mode)
-        ql.weight, *scales_biases = mx.quantize(
+        ql.weight, ql.scales, *biases = mx.quantize(
             embedding_layer.weight,
             group_size,
             bits,
             mode=mode,
         )
-        if mode == "affine":
+        ql.biases = biases[0] if biases else None
-            ql.scales, ql.biases = scales_biases
-        else:
-            (ql.scales,) = scales_biases
         return ql
 
 
@@ -214,11 +210,10 @@ class QuantizedLinear(Module):
             high=scale,
             shape=(output_dims, input_dims),
         )
-        self.weight, *scales_biases = mx.quantize(weight, group_size, bits, mode=mode)
+        self.weight, self.scales, *biases = mx.quantize(
-        if mode == "affine":
+            weight, group_size, bits, mode=mode
-            self.scales, self.biases = scales_biases
+        )
-        else:
+        self.biases = biases[0] if biases else None
-            (self.scales,) = scales_biases
 
         # And bias if needed
         if bias:
@@ -261,16 +256,13 @@ class QuantizedLinear(Module):
         """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""
         output_dims, input_dims = linear_layer.weight.shape
         ql = cls(input_dims, output_dims, False, group_size, bits, mode=mode)
-        ql.weight, *scales_biases = mx.quantize(
+        ql.weight, ql.scales, *biases = mx.quantize(
             linear_layer.weight,
             group_size,
             bits,
             mode=mode,
         )
-        if mode == "affine":
+        ql.biases = biases[0] if biases else None
-            ql.scales, ql.biases = scales_biases
-        else:
-            (ql.scales,) = scales_biases
 
         if "bias" in linear_layer:
             ql.bias = linear_layer.bias