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 new file mode 100644 index 0000000000..2a8035e700 --- /dev/null +++ b/mlx/backend/cuda/cross_entropy.cu @@ -0,0 +1,245 @@ +// 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: 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 + 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); + } + } + 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; + loss[row] = lse_val - static_cast(x[y[row]]); + } +} + +// get loss from the forward +template +__global__ void cross_entropy_vjp( + const T* x, // [M, N] + const int* y, // [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 +) { + 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 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{}); +#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 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 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 + 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; + + 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 0e05e9d19f..27a2849de9 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -163,6 +163,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..04bb4573ac 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -187,6 +187,107 @@ 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)}; + }; + + 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]; + + 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 9226c10eee..0efcfc435a 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 5dacaa605c..773e66b6b2 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,))