Add fused Metal cross entropy kernel

benchmarks/python/cross_entropy_bench.py

@@ -0,0 +1,86 @@
+"""Fused Metal cross entropy vs the decomposed fallback."""
+
+import time
+
+import mlx.core as mx
+
+
+def fallback(logits, targets):
+    score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze(-1)
+    return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype(mx.float32)
+
+
+def timeit(fn, *args, iters=50, warmup=10):
+    # Multi-GiB logits leave the buffer cache under pressure, which otherwise
+    # bleeds between measurements and makes later shapes look far slower.
+    mx.clear_cache()
+    for _ in range(warmup):
+        mx.eval(fn(*args))
+    mx.synchronize()
+    start = time.perf_counter()
+    for _ in range(iters):
+        mx.eval(fn(*args))
+    mx.synchronize()
+    return (time.perf_counter() - start) / iters * 1e3
+
+
+shapes = [
+    (8192, 4096),
+    (4096, 32000),
+    (2048, 128256),
+    (1024, 151936),
+]
+
+print(
+    f"{'rows x vocab':>20} {'dtype':>10} {'fallback':>10} {'fused':>10} {'speedup':>9}"
+)
+print("-" * 64)
+
+for dtype in [mx.float32, mx.bfloat16]:
+    for rows, V in shapes:
+        logits = mx.random.normal(shape=(rows, V), scale=2.0).astype(dtype)
+        targets = mx.random.randint(0, V, shape=(rows,))
+        mx.eval(logits, targets)
+
+        t_ref = timeit(lambda: fallback(logits, targets))
+        t_fused = timeit(lambda: mx.fast.cross_entropy(logits, targets))
+        print(
+            f"{rows:>7} x {V:<10} {str(dtype).split('.')[-1]:>10} "
+            f"{t_ref:>9.3f}ms {t_fused:>9.3f}ms {t_ref / t_fused:>8.2f}x"
+        )
+
+print()
+print("--- forward + backward ---")
+print(
+    f"{'rows x vocab':>20} {'dtype':>10} {'fallback':>10} {'fused':>10} {'speedup':>9}"
+)
+print("-" * 64)
+
+for dtype in [mx.float32, mx.bfloat16]:
+    for rows, V in shapes:
+        logits = mx.random.normal(shape=(rows, V), scale=2.0).astype(dtype)
+        targets = mx.random.randint(0, V, shape=(rows,))
+        mx.eval(logits, targets)
+
+        g_ref = mx.grad(lambda x, y: fallback(x, y).sum(), argnums=0)
+        g_fused = mx.grad(lambda x, y: mx.fast.cross_entropy(x, y).sum(), argnums=0)
+        t_ref = timeit(g_ref, logits, targets, iters=20)
+        t_fused = timeit(g_fused, logits, targets, iters=20)
+        print(
+            f"{rows:>7} x {V:<10} {str(dtype).split('.')[-1]:>10} "
+            f"{t_ref:>9.3f}ms {t_fused:>9.3f}ms {t_ref / t_fused:>8.2f}x"
+        )
+
+print()
+print("--- peak memory, forward + backward, 4096 x 128256 bf16 ---")
+for label, fn in [
+    ("fallback", lambda x, y: fallback(x, y).sum()),
+    ("fused", lambda x, y: mx.fast.cross_entropy(x, y).sum()),
+]:
+    logits = mx.random.normal(shape=(4096, 128256), scale=2.0).astype(mx.bfloat16)
+    targets = mx.random.randint(0, 128256, shape=(4096,))
+    mx.eval(logits, targets)
+    mx.clear_cache()
+    mx.reset_peak_memory()
+    mx.eval(mx.grad(fn, argnums=0)(logits, targets))
+    print(f"  {label:>10}: {mx.get_peak_memory() / 2**30:.2f} GiB")

mlx/backend/metal/CMakeLists.txt

@@ -39,6 +39,7 @@ if(MLX_METAL_JIT)
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jit_kernels.cpp)
   make_jit_source(arange)
   make_jit_source(copy)
+  make_jit_source(cross_entropy)
   make_jit_source(unary)
   make_jit_source(binary)
   make_jit_source(binary_two)
@@ -124,6 +125,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp

mlx/backend/metal/cross_entropy.cpp

