[Metal][Performance] Add radix select kernel for partition and argpartition

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)

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

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();
 

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,

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,

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})