From 75c22d6c4f9f1e6103b97b68355e3bae28f8688b Mon Sep 17 00:00:00 2001 From: Divkov575 Date: Sun, 13 Sep 2026 12:11:22 -0400 Subject: [PATCH 1/2] Add fused Metal cross entropy kernel CrossEntropy and CrossEntropyVJP had a fused CUDA kernel but threw "Metal cross entropy NYI" on Metal, where use_fallback always returned true and the op decomposed into logsumexp + take_along_axis. Port the CUDA kernel to Metal, following the logsumexp.h idiom: a block variant for rows up to 4096 and a looped variant above that, both reducing with simd_max/simd_sum and a cross-simdgroup pass through threadgroup memory. The accumulation is in float32 regardless of input type, so half precision inputs no longer materialize a float32 copy of the logits. The vjp writes the gradient into the logits buffer when that buffer is donatable. Each thread reads the target score into a register and the threadgroup barriers before any of them overwrite the row. Fusing also avoids a catastrophic cancellation the decomposed path suffers: forming (max - x_target) before adding log(normalizer) keeps the loss when the logits are large, where computing the full logsumexp first rounds it away. On an M2 Max, forward + backward is 2.4-6.1x faster and peak memory for 4096 x 128256 bfloat16 drops from 4.89 GiB to 1.96 GiB. Forward alone is 2.2-4.1x faster in half precision and roughly at parity in float32, where the fallback needs no cast. Validated against the decomposed fallback for float32/float16/bfloat16 across vocabularies spanning the block/looped boundary, degenerate rows (+inf, fully masked), non-contiguous logits and targets, donated buffers, and with both MLX_METAL_JIT on and off. --- benchmarks/python/cross_entropy_bench.py | 86 +++++++ mlx/backend/metal/CMakeLists.txt | 2 + mlx/backend/metal/cross_entropy.cpp | 157 +++++++++++++ mlx/backend/metal/jit/includes.h | 1 + mlx/backend/metal/jit_kernels.cpp | 22 ++ mlx/backend/metal/kernels.h | 5 + mlx/backend/metal/kernels/CMakeLists.txt | 1 + mlx/backend/metal/kernels/cross_entropy.h | 219 ++++++++++++++++++ mlx/backend/metal/kernels/cross_entropy.metal | 19 ++ mlx/backend/metal/nojit_kernels.cpp | 7 + mlx/backend/metal/primitives.cpp | 22 -- python/tests/test_fast.py | 89 +++++++ 12 files changed, 608 insertions(+), 22 deletions(-) create mode 100644 benchmarks/python/cross_entropy_bench.py create mode 100644 mlx/backend/metal/cross_entropy.cpp create mode 100644 mlx/backend/metal/kernels/cross_entropy.h create mode 100644 mlx/backend/metal/kernels/cross_entropy.metal diff --git a/benchmarks/python/cross_entropy_bench.py b/benchmarks/python/cross_entropy_bench.py new file mode 100644 index 0000000000..cc4667fe51 --- /dev/null +++ b/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") diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index ea4a995ade..598b45fd75 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/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 diff --git a/mlx/backend/metal/cross_entropy.cpp b/mlx/backend/metal/cross_entropy.cpp new file mode 100644 index 0000000000..4b78cbd88f --- /dev/null +++ b/mlx/backend/metal/cross_entropy.cpp @@ -0,0 +1,157 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +#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& inputs, + std::vector& 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& inputs, + std::vector& 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(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 diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index 4fb1be1110..9fe1923e6f 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/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(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 807a6550b7..6dd2cc163a 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/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, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 42888f0a78..d9589e7fc8 100644 --- a/mlx/backend/metal/kernels.h +++ b/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, diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index ecabbda550..7cc90c803c 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/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 diff --git a/mlx/backend/metal/kernels/cross_entropy.h b/mlx/backend/metal/kernels/cross_entropy.h new file mode 100644 index 0000000000..6cacc9cb3a --- /dev/null +++ b/mlx/backend/metal/kernels/cross_entropy.h @@ -0,0 +1,219 @@ +// Copyright © 2026 Apple Inc. + +// Fused logsumexp + gather. +// +// The accumulation is done in float32 regardless of the input type so that +// callers do not need to promote the logits with logits.astype(mx.float32). +// +// For each row: loss = logsumexp(x) - x[target] +// +// One threadgroup handles one row. The logsumexp is accumulated first, then a +// single thread does the gather and writes the loss. + +template +[[kernel]] void cross_entropy( + const device T* x, + const device int* y, + device float* loss, + constant int& axis_size, + 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]]) { + int lid = _lid; + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + AccT ld[N_READS]; + + const device T* row = x + gid * size_t(axis_size); + const device T* in = row + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + ld[i] = AccT(in[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + ld[i] = + ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits::min; + } + } + if (simd_group_id == 0) { + local_max[simd_lane_id] = Limits::min; + local_normalizer[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Get the max + AccT maxval = Limits::finite_min; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < ld[i]) ? ld[i] : maxval; + } + maxval = simd_max(maxval); + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + maxval = simd_max(local_max[simd_lane_id]); + if (simd_lane_id == 0) { + local_max[0] = maxval; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = local_max[0]; + + // Compute exp(x_i - maxval) and store the partial sums in local_normalizer + AccT normalizer = 0; + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(ld[i] - maxval); + } + normalizer = simd_sum(normalizer); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + normalizer = simd_sum(local_normalizer[simd_lane_id]); + if (simd_lane_id == 0) { + // Gather the score of the target class and subtract it from the lse. + AccT gap = maxval - AccT(row[y[gid]]); + loss[gid] = + static_cast(isinf(maxval) ? gap : log(normalizer) + gap); + } + } +} + +template +[[kernel]] void cross_entropy_looped( + const device T* x, + const device int* y, + device float* loss, + constant int& axis_size, + 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]]) { + const device T* row = x + gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + // The threadgroup may hold fewer than SIMD_SIZE simdgroups, so initialize + // every slot: the cross-simdgroup reduction below reads all of them. + if (simd_group_id == 0) { + local_max[simd_lane_id] = Limits::finite_min; + local_normalizer[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Get the max and the normalizer in one go + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(row[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = (offset + i < axis_size) ? AccT(row[offset + i]) + : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = simd_max(local_max[simd_lane_id]); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = simd_sum(local_normalizer[simd_lane_id]); + + if (lid == 0) { + AccT gap = maxval - AccT(row[y[gid]]); + loss[gid] = static_cast(isinf(maxval) ? gap : log(normalizer) + gap); + } +} + +// Gradient of the fused loss. The forward pass already folded the target score +// into the loss, so lse = loss + x[target] and the softmax probabilities come +// back as exp((x_i - x[target]) - loss). +template +[[kernel]] void cross_entropy_vjp( + const device T* x, + const device int* y, + const device float* loss, + const device float* cotan, + device T* grads, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + const device T* row = x + gid * size_t(axis_size); + device T* grads_row = grads + gid * size_t(axis_size); + + int y_n = y[gid]; + AccT g = AccT(cotan[gid]); + AccT loss_n = AccT(loss[gid]); + AccT x_t = AccT(row[y_n]); + + // grads aliases x when the logits buffer is donated, so hold every thread + // here until all of them have read the target score. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(row[offset + i]); + } + for (int i = 0; i < N_READS; i++) { + AccT p = fast::exp((vals[i] - x_t) - loss_n); + vals[i] = g * (p - ((offset + i) == y_n ? AccT(1) : AccT(0))); + } + for (int i = 0; i < N_READS; i++) { + grads_row[offset + i] = T(vals[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + if (offset + i < axis_size) { + AccT v = AccT(row[offset + i]); + AccT p = fast::exp((v - x_t) - loss_n); + grads_row[offset + i] = + T(g * (p - ((offset + i) == y_n ? AccT(1) : AccT(0)))); + } + } + } + } +} diff --git a/mlx/backend/metal/kernels/cross_entropy.metal b/mlx/backend/metal/kernels/cross_entropy.metal new file mode 100644 index 0000000000..d3a2b63848 --- /dev/null +++ b/mlx/backend/metal/kernels/cross_entropy.metal @@ -0,0 +1,19 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +using namespace metal; + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/cross_entropy.h" + +#define instantiate_cross_entropy(name, itype) \ + instantiate_kernel("block_cross_entropy_" #name, cross_entropy, itype) \ + instantiate_kernel("looped_cross_entropy_" #name, cross_entropy_looped, itype) \ + instantiate_kernel("vjp_cross_entropy_" #name, cross_entropy_vjp, itype) \ + +instantiate_cross_entropy(float32, float) +instantiate_cross_entropy(float16, half) +instantiate_cross_entropy(bfloat16, bfloat16_t) // clang-format on diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index f1141e9792..e07a06f441 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -79,6 +79,13 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index d1d0e781cc..68ba5fc2f5 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -215,26 +215,4 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } -namespace fast { - -// There is no fused Metal cross entropy kernel yet -bool CrossEntropy::use_fallback(Stream s) { - return true; -} - -void CrossEntropy::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error("[CrossEntropy::eval_gpu] Metal cross entropy NYI."); -} - -void CrossEntropyVJP::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error( - "[CrossEntropyVJP::eval_gpu] Metal cross entropy NYI."); -} - -} // namespace fast - } // namespace mlx::core diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 200781d372..0d48ef1be5 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -548,6 +548,95 @@ def cross_entropy_ref(logits, targets): self.assertEqual(out.shape, targets.shape) self.assertLess(mx.abs(out - expected).max().item(), tolerances[dtype]) + def test_cross_entropy_degenerate_rows(self): + def cross_entropy_ref(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 + ) + + targets = mx.array([0, 1]) + for name, logits in [ + ("has +inf", mx.array([[1.0, float("inf"), 2.0], [1.0, 2.0, 3.0]])), + ( + "masked", + mx.array([[0.5, float("-inf"), 2.0], [float("-inf"), 1.0, 3.0]]), + ), + ]: + out = mx.fast.cross_entropy(logits, targets) + expected = cross_entropy_ref(logits, targets) + self.assertTrue(mx.allclose(out, expected).item(), msg=name) + + # An entirely masked row has no finite logsumexp, so the loss is nan + # both fused and unfused. + all_masked = mx.array([[float("-inf")] * 3, [1.0, 2.0, 3.0]]) + out = mx.fast.cross_entropy(all_masked, targets) + self.assertTrue(mx.isnan(out[0]).item()) + self.assertTrue(mx.isfinite(out[1]).item()) + + def test_cross_entropy_catastrophic_cancellation(self): + # Every logit in the row is equal, so the loss is exactly log(V). The + # fused kernel forms (max - x_target) before adding log(normalizer) and + # keeps that value, where computing logsumexp first would round it away. + # Only the fused kernel has this property, so there is nothing to check + # on the cpu, which always takes the decomposed fallback. + if mx.default_device() == mx.cpu: + self.skipTest("cross entropy is not fused on the cpu") + V = 4096 + logits = mx.full((2, V), 1e30) + targets = mx.array([0, 5]) + out = mx.fast.cross_entropy(logits, targets) + self.assertTrue(mx.allclose(out, mx.full((2,), math.log(V)), atol=1e-4).item()) + + def test_cross_entropy_non_contiguous(self): + def cross_entropy_ref(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 + ) + + logits = mx.random.normal(shape=(1000, 8)).T + targets = mx.random.randint(0, 1000, shape=(8,)) + self.assertTrue( + mx.allclose( + mx.fast.cross_entropy(logits, targets), + cross_entropy_ref(logits, targets), + ).item() + ) + + logits = mx.random.normal(shape=(4, 2048))[:, ::2] + targets = mx.random.randint(0, 1024, shape=(4,)) + self.assertTrue( + mx.allclose( + mx.fast.cross_entropy(logits, targets), + cross_entropy_ref(logits, targets), + ).item() + ) + + def test_cross_entropy_grad_donated_input(self): + # The logits are produced inside the graph, so the vjp may write the + # gradient directly into that buffer. + def ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits, axis=-1) - score + + for V in [64, 4096, 32000]: + x = mx.random.normal(shape=(3, V), scale=2.0) + targets = mx.random.randint(0, V, shape=(3,)) + g1 = mx.grad(lambda a, y: ref(a * 2.0 + 1.0, y).sum(), argnums=0)( + x, targets + ) + g2 = mx.grad( + lambda a, y: mx.fast.cross_entropy(a * 2.0 + 1.0, y).sum(), argnums=0 + )(x, targets) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-5) + def test_cross_entropy_shape_checks(self): logits = mx.random.normal(shape=(4, 16)) with self.assertRaises(ValueError): From 996269bb8b25ef8dcc87552ab623bf5199917c33 Mon Sep 17 00:00:00 2001 From: Divkov575 Date: Sun, 13 Sep 2026 15:22:44 -0400 Subject: [PATCH 2/2] Route nn.losses.cross_entropy to the fused kernel on Metal The use_fast gate checked mx.cuda.is_available(), so the fused primitive was unreachable from nn.losses.cross_entropy on Metal even once the Metal kernel existed. Both GPU backends implement it now, so gate on the device rather than on the backend. The fast branch does not apply the max shift the fallback branch does, because the fused kernel shifts by the row max internally. The Metal kernel has that property, so the existing shift invariance tests hold on Metal: the loss stays log(n) for equal logits from 1e0 to 1e20, and matches the cpu at every offset. Add a test that the two devices agree. In float32 they agree exactly. In half precision they are allowed to differ, because the fused kernel accumulates in float32 while the fallback reduces in the dtype of the logits, so the test asserts the stronger property instead: the fused result is never further from the float32 answer. It is a regression guard rather than a demonstration of this change, since routing to a more accurate kernel cannot make an equivalence test fail. Also correct the docstring, which said Metal reduces in the dtype of the logits, and the dtype in the half precision example, which is the dtype of the logits rather than float32. --- python/mlx/nn/losses.py | 13 ++++---- python/tests/test_losses.py | 62 +++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 6 deletions(-) diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index b98d2765d6..8df7c29cd7 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -64,15 +64,15 @@ def cross_entropy( >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) >>> - >>> # Half precision logits with class indices as targets. On CUDA a + >>> # Half precision logits with class indices as targets. On a GPU a >>> # fused kernel accumulates the reduction in float32: >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16) >>> targets = mx.array([0, 1]) >>> nn.losses.cross_entropy(logits, targets) - array([0.0485873, 0.0485873], dtype=float32) + array([0.0485873, 0.0485873], dtype=bfloat16) >>> - >>> # Metal and the CPU reduce in the dtype of the logits, so upcast - >>> # them to get the same accuracy: + >>> # The cpu reduces in the dtype of the logits, so upcast them to get + >>> # the same accuracy: >>> nn.losses.cross_entropy(logits.astype(mx.float32), targets) array([0.0485873, 0.0485873], dtype=float32) """ @@ -95,9 +95,10 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) + # Both GPU backends implement the fused kernel, which shifts by the row max + # internally, so the explicit shift below is only needed off the fast path. use_fast = ( - mx.cuda.is_available() - and mx.default_device() == mx.gpu + mx.default_device() == mx.gpu and not targets_as_probs and label_smoothing == 0 and axis in (-1, logits.ndim - 1) diff --git a/python/tests/test_losses.py b/python/tests/test_losses.py index 96174f3c89..ffd86ea5c2 100644 --- a/python/tests/test_losses.py +++ b/python/tests/test_losses.py @@ -9,6 +9,68 @@ class TestLosses(mlx_tests.MLXTestCase): + def test_cross_entropy_gpu_matches_cpu(self): + # On a gpu the class index path is handed to the fused kernel, which + # shifts by the row max itself instead of using the explicit shift the + # fallback applies. The two have to agree, including for the arguments + # that fall back on either device. + if mx.default_device() == mx.cpu: + self.skipTest("needs a gpu to compare against the cpu") + + # cross_entropy picks the fast path off the default device, so switch + # the device rather than passing a stream. + def both_devices(*args, **kwargs): + gpu = nn.losses.cross_entropy(*args, **kwargs) + mx.set_default_device(mx.cpu) + try: + cpu = nn.losses.cross_entropy(*args, **kwargs) + mx.eval(cpu) + finally: + mx.set_default_device(mx.gpu) + mx.eval(gpu) + return gpu, cpu + + for V in [2, 7, 4096, 4097]: + logits = mx.random.normal(shape=(4, V), scale=3.0) + targets = mx.random.randint(0, V, shape=(4,)) + gpu, cpu = both_devices(logits, targets, reduction="none") + self.assertEqual(gpu.dtype, cpu.dtype) + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=f"V={V}") + + # In half precision the two are allowed to differ, because the + # fused kernel accumulates in float32 while the fallback reduces in + # the dtype of the logits. Hold it to the stronger property: the + # fused result is never further from the float32 answer. + for dtype in [mx.float16, mx.bfloat16]: + half = logits.astype(dtype) + reference = nn.losses.cross_entropy( + half.astype(mx.float32), targets, reduction="none" + ) + gpu, cpu = both_devices(half, targets, reduction="none") + self.assertEqual(gpu.dtype, dtype) + gpu_err = mx.abs(gpu.astype(mx.float32) - reference).max().item() + cpu_err = mx.abs(cpu.astype(mx.float32) - reference).max().item() + self.assertLessEqual(gpu_err, cpu_err + 1e-6, msg=f"V={V} {dtype}") + self.assertLess(gpu_err, 0.2, msg=f"V={V} {dtype}") + + # A large shared offset is where the fused and the decomposed paths + # could most easily disagree. + base = mx.array([[2.0, -1.0]]) + for offset in [0.0, 1e4, 1e6]: + gpu, cpu = both_devices(base + offset, mx.array([0]), reduction="none") + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=f"offset={offset}") + + # Arguments the fast path declines, so both devices decompose. + logits = mx.random.normal(shape=(4, 32)) + targets = mx.random.randint(0, 32, shape=(4,)) + for kwargs in [ + {"label_smoothing": 0.1}, + {"weights": mx.random.uniform(shape=(4,))}, + {"reduction": "mean"}, + ]: + gpu, cpu = both_devices(logits, targets, **kwargs) + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=str(kwargs)) + def test_cross_entropy(self): # No weights, no label smoothing logits = mx.array([[0.0, -float("inf")], [-float("inf"), 0.0]])