@@ -0,0 +1,157 @@
+// Copyright © 2026 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+namespace {
+
+// Above this the single pass kernel would need more than one read per thread,
+// so switch to the looped variant.
+constexpr int CROSS_ENTROPY_LOOPED_LIMIT = 4096;
+constexpr int SIMD_SIZE = 32;
+constexpr int N_READS = 4;
+
+size_t ceil_simd_multiple(size_t n) {
+  return SIMD_SIZE * ((n + SIMD_SIZE - 1) / SIMD_SIZE);
+}
+
+} // namespace
+
+bool CrossEntropy::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void CrossEntropy::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 2); // logits and targets
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& compute_encoder = metal::get_command_encoder(s);
+  auto& out = outputs[0];
+
+  auto ensure_row_contiguous = [&s, &compute_encoder](const array& x) {
+    if (x.flags().row_contiguous) {
+      return x;
+    }
+    array x_copy = contiguous_copy_gpu(x, s);
+    compute_encoder.add_temporary(x_copy);
+    return x_copy;
+  };
+
+  auto in = ensure_row_contiguous(inputs[0]); // [n_rows, V]
+  auto target = ensure_row_contiguous(inputs[1]); // [n_rows]
+  out.set_data(allocator::malloc(out.nbytes())); // [n_rows] in fp32
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  auto get_kernel = [&d, &in](bool looped) {
+    std::string kernel_name = looped ? "looped_" : "block_";
+    kernel_name += "cross_entropy_";
+    kernel_name += type_to_name(in);
+    return get_cross_entropy_kernel(d, kernel_name, in);
+  };
+
+  bool looped = axis_size > CROSS_ENTROPY_LOOPED_LIMIT;
+  auto kernel = get_kernel(looped);
+  size_t threadgroup_size = 0;
+  if (!looped) {
+    threadgroup_size = ceil_simd_multiple((axis_size + N_READS - 1) / N_READS);
+    // The single pass kernel needs one thread per N_READS elements. If the
+    // pipeline cannot hold that many threads, fall back to the looped variant.
+    if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) {
+      looped = true;
+      kernel = get_kernel(true);
+    }
+  }
+  if (looped) {
+    threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+  }
+
+  MTL::Size grid_dims(n_rows * threadgroup_size, 1, 1);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(target, 1);
+  compute_encoder.set_output_array(out, 2);
+  compute_encoder.set_bytes(axis_size, 3);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void CrossEntropyVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 4); // logits, targets, loss, cotangent
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& compute_encoder = metal::get_command_encoder(s);
+  auto& out = outputs[0];
+
+  auto ensure_row_contiguous = [&s, &compute_encoder](const array& x) {
+    if (x.flags().row_contiguous) {
+      return x;
+    }
+    array x_copy = contiguous_copy_gpu(x, s);
+    compute_encoder.add_temporary(x_copy);
+    return x_copy;
+  };
+
+  // The gradient has the same shape and type as the logits, so write it into
+  // the logits buffer whenever that buffer is ours to reuse. The kernel reads
+  // the target score into registers and barriers before overwriting anything.
+  auto set_output = [&s, &out](const array& x) {
+    if (x.flags().row_contiguous) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(allocator::malloc(out.nbytes()));
+      }
+      return x;
+    }
+    array x_copy = contiguous_copy_gpu(x, s);
+    out.copy_shared_buffer(x_copy);
+    return x_copy;
+  };
+
+  auto in = set_output(inputs[0]); // [n_rows, V]
+  auto target = ensure_row_contiguous(inputs[1]); // [n_rows]
+  auto loss = ensure_row_contiguous(inputs[2]); // [n_rows] fp32
+  auto cotan = ensure_row_contiguous(inputs[3]); // [n_rows] fp32
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  std::string kernel_name = "vjp_cross_entropy_";
+  kernel_name += type_to_name(in);
+  auto kernel = get_cross_entropy_kernel(d, kernel_name, in);
+
+  // The kernel loops over the row, so any threadgroup size is correct. Use
+  // just enough threads to cover the row without oversubscribing short rows.
+  size_t threadgroup_size = std::min(
+      static_cast<size_t>(kernel->maxTotalThreadsPerThreadgroup()),
+      ceil_simd_multiple((axis_size + N_READS - 1) / N_READS));
+
+  MTL::Size grid_dims(n_rows * threadgroup_size, 1, 1);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(target, 1);
+  compute_encoder.set_input_array(loss, 2);
+  compute_encoder.set_input_array(cotan, 3);
+  compute_encoder.set_output_array(out, 4);
+  compute_encoder.set_bytes(axis_size, 5);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+} // namespace mlx::core::fast

mlx/backend/metal/jit/includes.h

@@ -18,6 +18,7 @@ const char* unary();
 const char* binary();
 const char* binary_two();
 const char* copy();
+const char* cross_entropy();
 const char* fft();
 const char* gather_axis();
 const char* gather_front();

mlx/backend/metal/jit_kernels.cpp

@@ -360,6 +360,28 @@ MTL::ComputePipelineState* get_logsumexp_kernel(
   return d.get_kernel(kernel_name, lib);
 }
 
+MTL::ComputePipelineState* get_cross_entropy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&] {
+    // The library is keyed on the logits type: the loss is always float32.
+    auto t_str = get_type_string(in.dtype());
+    std::string kernel_source;
+    kernel_source = metal::utils();
+    kernel_source += metal::cross_entropy();
+    kernel_source +=
+        get_template_definition("block_" + lib_name, "cross_entropy", t_str);
+    kernel_source += get_template_definition(
+        "looped_" + lib_name, "cross_entropy_looped", t_str);
+    kernel_source +=
+        get_template_definition("vjp_" + lib_name, "cross_entropy_vjp", t_str);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
 MTL::ComputePipelineState* get_scan_kernel(
     metal::Device& d,
     const std::string& kernel_name,

mlx/backend/metal/kernels.h

@@ -64,6 +64,11 @@ MTL::ComputePipelineState* get_logsumexp_kernel(
     const std::string& kernel_name,
     const array& out);
 
+MTL::ComputePipelineState* get_cross_entropy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in);
+
 MTL::ComputePipelineState* get_scan_kernel(
     metal::Device& d,
     const std::string& kernel_name,

mlx/backend/metal/kernels/CMakeLists.txt

@@ -127,6 +127,7 @@ if(NOT MLX_METAL_JIT)
   build_kernel(binary binary.h binary_ops.h)
   build_kernel(binary_two binary_two.h)
   build_kernel(copy copy.h)
+  build_kernel(cross_entropy cross_entropy.h)
   build_kernel(fft fft.h fft/radix.h fft/readwrite.h)
   build_kernel(
     reduce