From 5959cf688eb66efb79b279521d60480aef2c653c Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 29 Jul 2026 14:46:45 +0200 Subject: [PATCH 1/4] cross entropy forward backward --- mlx/backend/cuda/cross_entropy.cu | 131 ++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 mlx/backend/cuda/cross_entropy.cu diff --git a/mlx/backend/cuda/cross_entropy.cu b/mlx/backend/cuda/cross_entropy.cu new file mode 100644 index 0000000000..d73b459586 --- /dev/null +++ b/mlx/backend/cuda/cross_entropy.cu @@ -0,0 +1,131 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/device/cast_op.cuh" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" + +#include +#include +#include + +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +// fused together logsumexp + gather +// cast to float32 inside the kernel +// to avoid logits.astype(mx.float32) +// for each row: logsumexp(x) - x_t +// first we accumulate logsumexp, then we do a gather +template +__global__ void cross_entropy( + const T* x, // [M, N] + const int* y, // [M,] + float* loss, // [M,] <- will be always in fp32 lse - x + float* lse, // logsumexp for backward [M,] in fp32 + int axis_size // N +) { + cg::greater max_op; + cg::plus plus_op; + + float prevmax; + float curmax = Limits::finite_min(); + float normalizer = 0; + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + x += grid.block_rank() * axis_size; // offset input + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto vals = load_vector(x, index, axis_size, Limits::min()); + prevmax = curmax; +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + curmax = max_op(curmax, static_cast(vals[i])); + } + // scale already accumulated normiliser + normalizer = normalizer * __expf(prevmax - curmax); + // add vals scaled by curmax +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + normalizer += __expf(static_cast(vals[i]) - curmax); + } + } + // here every thread has it's own normiliser : N_READS values with stride 32 + // and max for all this values we need to exchange it with other threads 1) in + // a warp 2) in a block first reduce in a warp + prevmax = curmax; + curmax = cg::reduce(warp, curmax, max_op); + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // second reduce in a block + __shared__ float warp_max[WARP_SIZE]; + __shared__ float warp_normaliser[WARP_SIZE]; + + if (warp.thread_rank() == 0) { + warp_max[warp.meta_group_rank()] = curmax; + warp_normaliser[warp.meta_group_rank()] = normalizer; + } + block.sync(); + bool is_valid = warp.thread_rank() < warp.meta_group_size(); + curmax = + is_valid ? warp_max[warp.thread_rank()] : Limits::finite_min(); + prevmax = curmax; + curmax = + cg::reduce(warp, curmax, max_op); // max within a block (global row max) + normalizer = is_valid ? warp_normaliser[warp.thread_rank()] : 0.0f; + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // gather and writing the output: + auto row = grid.block_rank(); + if (block.thread_rank() == 0) { + float lse_val = isinf(curmax) ? curmax : log(normalizer) + curmax; + lse[row] = lse_val; + loss[row] = lse_val - static_cast(x[y[row]]); + } +} + +// get lse from the forward, i think we will assume non negative indc +template +__global__ void cross_entropy_vjp( + const T* x, // [M, N] + const int* y, // [M,] + const float* lse, // [M,] + const float* gy, // cotangent [M,] + T* grads, // [M, N] lse is accumulated in float, x is casted to float + int axis_size // N +) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto row = grid.block_rank(); + + x += row * axis_size; // offset input + grads += row * axis_size; // offset output + auto lse_n = lse[row]; // logsumexp + auto y_n = y[row]; // target index [0, N) + auto g = gy[row]; // cotangent + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); // [0, N) + auto vals = load_vector(x, index, axis_size, T{}); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + // lse(x) >= max(x), so the exponent is <= 0 and __expf is safe + int col = index * N_READS + i; + float val = __expf(static_cast(vals[i]) - lse_n); + vals[i] = static_cast(g * (val - (col == y_n ? 1.0f : 0.0f))); + } + store_vector(grads, index, vals, axis_size); + } +} +} // namespace cu + +} // namespace mlx::core From 219d489f3659b82d6f1aa5482fac4e373d7646f4 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 29 Jul 2026 17:16:49 +0200 Subject: [PATCH 2/4] fast primitive --- docs/src/python/fast.rst | 1 + mlx/backend/cuda/CMakeLists.txt | 1 + mlx/backend/cuda/cross_entropy.cu | 112 ++++++++++++++++++++++++++++-- mlx/backend/metal/primitives.cpp | 21 ++++++ mlx/backend/no_gpu/primitives.cpp | 2 + mlx/fast.cpp | 105 ++++++++++++++++++++++++++++ mlx/fast.h | 4 ++ mlx/fast_primitives.h | 61 ++++++++++++++++ python/mlx/nn/losses.py | 44 +++++++----- python/src/fast.cpp | 29 ++++++++ python/tests/test_fast.py | 61 ++++++++++++++++ 11 files changed, 417 insertions(+), 24 deletions(-) diff --git a/docs/src/python/fast.rst b/docs/src/python/fast.rst index affeb444f8..c930c7bb4a 100644 --- a/docs/src/python/fast.rst +++ b/docs/src/python/fast.rst @@ -10,6 +10,7 @@ Fast rms_norm layer_norm + cross_entropy rope scaled_dot_product_attention metal_kernel diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 065220e24e..0f3c21eeb1 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -19,6 +19,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cu ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp diff --git a/mlx/backend/cuda/cross_entropy.cu b/mlx/backend/cuda/cross_entropy.cu index d73b459586..b64f091cfb 100644 --- a/mlx/backend/cuda/cross_entropy.cu +++ b/mlx/backend/cuda/cross_entropy.cu @@ -22,14 +22,13 @@ namespace cg = cooperative_groups; // fused together logsumexp + gather // cast to float32 inside the kernel // to avoid logits.astype(mx.float32) -// for each row: logsumexp(x) - x_t +// for each row: loss = logsumexp(x) - x_t // first we accumulate logsumexp, then we do a gather template __global__ void cross_entropy( const T* x, // [M, N] const int* y, // [M,] float* loss, // [M,] <- will be always in fp32 lse - x - float* lse, // logsumexp for backward [M,] in fp32 int axis_size // N ) { cg::greater max_op; @@ -89,17 +88,16 @@ __global__ void cross_entropy( auto row = grid.block_rank(); if (block.thread_rank() == 0) { float lse_val = isinf(curmax) ? curmax : log(normalizer) + curmax; - lse[row] = lse_val; loss[row] = lse_val - static_cast(x[y[row]]); } } -// get lse from the forward, i think we will assume non negative indc +// get loss from the forward template __global__ void cross_entropy_vjp( const T* x, // [M, N] const int* y, // [M,] - const float* lse, // [M,] + const float* loss, // [M,] const float* gy, // cotangent [M,] T* grads, // [M, N] lse is accumulated in float, x is casted to float int axis_size // N @@ -110,9 +108,9 @@ __global__ void cross_entropy_vjp( x += row * axis_size; // offset input grads += row * axis_size; // offset output - auto lse_n = lse[row]; // logsumexp auto y_n = y[row]; // target index [0, N) auto g = gy[row]; // cotangent + auto lse_n = loss[row] + static_cast(x[y_n]); for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { auto index = r * BLOCK_DIM + block.thread_rank(); // [0, N) auto vals = load_vector(x, index, axis_size, T{}); @@ -128,4 +126,106 @@ __global__ void cross_entropy_vjp( } } // namespace cu +namespace fast { + +bool CrossEntropy::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropy::eval_gpu"); + assert(inputs.size() == 2); // logits and target + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + 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(cu::malloc_async(out.nbytes(), encoder)); // [n_rows] in fp32 + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(out), + axis_size); + }); + }); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropyVJP::eval_gpu"); + assert(inputs.size() == 4); // logits, target, loss, cotangent + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + 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,] + auto loss = ensure_row_contiguous(inputs[2]); // [n_rows,] fp32 + auto cotan = ensure_row_contiguous(inputs[3]); // [n_rows,] fp32 + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows, V] + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_input_array(loss); + encoder.set_input_array(cotan); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy_vjp", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy_vjp; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(loss), + gpu_ptr(cotan), + gpu_ptr(out), + axis_size); + }); + }); +} + +} // namespace fast + } // namespace mlx::core diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index d5bbf797e4..5c6c856a37 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -239,4 +239,25 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } +namespace fast { + +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/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 4819ed2724..3cc8dad3ec 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -162,6 +162,8 @@ NO_GPU(View) NO_GPU(MaskedScatter) namespace fast { +NO_GPU_USE_FALLBACK(CrossEntropy) +NO_GPU_MULTI(CrossEntropyVJP) NO_GPU_USE_FALLBACK(LayerNorm) NO_GPU_MULTI(LayerNormVJP) NO_GPU_USE_FALLBACK(RMSNorm) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..925e6e030e 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -187,6 +187,111 @@ bool RMSNormVJP::is_equivalent(const Primitive& other) const { return eps_ == a_other.eps_; } +array cross_entropy( + const array& logits, + const array& targets, + StreamOrDevice s_ /* = {} */) { + if (logits.ndim() < 1) { + throw std::invalid_argument( + "[cross_entropy] logits must have at least 1 dimension but got input " + "with 0 dimensions."); + } + auto expected = logits.shape(); + expected.pop_back(); + if (targets.shape() != expected) { + std::ostringstream msg; + msg << "[cross_entropy] targets shape " << targets.shape() + << " does not match logits shape " << logits.shape() + << " with the last axis removed."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(logits.dtype(), floating)) { + std::ostringstream msg; + msg << "[cross_entropy] Received unsupported logits type " << logits.dtype() + << "."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(targets.dtype(), integer)) { + std::ostringstream msg; + msg << "[cross_entropy] targets must be integer class indices but got " + << targets.dtype() << "."; + throw std::invalid_argument(msg.str()); + } + + auto s = to_stream(s_); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto loss = subtract(logsumexp(x, -1, /* keepdims= */ false, s), score, s); + return std::vector{astype(loss, float32, s)}; + }; + + // The kernel indexes the targets directly, so normalize them here to keep + // the fused path and the fallback (which goes through take_along_axis) in + // agreement on negative indices. + auto passed_targets = astype(targets, int32, s); + + if (!CrossEntropy::use_fallback(s)) { + return array( + expected, + float32, + std::make_shared(s, fallback), + {logits, passed_targets}); + } + return fallback({logits, passed_targets})[0]; +} + +std::vector CrossEntropy::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + assert(primals.size() == 2); + assert(outputs.size() == 1); + assert(cotangents.size() == 1); + + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[cross_entropy] Cannot differentiate with respect to the targets."); + } + } + + auto s = stream(); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto& loss = inputs[2]; + auto& g = inputs[3]; + + // loss = lse - x_t, so lse is recovered without saving it. + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto lse = add(loss, astype(score, float32, s), s); + auto p = + exp(subtract(astype(x, float32, s), expand_dims(lse, -1, s), s), s); + Shape class_shape(x.ndim(), 1); + class_shape.back() = x.shape(-1); + auto onehot = astype( + equal( + expand_dims(y, -1, s), + reshape(arange(x.shape(-1), y.dtype(), s), class_shape, s), + s), + float32, + s); + auto gx = multiply(expand_dims(g, -1, s), subtract(p, onehot, s), s); + return std::vector{astype(gx, x.dtype(), s)}; + }; + + return {array( + primals[0].shape(), + primals[0].dtype(), + std::make_shared(s, fallback), + {primals[0], primals[1], outputs[0], cotangents[0]})}; +} + array layer_norm( const array& x, const std::optional& weight, diff --git a/mlx/fast.h b/mlx/fast.h index 934fadc2b7..1922545649 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -24,6 +24,10 @@ MLX_API array layer_norm( float eps, StreamOrDevice s = {}); +/** Fused cross entropy with class indices as targets. */ +MLX_API array +cross_entropy(const array& logits, const array& targets, StreamOrDevice s = {}); + MLX_API array rope( const array& x, int dims, diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 0d2f861045..3f61bb238e 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -97,6 +97,67 @@ class RMSNormVJP : public Custom { float eps_; }; +// loss is always fp32 and the logits never have to be upcast in the graph. +class CrossEntropy : public Custom { + public: + CrossEntropy( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback(Stream stream); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + + DEFINE_NAME(CrossEntropy) + bool is_equivalent(const Primitive& other) const override { + return true; + } + std::vector output_shapes(const std::vector& inputs) override { + return {inputs[1].shape()}; + } + + auto state() const { + return std::monostate{}; + } +}; + +class CrossEntropyVJP : public Custom { + public: + CrossEntropyVJP( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + DEFINE_NAME(CrossEntropyVJP) + bool is_equivalent(const Primitive& other) const override { + return true; + } + DEFINE_INPUT_OUTPUT_SHAPE() + + auto state() const { + return std::monostate{}; + } +}; + class LayerNorm : public Custom { public: LayerNorm( diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 85b11255bb..dffb37e694 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -83,26 +83,34 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) - if targets_as_probs: - score = mx.sum(logits * targets, axis=axis) - else: - score = mx.take_along_axis(logits, mx.expand_dims(targets, axis), axis).squeeze( - axis - ) - - logsumexp_logits = mx.logsumexp(logits, axis=axis) - if label_smoothing > 0: - # Adjust the true class score with label smoothing - adjusted_score = (1 - label_smoothing) * score - - # Calculate the mean logit across the classes for smoothed loss - mean_logits = logits.mean(axis=axis) - smoothed_loss = -mean_logits * label_smoothing + use_fast = ( + not targets_as_probs + and label_smoothing == 0 + and axis in (-1, logits.ndim - 1) + and mx.issubdtype(logits.dtype, mx.floating) + and mx.issubdtype(targets.dtype, mx.integer) + ) - # Combine the adjusted score and smoothed loss with the logsumexp logits - loss = logsumexp_logits - adjusted_score + smoothed_loss + if use_fast: + loss = mx.fast.cross_entropy(logits, targets).astype(logits.dtype) else: - loss = logsumexp_logits - score + if targets_as_probs: + score = mx.sum(logits * targets, axis=axis) + else: + score = mx.take_along_axis( + logits, mx.expand_dims(targets, axis), axis + ).squeeze(axis) + + logsumexp_logits = mx.logsumexp(logits, axis=axis) + if label_smoothing > 0: + adjusted_score = (1 - label_smoothing) * score + + mean_logits = logits.mean(axis=axis) + smoothed_loss = -mean_logits * label_smoothing + + loss = logsumexp_logits - adjusted_score + smoothed_loss + else: + loss = logsumexp_logits - score # Apply weights if provided if weights is not None: diff --git a/python/src/fast.cpp b/python/src/fast.cpp index cd30b0bacd..1566733c38 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -174,6 +174,35 @@ void init_fast(nb::module_& parent_module) { array: The output array. )pbdoc"); + m.def( + "cross_entropy", + &mx::fast::cross_entropy, + "logits"_a, + "targets"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def cross_entropy(logits: array, targets: array, *, stream: Union[None, Stream, Device] = None) -> array"), + R"pbdoc( + Cross entropy loss with class indices as targets. + + Computes ``logsumexp(logits, axis=-1) - logits[..., target]`` in a + single fused kernel. The log-sum-exp is accumulated in float32 + regardless of the dtype of ``logits``, so the returned loss is always + float32 and the logits do not need to be upcast beforehand. + + Args: + logits (array): The unnormalized logits. The loss is computed over + the last axis. + targets (array): Class indices. The shape should match the shape of + ``logits`` with the last axis removed. The indices must be in + ``[0, logits.shape[-1])``. + + Returns: + array: The per-element loss in float32, with the shape of + ``targets``. + )pbdoc"); + m.def( "rope", [](const mx::array& a, diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 8343c45941..1ba3d17ba7 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -525,6 +525,67 @@ def inner(x, w, y): self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + def test_cross_entropy(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 + ) + + tolerances = {mx.float32: 1e-5, mx.float16: 5e-3, mx.bfloat16: 3e-2} + + for V in [7, 32, 128, 255, 256, 1000, 4096, 8192]: + for dtype in [mx.float32, mx.float16, mx.bfloat16]: + logits = (mx.random.normal(shape=(4, 7, V), scale=3.0) * 2).astype( + dtype + ) + targets = mx.random.randint(0, V, shape=(4, 7)) + expected = cross_entropy_ref(logits, targets) + out = mx.fast.cross_entropy(logits, targets) + self.assertEqual(out.dtype, mx.float32) + self.assertEqual(out.shape, targets.shape) + self.assertLess(mx.abs(out - expected).max().item(), tolerances[dtype]) + + def test_cross_entropy_shape_checks(self): + logits = mx.random.normal(shape=(4, 16)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((5,), mx.int32)) + with self.assertRaises(ValueError): + # Probability targets are not supported by the fused op. + mx.fast.cross_entropy(logits, mx.zeros((4, 16), mx.int32)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((4,), mx.float32)) + + def test_cross_entropy_grad(self): + 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 + + f1 = lambda x, y: ref(x, y).mean() + f2 = lambda x, y: mx.fast.cross_entropy(x, y).mean() + + for V in [7, 128, 1000, 4096]: + logits = mx.random.normal(shape=(4, 7, V), scale=2.0) + targets = mx.random.randint(0, V, shape=(4, 7)) + g1 = mx.grad(f1, argnums=0)(logits, targets) + g2 = mx.grad(f2, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + + w = mx.random.uniform(shape=(4, 7)) + f3 = lambda x, y: (ref(x, y) * w).sum() + f4 = lambda x, y: (mx.fast.cross_entropy(x, y) * w).sum() + logits = mx.random.normal(shape=(4, 7, 512), scale=2.0) + targets = mx.random.randint(0, 512, shape=(4, 7)) + g1 = mx.grad(f3, argnums=0)(logits, targets) + g2 = mx.grad(f4, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + def test_layer_norm_dim_check(self): with self.assertRaises(ValueError): weight = mx.ones((129,)) From 9c8145d0ca51c91dd135db4aaf6a95125add8f91 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 29 Jul 2026 17:44:20 +0200 Subject: [PATCH 3/4] cleaned comments --- mlx/backend/cuda/cross_entropy.cu | 3 --- mlx/fast.cpp | 4 ---- 2 files changed, 7 deletions(-) diff --git a/mlx/backend/cuda/cross_entropy.cu b/mlx/backend/cuda/cross_entropy.cu index b64f091cfb..ddc1ecdf55 100644 --- a/mlx/backend/cuda/cross_entropy.cu +++ b/mlx/backend/cuda/cross_entropy.cu @@ -59,9 +59,6 @@ __global__ void cross_entropy( normalizer += __expf(static_cast(vals[i]) - curmax); } } - // here every thread has it's own normiliser : N_READS values with stride 32 - // and max for all this values we need to exchange it with other threads 1) in - // a warp 2) in a block first reduce in a warp prevmax = curmax; curmax = cg::reduce(warp, curmax, max_op); normalizer = normalizer * __expf(prevmax - curmax); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 925e6e030e..04bb4573ac 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -228,9 +228,6 @@ array cross_entropy( return std::vector{astype(loss, float32, s)}; }; - // The kernel indexes the targets directly, so normalize them here to keep - // the fused path and the fallback (which goes through take_along_axis) in - // agreement on negative indices. auto passed_targets = astype(targets, int32, s); if (!CrossEntropy::use_fallback(s)) { @@ -266,7 +263,6 @@ std::vector CrossEntropy::vjp( auto& loss = inputs[2]; auto& g = inputs[3]; - // loss = lse - x_t, so lse is recovered without saving it. auto score = squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); auto lse = add(loss, astype(score, float32, s), s); From 1cdd0c690423ecf0a2f189ba6217c0000706b3eb Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 29 Jul 2026 21:50:17 +0200 Subject: [PATCH 4/4] reuse logits buffer --- mlx/backend/cuda/cross_entropy.cu | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cuda/cross_entropy.cu b/mlx/backend/cuda/cross_entropy.cu index ddc1ecdf55..2a8035e700 100644 --- a/mlx/backend/cuda/cross_entropy.cu +++ b/mlx/backend/cuda/cross_entropy.cu @@ -108,6 +108,7 @@ __global__ void cross_entropy_vjp( auto y_n = y[row]; // target index [0, N) auto g = gy[row]; // cotangent auto lse_n = loss[row] + static_cast(x[y_n]); + block.sync(); for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { auto index = r * BLOCK_DIM + block.thread_rank(); // [0, N) auto vals = load_vector(x, index, axis_size, T{}); @@ -190,11 +191,27 @@ void CrossEntropyVJP::eval_gpu( return x_copy; } }; - auto in = ensure_row_contiguous(inputs[0]); // [n_rows, V] + + auto check_input = [&s](const array& x, bool& copied) { + if (x.flags().row_contiguous) { + copied = false; + return x; + } + copied = true; + return contiguous_copy_gpu(x, s); + }; + bool donate_x = inputs[0].is_donatable(); + bool copied; + auto in = check_input(inputs[0], copied); // [n_rows, V] + donate_x |= copied; 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 - out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows, V] + if (donate_x) { + out.copy_shared_buffer(in); + } else { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows, V] + } int axis_size = in.shape().back(); int n_rows = in.data_size() / axis_size;