From b624b0f4e6eed31e7e0abc791ca10e8366de0832 Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 24 Sep 2026 17:33:42 +0200 Subject: [PATCH 1/4] Add radix partition kernel --- mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/jit/includes.h | 1 + mlx/backend/metal/jit_kernels.cpp | 22 + mlx/backend/metal/kernels.h | 8 + mlx/backend/metal/kernels/CMakeLists.txt | 1 + mlx/backend/metal/kernels/partition.h | 580 ++++++++++++++++++++++ mlx/backend/metal/kernels/partition.metal | 31 ++ mlx/backend/metal/nojit_kernels.cpp | 10 + mlx/backend/metal/sort.cpp | 67 ++- 9 files changed, 717 insertions(+), 4 deletions(-) create mode 100644 mlx/backend/metal/kernels/partition.h create mode 100644 mlx/backend/metal/kernels/partition.metal diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index c86cfdeb07..9f29a72c42 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -48,6 +48,7 @@ if(MLX_METAL_JIT) make_jit_source(softmax) make_jit_source(scan) make_jit_source(sort) + make_jit_source(partition) make_jit_source(searchsorted kernels/sort.h) make_jit_source( reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index f5a5fe45ef..2f4e997d78 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -31,6 +31,7 @@ const char* scan(); const char* scatter_axis(); const char* softmax(); const char* sort(); +const char* partition(); const char* searchsorted(); const char* reduce(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 1c052e6314..966fb652ad 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -436,6 +436,28 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_partition_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + const array& out, + bool arg_partition, + int bn) { + auto lib = d.get_library(kernel_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::utils() << metal::partition(); + kernel_source << get_template_definition( + kernel_name, + "radix_partition", + get_type_string(in.dtype()), + get_type_string(out.dtype()), + arg_partition ? "true" : "false", + bn); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_searchsorted_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 28ac3ba940..67bb859d4f 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -81,6 +81,14 @@ MTL::ComputePipelineState* get_sort_kernel( int bn, int tn); +MTL::ComputePipelineState* get_partition_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + const array& out, + bool arg_partition, + int bn); + MTL::ComputePipelineState* get_searchsorted_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 00f90ac862..5b2af922af 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -145,6 +145,7 @@ if(NOT MLX_METAL_JIT) build_kernel(logsumexp logsumexp.h) build_kernel(searchsorted searchsorted.h sort.h) build_kernel(sort sort.h) + build_kernel(partition partition.h) build_kernel(ternary ternary.h ternary_ops.h) build_kernel(unary unary.h unary_ops.h) build_kernel(steel/conv/kernels/steel_conv ${STEEL_HEADERS}) diff --git a/mlx/backend/metal/kernels/partition.h b/mlx/backend/metal/kernels/partition.h new file mode 100644 index 0000000000..df642fe121 --- /dev/null +++ b/mlx/backend/metal/kernels/partition.h @@ -0,0 +1,580 @@ +// Copyright © 2026 Apple Inc. + +// Radix select partition: one threadgroup per row finds the k-th key with a +// histogram pass per digit, then writes the row as below | equal | above it. + +using namespace metal; + +static constant constexpr int radix_n_reads = 8; +// Keys a thread loads together in a select pass, so the loads are in flight +// at the same time +static constant constexpr int radix_n_loads = 16; +static constant constexpr int radix_simd_size = 32; +static constant constexpr int radix_threshold_sides = 3; + +template +struct RadixLayout { + static constant constexpr int simd_groups = BLOCK_THREADS / radix_simd_size; + // Wide digits need fewer passes and leave fewer matching keys per pass, + // 256 threads spread the 2048 bins at no extra cost per thread + static constant constexpr int digit_bits = BLOCK_THREADS >= 256 ? 11 : 8; + static constant constexpr int bins = 1 << digit_bits; + static constant constexpr int bins_per_thread = bins / BLOCK_THREADS; + // Scan scratch per value: one sum per simdgroup and the total after them + static constant constexpr int scan_size = simd_groups + 1; + // Candidates cache, 8 keys per thread so small groups keep more rows per core + static constant constexpr int candidates_cache_size = 8 * BLOCK_THREADS; + + // Digit passes over a key of `key_bits`, the last digit may be narrower + static METAL_FUNC constexpr int passes(int key_bits) { + return (key_bits + digit_bits - 1) / digit_bits; + } + static METAL_FUNC constexpr int shift(int key_bits, int pass) { + int shift = key_bits - digit_bits * (pass + 1); + return shift > 0 ? shift : 0; + } + + static_assert( + BLOCK_THREADS % radix_simd_size == 0, + "BLOCK_THREADS must be a multiple of the simd size"); + static_assert( + bins % BLOCK_THREADS == 0, + "BLOCK_THREADS must divide the number of bins"); +}; + +/////////////////////////////////////////////////////////////////////////////// +// Radix keys +/////////////////////////////////////////////////////////////////////////////// + +// Unsigned key of the same width as T, so the digits cover the whole value. +// Signed ints and positive floats flip the sign bit, negative floats flip every +// bit, NaN gets the largest key so it sorts last. +template +struct RadixKeyTraits { + using KeyT = metal::conditional_t< + sizeof(T) == 1, + uint8_t, + metal::conditional_t< + sizeof(T) == 2, + uint16_t, + metal::conditional_t>>; + static METAL_FUNC KeyT to_key(T v) { + constexpr KeyT sign_bit = KeyT(1) << (8 * sizeof(T) - 1); + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(v)) { + return KeyT(~KeyT(0)); + } + KeyT bits = as_type(v); + return bits ^ ((bits & sign_bit) ? KeyT(~KeyT(0)) : sign_bit); + } else if constexpr (metal::is_signed_v) { + return as_type(v) ^ sign_bit; + } else { + return v; + } + } +}; + +// Complex keys order by the real part, then the imaginary part, and a NaN in +// either part sorts last +template <> +struct RadixKeyTraits { + using KeyT = uint64_t; + static METAL_FUNC KeyT to_key(complex64_t v) { + if (metal::isnan(v.real) || metal::isnan(v.imag)) { + return ~KeyT(0); + } + return (KeyT(RadixKeyTraits::to_key(v.real)) + << (8 * sizeof(float))) | + RadixKeyTraits::to_key(v.imag); + } +}; + +// Side of the threshold a key goes to: 0 below, 1 equal, 2 above +template +METAL_FUNC int radix_threshold_side(KeyT key, KeyT threshold) { + return key < threshold ? 0 : (key == threshold ? 1 : 2); +} + +/////////////////////////////////////////////////////////////////////////////// +// Threadgroup scan +/////////////////////////////////////////////////////////////////////////////// + +// Exclusive scan of N values per thread over the threadgroup. `values` become +// the sums before this thread, `totals` the sums over all threads. +template +METAL_FUNC void threadgroup_exclusive_sum( + thread uint* values, + thread uint* totals, + uint simd_lane_id, + uint simd_group_id, + threadgroup uint* sums) { + constexpr int simd_groups = RadixLayout::simd_groups; + constexpr int scan_size = RadixLayout::scan_size; + constexpr int last_lane = radix_simd_size - 1; + + // Scan inside each simdgroup, the last lane stores the simdgroup sum + uint prev[N]; + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < N; c++) { + prev[c] = simd_prefix_exclusive_sum(values[c]); + if (simd_lane_id == last_lane) { + sums[c * scan_size + simd_group_id] = prev[c] + values[c]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Scan the simdgroup sums, the total lands right after them + if (simd_group_id == 0) { + bool valid = simd_lane_id < simd_groups; + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < N; c++) { + threadgroup uint* group_sums = sums + c * scan_size; + uint sum = valid ? group_sums[simd_lane_id] : 0; + uint scanned = simd_prefix_exclusive_sum(sum); + if (valid) { + group_sums[simd_lane_id] = scanned; + } + if (simd_lane_id == last_lane) { + group_sums[simd_groups] = scanned + sum; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < N; c++) { + threadgroup uint* group_sums = sums + c * scan_size; + values[c] = prev[c] + group_sums[simd_group_id]; + totals[c] = group_sums[simd_groups]; + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Radix select +/////////////////////////////////////////////////////////////////////////////// + +struct RadixBin { + uint digit; + uint below; + uint count; +}; + +// Prefix of the k-th key fixed so far with the counts around it, k counts +// from the bottom of the keys that still match the prefix +template +struct RadixSelectState { + KeyT prefix; + KeyT fixed; + uint k; + uint n_below; + uint n_equal; +}; + +template +METAL_FUNC RadixSelectState radix_select_start(uint k) { + RadixSelectState state; + state.prefix = 0; + state.fixed = 0; + state.k = k; + state.n_below = 0; + state.n_equal = 0; + return state; +} + +template +METAL_FUNC void radix_select_advance( + thread RadixSelectState& state, + RadixBin bin, + int shift) { + state.n_below += bin.below; + state.k -= bin.below; + state.prefix |= KeyT(bin.digit) << shift; + state.fixed |= KeyT((1 << DIGIT_BITS) - 1) << shift; + state.n_equal = bin.count; +} + +// Counts the key if its fixed digits match the prefix. Returns whether it did. +template +METAL_FUNC bool radix_count_key( + KeyT key, + const thread RadixSelectState& state, + int shift, + threadgroup atomic_uint* hist) { + if ((key & state.fixed) != state.prefix) { + return false; + } + uint digit = uint((key >> shift) & KeyT((1 << DIGIT_BITS) - 1)); + atomic_fetch_add_explicit(&hist[digit], 1u, memory_order_relaxed); + return true; +} + +METAL_FUNC uint radix_bin_count(threadgroup atomic_uint* hist, uint bin) { + return atomic_load_explicit(&hist[bin], memory_order_relaxed); +} + +template +METAL_FUNC uint radix_bins_sum(threadgroup atomic_uint* hist, uint first) { + uint sum = 0; + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < BINS_PER_THREAD; j++) { + sum += radix_bin_count(hist, first + j); + } + return sum; +} + +// Walks the bins of one thread with `below` keys before them. Returns whether +// the sum from the bottom reaches k inside them, then `bin` is that bin. +template +METAL_FUNC bool radix_bins_walk( + threadgroup atomic_uint* hist, + uint first, + uint below, + uint k, + thread RadixBin& bin) { + bool found = false; + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < BINS_PER_THREAD; j++) { + uint count = radix_bin_count(hist, first + j); + if (!found && below < k && below + count >= k) { + found = true; + bin.digit = first + j; + bin.below = below; + bin.count = count; + } + below += count; + } + return found; +} + +// Scans the histogram from the bottom with the whole threadgroup, +// `bins_per_thread` bins per thread. Returns the bin of the k-th smallest key. +template +METAL_FUNC RadixBin radix_find_bin( + threadgroup atomic_uint* hist, + uint k, + uint lid_x, + uint simd_lane_id, + uint simd_group_id, + threadgroup uint* bin_count_simdgroup_sums, + threadgroup RadixBin& kth_bin) { + constexpr int bins_per_thread = RadixLayout::bins_per_thread; + + uint first = lid_x * bins_per_thread; + uint below = radix_bins_sum(hist, first); + uint total; + threadgroup_exclusive_sum( + &below, &total, simd_lane_id, simd_group_id, bin_count_simdgroup_sums); + + // The sum from the bottom reaches k inside exactly one bin + RadixBin bin; + if (radix_bins_walk(hist, first, below, k, bin)) { + kth_bin = bin; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + return kth_bin; +} + +// Source of a select pass: the row, the row while it fills the candidates +// cache, or the cache only +enum class RadixPass { Row, FillCache, Cache }; + +template +METAL_FUNC void radix_load_keys( + const device T* in, + int i, + thread typename RadixKeyTraits::KeyT* keys) { + MLX_MTL_PRAGMA_UNROLL + for (int r = 0; r < radix_n_loads; r++) { + keys[r] = RadixKeyTraits::to_key(in[i + r * BLOCK_THREADS]); + } +} + +template +METAL_FUNC void radix_count_row_key( + KeyT key, + const thread RadixSelectState& state, + int shift, + RadixPass pass, + threadgroup atomic_uint* hist, + threadgroup KeyT* candidates_cache, + threadgroup atomic_uint& candidates_cache_size) { + if (radix_count_key(key, state, shift, hist) && + pass == RadixPass::FillCache) { + uint slot = atomic_fetch_add_explicit( + &candidates_cache_size, 1u, memory_order_relaxed); + candidates_cache[slot] = key; + } +} + +// Finds the k-th smallest element of the row, one histogram pass per digit +// from the top. Once the keys that still match the prefix fit in the +// candidates cache, the later passes read the cache instead of the row. +template +METAL_FUNC RadixSelectState::KeyT> radix_select_row( + const device T* in, + int n, + uint k, + uint lid_x, + uint simd_lane_id, + uint simd_group_id, + threadgroup atomic_uint* hist, + threadgroup typename RadixKeyTraits::KeyT* candidates_cache, + threadgroup atomic_uint& candidates_cache_size, + threadgroup uint* bin_count_simdgroup_sums, + threadgroup RadixBin& kth_bin) { + using KeyT = typename RadixKeyTraits::KeyT; + using Layout = RadixLayout; + constexpr int key_bits = 8 * sizeof(KeyT); + + auto state = radix_select_start(k); + uint n_candidates = 0; + RadixPass pass = RadixPass::Row; + + for (int p = 0; p < Layout::passes(key_bits); p++) { + int shift = Layout::shift(key_bits, p); + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < Layout::bins_per_thread; j++) { + atomic_store_explicit( + &hist[lid_x + j * BLOCK_THREADS], 0u, memory_order_relaxed); + } + if (lid_x == 0) { + atomic_store_explicit(&candidates_cache_size, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Count the keys that still match the fixed prefix + if (pass == RadixPass::Cache) { + for (uint i = lid_x; i < n_candidates; i += BLOCK_THREADS) { + radix_count_key( + candidates_cache[i], state, shift, hist); + } + } else { + int i = lid_x; + for (; i + (radix_n_loads - 1) * BLOCK_THREADS < n; + i += radix_n_loads * BLOCK_THREADS) { + KeyT keys[radix_n_loads]; + radix_load_keys(in, i, keys); + MLX_MTL_PRAGMA_UNROLL + for (int r = 0; r < radix_n_loads; r++) { + radix_count_row_key( + keys[r], + state, + shift, + pass, + hist, + candidates_cache, + candidates_cache_size); + } + } + for (; i < n; i += BLOCK_THREADS) { + radix_count_row_key( + RadixKeyTraits::to_key(in[i]), + state, + shift, + pass, + hist, + candidates_cache, + candidates_cache_size); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (pass == RadixPass::FillCache) { + n_candidates = + atomic_load_explicit(&candidates_cache_size, memory_order_relaxed); + pass = RadixPass::Cache; + } + + RadixBin bin = radix_find_bin( + hist, + state.k, + lid_x, + simd_lane_id, + simd_group_id, + bin_count_simdgroup_sums, + kth_bin); + radix_select_advance(state, bin, shift); + + // Once the candidates fit, the next pass copies them to the cache + if (pass == RadixPass::Row && bin.count <= Layout::candidates_cache_size) { + pass = RadixPass::FillCache; + } + } + return state; +} + +/////////////////////////////////////////////////////////////////////////////// +// Partition write +/////////////////////////////////////////////////////////////////////////////// + +// Writes one chunk of BLOCK_THREADS * radix_n_reads elements from `j`, one +// threadgroup scan gives each thread its output positions. A FULL chunk +// ends before `end`, so it needs no bounds checks. +template < + typename T, + typename U, + bool ARG_PARTITION, + int BLOCK_THREADS, + bool FULL> +METAL_FUNC void radix_partition_write_chunk( + const device T* in, + device U* out, + int j, + int end, + typename RadixKeyTraits::KeyT threshold, + thread uint* side_offset, + uint lid_x, + uint simd_lane_id, + uint simd_group_id, + threadgroup uint* threshold_side_count_simdgroup_sums) { + using KeyT = typename RadixKeyTraits::KeyT; + + T vals[radix_n_reads]; + int threshold_side[radix_n_reads]; + uint pos[radix_threshold_sides] = {0, 0, 0}; + + // Side of the threshold for radix_n_reads consecutive elements + MLX_MTL_PRAGMA_UNROLL + for (int r = 0; r < radix_n_reads; r++) { + int i = j + lid_x * radix_n_reads + r; + if (FULL || i < end) { + vals[r] = in[i]; + KeyT key = RadixKeyTraits::to_key(vals[r]); + threshold_side[r] = radix_threshold_side(key, threshold); + pos[threshold_side[r]]++; + } + } + + // Offsets of this thread inside the chunk and the chunk totals + uint totals[radix_threshold_sides]; + threadgroup_exclusive_sum( + pos, + totals, + simd_lane_id, + simd_group_id, + threshold_side_count_simdgroup_sums); + + // Add the offsets of each side, then write + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < radix_threshold_sides; c++) { + pos[c] += side_offset[c]; + side_offset[c] += totals[c]; + } + MLX_MTL_PRAGMA_UNROLL + for (int r = 0; r < radix_n_reads; r++) { + int i = j + lid_x * radix_n_reads + r; + if (FULL || i < end) { + if constexpr (ARG_PARTITION) { + out[pos[threshold_side[r]]] = U(i); + } else { + out[pos[threshold_side[r]]] = U(vals[r]); + } + pos[threshold_side[r]]++; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// Writes the elements in [start, end) of the row below, equal and above the +// threshold, keeping the input order inside each side. `side_offset` holds +// where each side continues in the output and moves along. +template +METAL_FUNC void radix_partition_write( + const device T* in, + device U* out, + int start, + int end, + typename RadixKeyTraits::KeyT threshold, + thread uint* side_offset, + uint lid_x, + uint simd_lane_id, + uint simd_group_id, + threadgroup uint* threshold_side_count_simdgroup_sums) { + constexpr int chunk_size = BLOCK_THREADS * radix_n_reads; + + for (int j = start; j < end; j += chunk_size) { + if (j + chunk_size <= end) { + radix_partition_write_chunk( + in, + out, + j, + end, + threshold, + side_offset, + lid_x, + simd_lane_id, + simd_group_id, + threshold_side_count_simdgroup_sums); + } else { + radix_partition_write_chunk( + in, + out, + j, + end, + threshold, + side_offset, + lid_x, + simd_lane_id, + simd_group_id, + threshold_side_count_simdgroup_sums); + } + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Partition kernel +/////////////////////////////////////////////////////////////////////////////// + +template +[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void +radix_partition( + const device T* in [[buffer(0)]], + device U* out [[buffer(1)]], + const constant int& axis_size [[buffer(2)]], + const constant int& kth [[buffer(3)]], + uint3 gid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + using KeyT = typename RadixKeyTraits::KeyT; + using Layout = RadixLayout; + + threadgroup atomic_uint hist[Layout::bins]; + threadgroup KeyT candidates_cache[Layout::candidates_cache_size]; + threadgroup atomic_uint candidates_cache_size; + threadgroup uint bin_count_simdgroup_sums[Layout::scan_size]; + threadgroup uint threshold_side_count_simdgroup_sums + [radix_threshold_sides * Layout::scan_size]; + threadgroup RadixBin kth_bin; + + int64_t offset = int64_t(gid.y) * axis_size; + in += offset; + out += offset; + + // kth is zero based, k is a count + uint k = kth + 1; + auto state = radix_select_row( + in, + axis_size, + k, + lid.x, + simd_lane_id, + simd_group_id, + hist, + candidates_cache, + candidates_cache_size, + bin_count_simdgroup_sums, + kth_bin); + uint side_offset[radix_threshold_sides] = { + 0, state.n_below, state.n_below + state.n_equal}; + radix_partition_write( + in, + out, + 0, + axis_size, + state.prefix, + side_offset, + lid.x, + simd_lane_id, + simd_group_id, + threshold_side_count_simdgroup_sums); +} diff --git a/mlx/backend/metal/kernels/partition.metal b/mlx/backend/metal/kernels/partition.metal new file mode 100644 index 0000000000..8e866e50a7 --- /dev/null +++ b/mlx/backend/metal/kernels/partition.metal @@ -0,0 +1,31 @@ +// Copyright © 2026 Apple Inc. + +#include + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/partition.h" + +#define instantiate_radix_partition_bn(tname, type, bn) \ + instantiate_kernel("radix_partition_" #tname "_bn" #bn, \ + radix_partition, type, type, false, bn) \ + instantiate_kernel("radix_argpartition_" #tname "_bn" #bn, \ + radix_partition, type, uint32_t, true, bn) + +#define instantiate_radix_partition(tname, type) \ + instantiate_radix_partition_bn(tname, type, 64) \ + instantiate_radix_partition_bn(tname, type, 128) \ + instantiate_radix_partition_bn(tname, type, 256) + +instantiate_radix_partition(uint8, uint8_t) +instantiate_radix_partition(uint16, uint16_t) +instantiate_radix_partition(uint32, uint32_t) +instantiate_radix_partition(uint64, uint64_t) +instantiate_radix_partition(int8, int8_t) +instantiate_radix_partition(int16, int16_t) +instantiate_radix_partition(int32, int32_t) +instantiate_radix_partition(int64, int64_t) +instantiate_radix_partition(float16, half) +instantiate_radix_partition(float32, float) +instantiate_radix_partition(bfloat16, bfloat16_t) +instantiate_radix_partition(complex64, complex64_t) // clang-format on diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 4f48a25b77..d3cbb05f4e 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -100,6 +100,16 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_partition_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + const array&, + bool, + int) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_searchsorted_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp index 15d6bdfb4e..d01e6f7231 100644 --- a/mlx/backend/metal/sort.cpp +++ b/mlx/backend/metal/sort.cpp @@ -317,6 +317,59 @@ void gpu_merge_sort( } } +constexpr int RADIX_PARTITION_MIN_AXIS_SIZE = 129; +constexpr size_t RADIX_PARTITION_MIN_SIZE = 256 * 1024; +constexpr int RADIX_PARTITION_LONG_AXIS_SIZE = 4096; + +bool use_radix_partition(const array& in, int axis) { + if (axis != in.ndim() - 1 || !in.flags().row_contiguous) { + return false; + } + int axis_size = in.shape(axis); + if (axis_size < RADIX_PARTITION_MIN_AXIS_SIZE) { + return false; + } + return in.size() >= RADIX_PARTITION_MIN_SIZE || + axis_size >= RADIX_PARTITION_LONG_AXIS_SIZE; +} + +void gpu_radix_partition( + const Stream& s, + metal::Device& d, + const array& in, + array& out, + int axis, + int kth, + bool arg_partition) { + int axis_size = in.shape(axis); + int n_rows = in.size() / axis_size; + + int bn = 64; + if (axis_size >= 1024) { + bn = 128; + } + if (axis_size >= 4096) { + bn = 256; + } + + std::string kernel_name = + arg_partition ? "radix_argpartition_" : "radix_partition_"; + concatenate(kernel_name, type_to_name(in), "_bn", bn); + auto kernel = + get_partition_kernel(d, kernel_name, in, out, arg_partition, bn); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(axis_size, 2); + compute_encoder.set_bytes(kth, 3); + + MTL::Size group_dims = MTL::Size(bn, 1, 1); + MTL::Size grid_dims = MTL::Size(1, n_rows, 1); + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + } // namespace void ArgSort::eval_gpu(const std::vector& inputs, array& out) { @@ -344,7 +397,6 @@ void Sort::eval_gpu(const std::vector& inputs, array& out) { } void ArgPartition::eval_gpu(const std::vector& inputs, array& out) { - // We direct arg partition to sort for now assert(inputs.size() == 1); out.set_data(allocator::malloc(out.nbytes())); @@ -353,11 +405,14 @@ void ArgPartition::eval_gpu(const std::vector& inputs, array& out) { auto& d = metal::device(s.device); auto& in = inputs[0]; - gpu_merge_sort(s, d, in, out, axis_, true); + if (use_radix_partition(in, axis_)) { + gpu_radix_partition(s, d, in, out, axis_, kth_, true); + } else { + gpu_merge_sort(s, d, in, out, axis_, true); + } } void Partition::eval_gpu(const std::vector& inputs, array& out) { - // We direct partition to sort for now assert(inputs.size() == 1); out.set_data(allocator::malloc(out.nbytes())); @@ -366,7 +421,11 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { auto& d = metal::device(s.device); auto& in = inputs[0]; - gpu_merge_sort(s, d, in, out, axis_, false); + if (use_radix_partition(in, axis_)) { + gpu_radix_partition(s, d, in, out, axis_, kth_, false); + } else { + gpu_merge_sort(s, d, in, out, axis_, false); + } } void SearchSorted::eval_gpu(const std::vector& inputs, array& out) { From a76f140a6444b17855b0ea14eff9413a8971b5dd Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 24 Sep 2026 17:57:25 +0200 Subject: [PATCH 2/4] Add tests for the radix partition kernel --- python/tests/test_ops.py | 125 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d5568610e8..7ac738f2f9 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3007,6 +3007,131 @@ def test_partition(self): M = top_k_mx.shape[axis or 0] self.assertEqual(M, (kth + N) % N) + def test_partition_long_rows(self): + # Rows long enough for the Metal radix select path + np.random.seed(0) + + def check(a_np, a_mx, kth, ref): + # a_np is the float32 reference of a_mx when a_mx is bfloat16, ref + # its sorted rows + rows, n = a_np.shape + b_mx = mx.partition(a_mx, kth, axis=-1) + self.assertEqual(b_mx.dtype, a_mx.dtype) + b_np = np.array( + b_mx.astype(mx.float32) if a_mx.dtype == mx.bfloat16 else b_mx + ) + i_mx = mx.argpartition(a_mx, kth, axis=-1) + self.assertEqual(i_mx.dtype, mx.uint32) + i_np = np.array(i_mx) + self.assertTrue( + np.array_equal(np.sort(i_np, axis=-1), np.tile(np.arange(n), (rows, 1))) + ) + for out in (b_np, np.take_along_axis(a_np, i_np, axis=-1)): + self.assertTrue( + np.array_equal(np.sort(out, axis=-1), ref, equal_nan=True) + ) + finite = np.where(np.isnan(out), np.inf, out) + pivot = finite[:, kth : kth + 1] + self.assertTrue(np.all(finite[:, :kth] <= pivot)) + self.assertTrue(np.all(finite[:, kth + 1 :] >= pivot)) + + def data(shape, dtype): + low = 0 if np.dtype(dtype).kind == "u" else -100 + return np.random.uniform(low, 100, size=shape).astype(dtype) + + # The widest threadgroup, kth at both ends and in the middle + shape = (16, 4096) + for dtype in ("float32", "float16", "int8", "int64"): + for kth in (0, 32, 2048, 4095): + a_np = data(shape, dtype) + ref = np.sort(a_np, axis=-1) + with self.subTest(dtype=dtype, kth=kth): + check(a_np, mx.array(a_np), kth, ref) + + # Row lengths that pick the smaller threadgroups, keys of every width + for short_shape, dtype in ( + ((512, 512), "uint8"), + ((512, 1000), "int16"), + ((2048, 129), "uint32"), + ((128, 2048), "uint64"), + ): + n = short_shape[1] + a_np = data(short_shape, dtype) + ref = np.sort(a_np, axis=-1) + for kth in (0, n // 2, n - 1): + with self.subTest(shape=short_shape, dtype=dtype, kth=kth): + check(a_np, mx.array(a_np), kth, ref) + + # bfloat16 goes through float32 for the reference + a_np = data(shape, np.float32) + a_np[:, ::97] = np.nan + a_mx = mx.array(a_np).astype(mx.bfloat16) + a_np = np.array(a_mx.astype(mx.float32)) + ref = np.sort(a_np, axis=-1) + for kth in (0, 32, 4095): + with self.subTest(dtype="bfloat16", kth=kth): + check(a_np, a_mx, kth, ref) + + # Ties and NaN + a_np = np.random.randint(0, 5, size=shape).astype(np.float32) + a_np[:, ::7] = np.nan + ref = np.sort(a_np, axis=-1) + for kth in (0, 2048, 4095): + with self.subTest(case="ties_nan", kth=kth): + check(a_np, mx.array(a_np), kth, ref) + + # Longer rows than the candidate cache, 16-bit and 32-bit keys + for dtype in ("float16", "float32"): + a_np = data((8, 40000), dtype) + with self.subTest(dtype=dtype, kth=40000 - 2048): + check(a_np, mx.array(a_np), 40000 - 2048, np.sort(a_np, axis=-1)) + + # complex64 against the sort of the same device, NaN cases on the GPU + a_np = ( + np.random.uniform(-100, 100, size=shape) + + 1j * np.random.randint(-3, 3, size=shape) + ).astype(np.complex64) + with_nan = a_np.copy() + with_nan.real[:, ::53] = np.nan + with_nan.imag[:, 5::71] = np.nan + cases = [(a_np, False)] + ( + [(with_nan, True)] if mx.default_device() == mx.gpu else [] + ) + for c_np, has_nan in cases: + a_mx = mx.array(c_np) + ref = np.array(mx.sort(a_mx, axis=-1)) + for kth in (0, 32, 2048, 4095): + with self.subTest(dtype="complex64", nan=has_nan, kth=kth): + i_mx = mx.argpartition(a_mx, kth, axis=-1) + self.assertTrue( + np.array_equal( + np.sort(np.array(i_mx), axis=-1), + np.tile(np.arange(shape[1]), (shape[0], 1)), + ) + ) + for out in ( + mx.partition(a_mx, kth, axis=-1), + mx.take_along_axis(a_mx, i_mx, axis=-1), + ): + for part, ref_part in ( + (out[:, : kth + 1], ref[:, : kth + 1]), + (out[:, kth:], ref[:, kth:]), + ): + self.assertTrue( + np.array_equal( + np.array(mx.sort(part, axis=-1)), + ref_part, + equal_nan=True, + ) + ) + + # A transposed view stays correct + a_np = data((4096, 64), np.float32) + b_t = np.array(mx.partition(mx.array(a_np).T, 3, axis=-1)) + self.assertTrue( + np.array_equal(b_t[:, 3], np.partition(a_np.T, 3, axis=-1)[:, 3]) + ) + def test_argpartition(self): x = mx.broadcast_to(mx.array([1, 2, 3]), (2, 3)) out = mx.argpartition(x, kth=1, axis=0) From a26929f85235df756b5a3af3f563c36e883f99ca Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 24 Sep 2026 18:24:29 +0200 Subject: [PATCH 3/4] Add partition benchmarks --- benchmarks/python/partition_bench.py | 36 ++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 benchmarks/python/partition_bench.py diff --git a/benchmarks/python/partition_bench.py b/benchmarks/python/partition_bench.py new file mode 100644 index 0000000000..f032a62f01 --- /dev/null +++ b/benchmarks/python/partition_bench.py @@ -0,0 +1,36 @@ +# Copyright © 2026 Apple Inc. + +import mlx.core as mx +from time_utils import time_fn + + +def time_partition(shape, k, dtype): + x = mx.random.normal(shape, dtype=dtype) + mx.eval(x) + name = f"{shape[0]}x{shape[1]} {dtype} k={k}" + time_fn(mx.partition, x, -k, axis=-1, msg=f"partition {name}") + time_fn(mx.argpartition, x, -k, axis=-1, msg=f"argpartition {name}") + # The merge sort is the path partition took before the radix kernel + time_fn(mx.sort, x, axis=-1, msg=f"sort {name}") + + +if __name__ == "__main__": + cases = ( + ((4096, 8192), 2048), + ((2048, 32768), 2048), + ((8, 131072), 2048), + ((64, 128000), 40), + ((2048, 8192), 32), + ((4096, 256), 8), + ((8192, 128), 8), + ((8192, 129), 8), + ((511, 512), 32), + ((512, 512), 32), + ((2048, 1023), 32), + ((2048, 1024), 32), + ((32, 4095), 32), + ((32, 4096), 32), + ) + for dtype in (mx.float32, mx.bfloat16): + for shape, k in cases: + time_partition(shape, k, dtype) From a7edc137b7227f51eefcbeb913510b17ad71f4ff Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 24 Sep 2026 22:43:41 +0200 Subject: [PATCH 4/4] Relax radix dispatch for contiguous inputs with stride-1 partition axes --- mlx/backend/metal/sort.cpp | 23 ++++-- python/tests/test_ops.py | 161 ++++++++++++++++++------------------- 2 files changed, 95 insertions(+), 89 deletions(-) diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp index d01e6f7231..a1ca73ac84 100644 --- a/mlx/backend/metal/sort.cpp +++ b/mlx/backend/metal/sort.cpp @@ -322,7 +322,7 @@ constexpr size_t RADIX_PARTITION_MIN_SIZE = 256 * 1024; constexpr int RADIX_PARTITION_LONG_AXIS_SIZE = 4096; bool use_radix_partition(const array& in, int axis) { - if (axis != in.ndim() - 1 || !in.flags().row_contiguous) { + if (!in.flags().contiguous || in.strides()[axis] != 1) { return false; } int axis_size = in.shape(axis); @@ -342,7 +342,8 @@ void gpu_radix_partition( int kth, bool arg_partition) { int axis_size = in.shape(axis); - int n_rows = in.size() / axis_size; + // Broadcast rows share their data. The output keeps the same strides + int n_rows = in.data_size() / axis_size; int bn = 64; if (axis_size >= 1024) { @@ -399,15 +400,20 @@ void Sort::eval_gpu(const std::vector& inputs, array& out) { void ArgPartition::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); - out.set_data(allocator::malloc(out.nbytes())); - auto& s = stream(); auto& d = metal::device(s.device); auto& in = inputs[0]; if (use_radix_partition(in, axis_)) { + // The kernel keeps the layout of the input + out.set_data( + allocator::malloc(in.data_size() * out.itemsize()), + in.data_size(), + in.strides(), + in.flags()); gpu_radix_partition(s, d, in, out, axis_, kth_, true); } else { + out.set_data(allocator::malloc(out.nbytes())); gpu_merge_sort(s, d, in, out, axis_, true); } } @@ -415,15 +421,20 @@ void ArgPartition::eval_gpu(const std::vector& inputs, array& out) { void Partition::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); - out.set_data(allocator::malloc(out.nbytes())); - auto& s = stream(); auto& d = metal::device(s.device); auto& in = inputs[0]; if (use_radix_partition(in, axis_)) { + // The kernel keeps the layout of the input + out.set_data( + allocator::malloc(in.data_size() * out.itemsize()), + in.data_size(), + in.strides(), + in.flags()); gpu_radix_partition(s, d, in, out, axis_, kth_, false); } else { + out.set_data(allocator::malloc(out.nbytes())); gpu_merge_sort(s, d, in, out, axis_, false); } } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 7ac738f2f9..23b33d031a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3011,29 +3011,40 @@ def test_partition_long_rows(self): # Rows long enough for the Metal radix select path np.random.seed(0) - def check(a_np, a_mx, kth, ref): - # a_np is the float32 reference of a_mx when a_mx is bfloat16, ref - # its sorted rows - rows, n = a_np.shape - b_mx = mx.partition(a_mx, kth, axis=-1) - self.assertEqual(b_mx.dtype, a_mx.dtype) - b_np = np.array( - b_mx.astype(mx.float32) if a_mx.dtype == mx.bfloat16 else b_mx - ) - i_mx = mx.argpartition(a_mx, kth, axis=-1) - self.assertEqual(i_mx.dtype, mx.uint32) - i_np = np.array(i_mx) - self.assertTrue( - np.array_equal(np.sort(i_np, axis=-1), np.tile(np.arange(n), (rows, 1))) - ) - for out in (b_np, np.take_along_axis(a_np, i_np, axis=-1)): - self.assertTrue( - np.array_equal(np.sort(out, axis=-1), ref, equal_nan=True) - ) - finite = np.where(np.isnan(out), np.inf, out) - pivot = finite[:, kth : kth + 1] - self.assertTrue(np.all(finite[:, :kth] <= pivot)) - self.assertTrue(np.all(finite[:, kth + 1 :] >= pivot)) + def check(a_np, a_mx, kths, axis=-1): + # a_np is the float32 reference of a_mx when a_mx is bfloat16. A + # NaN in any part of a value sorts last + n = a_np.shape[axis] + ref = np.sort(a_np, axis=axis) + ranks = np.indices(a_np.shape)[axis] + for kth in kths: + with self.subTest( + shape=a_np.shape, dtype=a_mx.dtype, axis=axis, kth=kth + ): + b_mx = mx.partition(a_mx, kth, axis=axis) + self.assertEqual(b_mx.dtype, a_mx.dtype) + b_np = np.array( + b_mx.astype(mx.float32) if a_mx.dtype == mx.bfloat16 else b_mx + ) + i_mx = mx.argpartition(a_mx, kth, axis=axis) + self.assertEqual(i_mx.dtype, mx.uint32) + i_np = np.array(i_mx) + self.assertTrue(np.array_equal(np.sort(i_np, axis=axis), ranks)) + for out in (b_np, np.take_along_axis(a_np, i_np, axis=axis)): + self.assertTrue( + np.array_equal(np.sort(out, axis=axis), ref, equal_nan=True) + ) + finite = np.where(np.isnan(out), np.inf, out) + pivot = np.take(finite, [kth], axis=axis) + self.assertTrue( + np.all(np.take(finite, np.arange(kth), axis=axis) <= pivot) + ) + self.assertTrue( + np.all( + np.take(finite, np.arange(kth + 1, n), axis=axis) + >= pivot + ) + ) def data(shape, dtype): low = 0 if np.dtype(dtype).kind == "u" else -100 @@ -3042,11 +3053,8 @@ def data(shape, dtype): # The widest threadgroup, kth at both ends and in the middle shape = (16, 4096) for dtype in ("float32", "float16", "int8", "int64"): - for kth in (0, 32, 2048, 4095): - a_np = data(shape, dtype) - ref = np.sort(a_np, axis=-1) - with self.subTest(dtype=dtype, kth=kth): - check(a_np, mx.array(a_np), kth, ref) + a_np = data(shape, dtype) + check(a_np, mx.array(a_np), (0, 32, 2048, 4095)) # Row lengths that pick the smaller threadgroups, keys of every width for short_shape, dtype in ( @@ -3057,80 +3065,67 @@ def data(shape, dtype): ): n = short_shape[1] a_np = data(short_shape, dtype) - ref = np.sort(a_np, axis=-1) - for kth in (0, n // 2, n - 1): - with self.subTest(shape=short_shape, dtype=dtype, kth=kth): - check(a_np, mx.array(a_np), kth, ref) + check(a_np, mx.array(a_np), (0, n // 2, n - 1)) # bfloat16 goes through float32 for the reference a_np = data(shape, np.float32) a_np[:, ::97] = np.nan a_mx = mx.array(a_np).astype(mx.bfloat16) - a_np = np.array(a_mx.astype(mx.float32)) - ref = np.sort(a_np, axis=-1) - for kth in (0, 32, 4095): - with self.subTest(dtype="bfloat16", kth=kth): - check(a_np, a_mx, kth, ref) + check(np.array(a_mx.astype(mx.float32)), a_mx, (0, 32, 4095)) # Ties and NaN a_np = np.random.randint(0, 5, size=shape).astype(np.float32) a_np[:, ::7] = np.nan - ref = np.sort(a_np, axis=-1) - for kth in (0, 2048, 4095): - with self.subTest(case="ties_nan", kth=kth): - check(a_np, mx.array(a_np), kth, ref) + check(a_np, mx.array(a_np), (0, 2048, 4095)) # Longer rows than the candidate cache, 16-bit and 32-bit keys for dtype in ("float16", "float32"): a_np = data((8, 40000), dtype) - with self.subTest(dtype=dtype, kth=40000 - 2048): - check(a_np, mx.array(a_np), 40000 - 2048, np.sort(a_np, axis=-1)) + check(a_np, mx.array(a_np), (40000 - 2048,)) - # complex64 against the sort of the same device, NaN cases on the GPU + # complex64, the NaN case only on the GPU a_np = ( np.random.uniform(-100, 100, size=shape) + 1j * np.random.randint(-3, 3, size=shape) ).astype(np.complex64) - with_nan = a_np.copy() - with_nan.real[:, ::53] = np.nan - with_nan.imag[:, 5::71] = np.nan - cases = [(a_np, False)] + ( - [(with_nan, True)] if mx.default_device() == mx.gpu else [] - ) - for c_np, has_nan in cases: - a_mx = mx.array(c_np) - ref = np.array(mx.sort(a_mx, axis=-1)) - for kth in (0, 32, 2048, 4095): - with self.subTest(dtype="complex64", nan=has_nan, kth=kth): - i_mx = mx.argpartition(a_mx, kth, axis=-1) - self.assertTrue( - np.array_equal( - np.sort(np.array(i_mx), axis=-1), - np.tile(np.arange(shape[1]), (shape[0], 1)), - ) - ) - for out in ( - mx.partition(a_mx, kth, axis=-1), - mx.take_along_axis(a_mx, i_mx, axis=-1), - ): - for part, ref_part in ( - (out[:, : kth + 1], ref[:, : kth + 1]), - (out[:, kth:], ref[:, kth:]), - ): - self.assertTrue( - np.array_equal( - np.array(mx.sort(part, axis=-1)), - ref_part, - equal_nan=True, - ) - ) + check(a_np, mx.array(a_np), (0, 32, 2048, 4095)) + if mx.default_device() == mx.gpu: + a_np.real[:, ::53] = np.nan + a_np.imag[:, 5::71] = np.nan + check(a_np, mx.array(a_np), (0, 32, 2048, 4095)) - # A transposed view stays correct - a_np = data((4096, 64), np.float32) - b_t = np.array(mx.partition(mx.array(a_np).T, 3, axis=-1)) - self.assertTrue( - np.array_equal(b_t[:, 3], np.partition(a_np.T, 3, axis=-1)[:, 3]) - ) + # Views with a partition axis of stride 1 take the radix path. The + # axis of 4096 selects it for any row count + a_np = data((8, 4096), np.float32) + a_mx = mx.array(a_np) + a3_np = a_np.reshape(2, 4, 4096) + a3_mx = a_mx.reshape(2, 4, 4096) + int_np = a_np.astype(np.int32) + for v_np, v_mx, axis in ( + (a_np.T, a_mx.T, 0), + (np.moveaxis(a3_np, -1, 0), mx.moveaxis(a3_mx, -1, 0), 0), + (np.moveaxis(a3_np, -1, 1), mx.moveaxis(a3_mx, -1, 1), 1), + (a_np[1:].T, a_mx[1:].T, 0), + (int_np.T, mx.array(int_np).T, 0), + (a_np[0], a_mx[0], 0), + ( + np.broadcast_to(a_np[:1], (64, 4096)), + mx.broadcast_to(a_mx[:1], (64, 4096)), + -1, + ), + ( + np.broadcast_to(a_np[:, None, :], (8, 3, 4096)), + mx.broadcast_to(a_mx[:, None, :], (8, 3, 4096)), + -1, + ), + ): + n = v_np.shape[axis] + check(v_np, v_mx, (0, n // 2, n - 1), axis) + + # The last axis of a transposed view has stride 4096 and stays on the + # sort path. The row count keeps the other dispatch checks true + a_np = data((160, 4096), np.float32) + check(a_np.T, mx.array(a_np).T, (0, 80, 159)) def test_argpartition(self): x = mx.broadcast_to(mx.array([1, 2, 3]), (2, 3))