From b7c3dd6d27f45b5365b08a840310187dc503f1db Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 17 Jul 2026 18:53:03 +0200 Subject: [PATCH 001/222] [CUDA][Improvement] RMSNorm forward speed up (#3850) --- mlx/backend/cuda/rms_norm.cu | 153 ++++++++++++++++++++++++++--------- 1 file changed, 116 insertions(+), 37 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 365eb57b1c..97f3e94a5e 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -51,7 +51,18 @@ struct BlockBroadcastReduce { } }; -template +// xs and ws stay in registers +// Each thread does N_CHUNKS_THREAD vectorised interleaved loads of width +// N_READS x, w, out in majority of cases should be 16 bytes aligned -> using +// unsafe_load_vector if not aligned -> fall back to load_vector we do it like +// this because of the register presure: load_vector allocates registers for +// loop load fall back and occupancy drops to 30-40% -> x2 slow down +template < + typename T, + int BLOCK_SIZE, + int N_CHUNKS_THREAD, + bool ALIGNED, + int N_READS = 8> __global__ void rms_norm_small( const T* x, const T* w, @@ -63,40 +74,61 @@ __global__ void rms_norm_small( auto grid = cg::this_grid(); auto block = cg::this_thread_block(); - using BlockReduceT = BlockBroadcastReduce; + using BlockReduceT = BlockBroadcastReduce; __shared__ typename BlockReduceT::TempStorage temp; - auto row = - (grid.block_rank() * block.dim_threads().y) + block.thread_index().y; + auto row = grid.block_rank(); if (row >= n_rows) { return; } x += row * axis_size; out += row * axis_size; - // Normalizer. + AlignedVector xn[N_CHUNKS_THREAD]; + float normalizer = 0; auto index = block.thread_index().x; - auto xn = load_vector(x, index, axis_size, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float t = static_cast(xn[i]); - normalizer += t * t; + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + if constexpr (ALIGNED) { + xn[i] = unsafe_load_vector(x, offset); + } else { + xn[i] = load_vector(x, offset, axis_size, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; ++j) { + float t = static_cast(xn[i][j]); + normalizer += t * t; + } } - normalizer = BlockReduceT{block, temp}.Sum(normalizer); normalizer = rsqrt(normalizer / axis_size + eps); - // Outputs. - auto wn = load_vector(w, index, axis_size, w_stride, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float y = static_cast(xn[i]) * normalizer; - xn[i] = wn[i] * static_cast(y); + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + AlignedVector wn; + if constexpr (ALIGNED) { + wn = unsafe_load_vector(w, offset); + } else { + wn = load_vector(w, offset, axis_size, w_stride, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; j++) { + float y = static_cast(xn[i][j]) * normalizer; + xn[i][j] = wn[j] * static_cast(y); + } + if constexpr (ALIGNED) { + unsafe_store_vector(out, offset, xn[i]); + } else { + store_vector(out, offset, xn[i], axis_size); + } } - store_vector(out, index, xn, axis_size); } +// TODO: load x to the shared memory and reload from shared memory after the +// reduction template __global__ void rms_norm( const T* x, @@ -318,6 +350,42 @@ void dispatch_group_dim(int axis_size, F&& f) { } } +template +void dispatch_chunks(int n_chunks, F&& f) { + auto block_size = std::integral_constant{}; + if (n_chunks <= 1) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 2) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 3) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 4) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 5) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 6) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 7) { + f(block_size, std::integral_constant{}); + } else { + f(block_size, std::integral_constant{}); + } +} + +template +void dispatch_num_chunks(int axis_size, F&& f) { + int nvec = (axis_size + N_READS - 1) / N_READS; + if (axis_size <= N_READS * 64) { + f(std::integral_constant{}, std::integral_constant{}); + } else if (nvec % 128 == 0 && nvec / 128 <= 8) { + dispatch_chunks<128>(nvec / 128, f); + } else if (nvec % 64 == 0 && nvec / 64 <= 8) { + dispatch_chunks<64>(nvec / 64, f); + } else { + dispatch_chunks<128>((nvec + 127) / 128, f); + } +} + // TODO: There are duplicate code with backend/metal/normalization.cpp void RMSNorm::eval_gpu( const std::vector& inputs, @@ -365,27 +433,38 @@ void RMSNorm::eval_gpu( dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) { using DataType = cuda_type_t; constexpr int N_READS = 16 / sizeof(DataType); - if (axis_size <= N_READS * 1024) { - dispatch_group_dim( - axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { - constexpr int block_dim = n_groups() * group_dim(); - static_assert(block_dim <= 32 || groups_per_block() == 1); - auto kernel = - cu::rms_norm_small; - auto n_blocks = - (n_rows + groups_per_block() - 1) / groups_per_block(); - encoder.add_kernel_node( - kernel, - n_blocks, - {block_dim, groups_per_block()}, - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(out), - eps_, - axis_size, - n_rows, - w_stride); - }); + if (axis_size <= N_READS * 128 * 8) { + dispatch_num_chunks< + N_READS>(axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + constexpr int N_CHUNKS = n_chunks(); + bool aligned = (axis_size == N_READS * BLOCK_SIZE * N_CHUNKS) && + (w_stride == 1) && + (reinterpret_cast(gpu_ptr(x)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(w)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(out)) % 16 == 0); + // MSVC can't find N_READS two lambda levels down, dispatch_bool would + // add that second level + auto kernel = aligned + ? cu::rms_norm_small + : cu::rms_norm_small< + DataType, + BLOCK_SIZE, + N_CHUNKS, + false, + N_READS>; + encoder.add_kernel_node( + kernel, + n_rows, + BLOCK_SIZE, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(out), + eps_, + axis_size, + n_rows, + w_stride); + }); } else { auto kernel = cu::rms_norm; encoder.add_kernel_node( From ce3073389e92d8f7b7ecd3d7697523795786e046 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Mon, 20 Jul 2026 13:09:32 -0700 Subject: [PATCH 002/222] Fix captured random state in compile (#3828) --- python/src/random.cpp | 42 +++++++++++++++++++++- python/src/random.h | 17 +++++++++ python/src/trees.cpp | 27 +++++++++----- python/tests/test_compile.py | 68 ++++++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+), 10 deletions(-) create mode 100644 python/src/random.h diff --git a/python/src/random.cpp b/python/src/random.cpp index ddddb56eed..ceebb52e7c 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -9,6 +9,7 @@ #include "mlx/ops.h" #include "mlx/random.h" +#include "python/src/random.h" #include "python/src/small_vector.h" #include "python/src/utils.h" @@ -63,15 +64,54 @@ PyKeySequence& default_key() { return ks; } +// A process-global sentinel for `mx.random.state`. Since it is the same object +// on every thread, capturing it (e.g. with `mx.compile`) is thread-independent; +// the pytree traversal in trees.cpp resolves it to the calling thread's key. +class RandomState {}; + +nb::object random_state_sentinel() { + static nb::object sentinel = []() { + auto sentinel = nb::cast(RandomState{}); + sentinel.inc_ref(); + return sentinel; + }(); + + return sentinel; +} + +mx::array random_state_key() { + return nb::cast(default_key().state()[0]); +} + +void set_random_state_key(const mx::array& key) { + default_key().state()[0] = nb::cast(key); +} + void init_random(nb::module_& parent_module) { auto m = parent_module.def_submodule( "random", "mlx.core.random: functionality related to random number generation"); + nb::class_(m, "_RandomState") + .def("__len__", [](const RandomState&) { return 1; }) + .def( + "__getitem__", + [](const RandomState&, int i) -> nb::object { + if (i != 0 && i != -1) { + throw nb::index_error("random state index out of range"); + } + return default_key().state()[0]; + }, + "index"_a) + .def("__iter__", [](const RandomState&) { + return nb::iter(default_key().state()); + }); + m.def("__getattr__", [&](nb::handle key) -> nb::object { // Create random.state lazily to avoid initializing device during import. if (nb::isinstance(key) && nb::cast(key) == "state") { - return default_key().state(); + default_key().state(); + return random_state_sentinel(); } return nb::steal(PyErr_Format( PyExc_AttributeError, diff --git a/python/src/random.h b/python/src/random.h new file mode 100644 index 0000000000..2baf9d92f1 --- /dev/null +++ b/python/src/random.h @@ -0,0 +1,17 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "mlx/array.h" + +namespace mx = mlx::core; +namespace nb = nanobind; + +// The process-global `mx.random.state` sentinel. +nb::object random_state_sentinel(); + +// Read/write the calling thread's current PRNG key. +mx::array random_state_key(); +void set_random_state_key(const mx::array& key); diff --git a/python/src/trees.cpp b/python/src/trees.cpp index 4b9ca9e123..cf05d4e25a 100644 --- a/python/src/trees.cpp +++ b/python/src/trees.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include "python/src/trees.h" +#include "python/src/random.h" template void validate_subtrees(const std::vector& subtrees) { @@ -152,8 +153,12 @@ void tree_visit( void tree_visit(nb::handle tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (nb::isinstance(subtree) || + if (subtree.is(random_state)) { + visitor(nb::cast(random_state_key())); + } else if ( + nb::isinstance(subtree) || nb::isinstance(subtree)) { for (auto item : subtree) { recurse(item); @@ -174,8 +179,14 @@ void tree_visit_update( nb::object tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (nb::isinstance(subtree)) { + if (subtree.is(random_state)) { + // Read/write the calling thread's key; keep the sentinel in the tree. + set_random_state_key( + nb::cast(visitor(nb::cast(random_state_key())))); + return nb::cast(subtree); + } else if (nb::isinstance(subtree)) { auto l = nb::cast(subtree); for (int i = 0; i < l.size(); ++i) { l[i] = recurse(l[i]); @@ -262,14 +273,12 @@ nb::object tree_unflatten( } nb::object structure_sentinel() { - static nb::object sentinel; - - if (sentinel.ptr() == nullptr) { - sentinel = nb::capsule(&sentinel); - // probably not needed but this should make certain that we won't ever - // delete the sentinel + static nb::object sentinel = []() { + PyObject* raw_obj = PyObject_New(PyObject, &PyBaseObject_Type); + nb::object sentinel = nb::steal(raw_obj); sentinel.inc_ref(); - } + return sentinel; + }(); return sentinel; } diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 35e51600e1..7a2c6b9d0d 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -439,6 +439,74 @@ def fun(): self.assertFalse(mx.allclose(fun(), fun(), 1e-2, 1e-2)) + def test_compile_rng_across_threads(self): + # A function compiled with inputs/outputs=mx.random.state on one thread + # must still use (and advance/seed) the calling thread's RNG state when + # invoked from another thread, whether captured directly or nested. + + # The state sentinel is a single global object shared across threads. + state_from_thread = {} + + def grab(): + state_from_thread["s"] = mx.random.state + + t = threading.Thread(target=grab) + t.start() + t.join() + self.assertIs(mx.random.state, state_from_thread["s"]) + + direct = partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)( + lambda: mx.random.uniform(shape=(10, 10)) + ) + + nested_state = [{"unused": mx.array(0.0)}, mx.random.state] + nested = partial(mx.compile, inputs=nested_state, outputs=nested_state)( + lambda: mx.random.uniform(shape=(10, 10)) + ) + + for fun in (direct, nested): + results = {} + + def worker(): + with mx.stream(mx.cpu): + a = fun() + b = fun() + results["advances"] = not bool(mx.allclose(a, b, 1e-2, 1e-2).item()) + mx.random.seed(42) + c = fun() + mx.random.seed(42) + d = fun() + results["seed_reproducible"] = bool(mx.allclose(c, d).item()) + mx.random.seed(1234) + e = fun() + results["seed_changes"] = not bool( + mx.allclose(c, e, 1e-2, 1e-2).item() + ) + + t = threading.Thread(target=worker) + t.start() + t.join() + + self.assertTrue(results["advances"]) + self.assertTrue(results["seed_reproducible"]) + self.assertTrue(results["seed_changes"]) + + def test_compile_state_capture_with_rng_updates_in_place(self): + # Capturing mx.random.state alongside other state via outputs= must not + # break in-place updates of the other captured containers. + counter = {"v": mx.array(0.0)} + state = [counter, mx.random.state] + + @partial(mx.compile, inputs=state, outputs=state) + def step(): + counter["v"] = counter["v"] + 1.0 + return mx.random.uniform(shape=(2,)) + + for _ in range(3): + step() + mx.eval(counter["v"]) + self.assertEqual(counter["v"].item(), 3.0) + def test_compile_kwargs(self): @mx.compile def fun(x, y, z): From 353440c378928ccc1c92c504b434c6054156c14d Mon Sep 17 00:00:00 2001 From: Kolja Wawrowsky <3075215+apocryphx@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:34:55 -0700 Subject: [PATCH 003/222] Fix JIT preamble header filter matching project paths containing "Xcode" (#3873) --- mlx/backend/metal/make_compiled_preamble.sh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh index 07c8863bc6..1b1d9cf45a 100644 --- a/mlx/backend/metal/make_compiled_preamble.sh +++ b/mlx/backend/metal/make_compiled_preamble.sh @@ -50,8 +50,10 @@ if [ -n "$HDRS" ]; then fi fi -# Remove any included system frameworks (for MetalPerformancePrimitive headers) -HDRS=$(echo "$HDRS" | grep -v "Xcode") +# Remove any included system frameworks (for MetalPerformancePrimitive headers). +# Match "Xcode.app" rather than bare "Xcode" so project checkouts that live under +# a directory named Xcode/ are not filtered out along with the SDK headers. +HDRS=$(echo "$HDRS" | grep -v "Xcode.app") # Use the header depth to sort the files in order of inclusion declare -a HDRS_LIST=() From 30a19f7234abfb4d136af58623eaac289e63f5d2 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 20 Jul 2026 23:36:11 -0700 Subject: [PATCH 004/222] Document default value of p in dropout layer docstrings (#3870) --- python/mlx/nn/layers/dropout.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/dropout.py b/python/mlx/nn/layers/dropout.py index 78f8c7dbf5..ef74ba9b60 100644 --- a/python/mlx/nn/layers/dropout.py +++ b/python/mlx/nn/layers/dropout.py @@ -12,7 +12,7 @@ class Dropout(Module): expected value of a given element will remain the same. Args: - p (float): The probability to zero an element + p (float): The probability to zero an element. Default: ``0.5``. """ def __init__(self, p: float = 0.5): @@ -56,6 +56,7 @@ class Dropout2d(Module): Args: p (float): Probability of zeroing a channel during training. + Default: ``0.5``. """ def __init__(self, p: float = 0.5): @@ -105,6 +106,7 @@ class Dropout3d(Module): Args: p (float): Probability of zeroing a channel during training. + Default: ``0.5``. """ def __init__(self, p: float = 0.5): From 8f64abc9b9d66053ed591fd4f5fca81a0dfc28c5 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Tue, 21 Jul 2026 23:10:05 +0200 Subject: [PATCH 005/222] [WIP] [CUDA] fsdp (#3768) --- docs/src/python/nn.rst | 1 - docs/src/python/nn/distributed.rst | 2 + python/mlx/nn/__init__.py | 1 - python/mlx/nn/layers/__init__.py | 2 + python/mlx/nn/layers/distributed.py | 148 +++++++++++++- python/mlx/nn/utils.py | 165 +++------------ python/tests/mlx_distributed_tests.py | 40 +++- python/tests/nccl_test_distributed.py | 279 ++++++-------------------- 8 files changed, 273 insertions(+), 365 deletions(-) diff --git a/docs/src/python/nn.rst b/docs/src/python/nn.rst index 0ea3a56fbb..00f7f95456 100644 --- a/docs/src/python/nn.rst +++ b/docs/src/python/nn.rst @@ -175,7 +175,6 @@ In detail: value_and_grad quantize average_gradients - fsdp_apply_gradients .. toctree:: diff --git a/docs/src/python/nn/distributed.rst b/docs/src/python/nn/distributed.rst index 07dd8e2308..23998b71a3 100644 --- a/docs/src/python/nn/distributed.rst +++ b/docs/src/python/nn/distributed.rst @@ -15,6 +15,7 @@ create sharded layers from existing :class:`Modules `. shard_linear shard_inplace + fully_shard Layers ^^^^^^ @@ -28,3 +29,4 @@ Layers ShardedToAllLinear QuantizedAllToShardedLinear QuantizedShardedToAllLinear + FullyShardedModule diff --git a/python/mlx/nn/__init__.py b/python/mlx/nn/__init__.py index 6c3bd0aa24..5b51f8a01d 100644 --- a/python/mlx/nn/__init__.py +++ b/python/mlx/nn/__init__.py @@ -4,6 +4,5 @@ from mlx.nn.layers import * from mlx.nn.utils import ( average_gradients, - fsdp_apply_gradients, value_and_grad, ) diff --git a/python/mlx/nn/layers/__init__.py b/python/mlx/nn/layers/__init__.py index c2fba58347..25b96c0a7f 100644 --- a/python/mlx/nn/layers/__init__.py +++ b/python/mlx/nn/layers/__init__.py @@ -64,9 +64,11 @@ ) from mlx.nn.layers.distributed import ( AllToShardedLinear, + FullyShardedModule, QuantizedAllToShardedLinear, QuantizedShardedToAllLinear, ShardedToAllLinear, + fully_shard, ) from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d from mlx.nn.layers.embedding import Embedding diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index c4799bd073..69804756c9 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -1,14 +1,14 @@ # Copyright © 2024 Apple Inc. import math -from functools import lru_cache +from functools import lru_cache, reduce from typing import Callable, Optional, Union import mlx.core as mx from mlx.nn.layers.base import Module from mlx.nn.layers.linear import Linear from mlx.nn.layers.quantized import QuantizedLinear -from mlx.utils import tree_map_with_path +from mlx.utils import tree_flatten, tree_map_with_path, tree_unflatten @lru_cache @@ -617,3 +617,147 @@ def from_quantized_linear( ) return sl + + +def _make_gather_fn(group, full_shapes, shard_sizes, compute_dtype): + N = group.size() + indices = reduce(lambda acc, w: acc + [acc[-1] + w], shard_sizes, [0]) + split_indices = indices[1:-1] + shard_shapes = [(shape[0] // N,) + tuple(shape[1:]) for shape in full_shapes] + + def _maybe_cast(x, dtype): + if dtype is None or x.dtype == dtype: + return x + return x.astype(dtype) + + @mx.custom_function + def gather(shards): + shard = mx.concatenate( + [_maybe_cast(s.reshape(1, -1), compute_dtype) for s in shards], axis=1 + ) + full = mx.distributed.all_gather(shard, group=group) + parts = mx.split(full, split_indices, axis=1) + return [p.reshape(shape) for p, shape in zip(parts, full_shapes)] + + @gather.vjp + def gather_vjp(shards, cotangents, _): + local_full = mx.concatenate([c.reshape(N, -1) for c in cotangents], axis=1) + local_shard = mx.distributed.sum_scatter(local_full, group=group) / N + parts = mx.split(local_shard, split_indices, axis=1) + return [ + _maybe_cast(p.reshape(shape), s.dtype) + for p, shape, s in zip(parts, shard_shapes, shards) + ] + + return gather + + +def _maybe_shard(m, k, v): + if isinstance(v, FullyShardedModule): + return False + return Module.valid_parameter_filter(m, k, v) + + +class FullyShardedModule(Module): + """Wrap a module so each member of the group holds only a shard of its + parameters. + + The full parameters are gathered for the forward pass and the gradients + are reduce-scattered in the backward pass, so during training + each member of the group stores and updates only its own shard. + + Every parameter is sharded along axis 0, so each parameter's size along + that axis must be divisible by the size of ``group``. + + Use :func:`fully_shard` to wrap a module. + + Args: + module (mlx.nn.Module): The module whose parameters will be sharded. + group (mlx.core.distributed.Group, optional): The group to shard + across. If not set, the global group is used. Default: ``None``. + compute_dtype (mlx.core.Dtype, optional): If set, the gathered + parameters are cast to this dtype for the forward pass. + Default: ``None``. + """ + + def __init__( + self, + module: Module, + group: Optional[mx.distributed.Group] = None, + compute_dtype: Optional[mx.Dtype] = None, + ): + super().__init__() + group = group or mx.distributed.init() + N = group.size() + + shard_params = module.filter_and_map(_maybe_shard) + flat = tree_flatten(shard_params) + for path, a in flat: + if a.ndim == 0: + raise ValueError( + f"Cannot shard parameter '{path}' because it is a scalar." + ) + if a.shape[0] % N != 0: + raise ValueError( + f"Cannot shard parameter '{path}' with shape {a.shape} " + f"across {N} devices: axis 0 must be divisible by {N}." + ) + + super(Module, self).__setattr__("_paths", [k for k, _ in flat]) + full_shapes = [a.shape for _, a in flat] + shard_sizes = [a.size // N for _, a in flat] + + module.update(_shard(shard_params, lambda p, w: 0, group)) + + self.module = module + self._gather_fn = _make_gather_fn( + group, full_shapes, shard_sizes, compute_dtype + ) + + def _extra_repr(self) -> str: + return f"num_sharded_params={len(self._paths)}" + + def _gathered_call(self, fn, *args, **kwargs): + shard_tree = self.module.filter_and_map(_maybe_shard) + shards = [a for _, a in tree_flatten(shard_tree)] + fulls = self._gather_fn(shards) + self.module.update(tree_unflatten(list(zip(self._paths, fulls)))) + try: + return fn(*args, **kwargs) + finally: + self.module.update(shard_tree) + + def __call__(self, *args, **kwargs): + return self._gathered_call(self.module, *args, **kwargs) + + def as_linear(self, *args, **kwargs): + return self._gathered_call(self.module.as_linear, *args, **kwargs) + + +def fully_shard( + module: Module, + *, + group: Optional[mx.distributed.Group] = None, + compute_dtype: Optional[mx.Dtype] = None, +) -> Module: + """Wrap ``module`` in a :class:`FullyShardedModule`. + + Args: + module (mlx.nn.Module): The module to wrap. + group (mlx.core.distributed.Group, optional): The group to shard + across. If not set, the global group is used. Default: ``None``. + compute_dtype (mlx.core.Dtype, optional): If set, the gathered + parameters are cast to this dtype for the forward pass. + Default: ``None``. + + Returns: + The wrapped :class:`FullyShardedModule`, or ``module`` unchanged. + """ + group = group or mx.distributed.init() + if group.size() == 1: + return module + if isinstance(module, FullyShardedModule): + return module + + wrapped = FullyShardedModule(module, group=group, compute_dtype=compute_dtype) + return wrapped if wrapped._paths else module diff --git a/python/mlx/nn/utils.py b/python/mlx/nn/utils.py index b53e9efe21..a60468db64 100644 --- a/python/mlx/nn/utils.py +++ b/python/mlx/nn/utils.py @@ -173,149 +173,36 @@ def average_gradients( return tree_unflatten(new_flat_grads) -def _clip_grads_fsdp(grads_slice, max_norm, group=None): - local_norm_sq = tree_reduce(lambda acc, g: acc + g.square().sum(), grads_slice, 0.0) - global_norm_sq = mx.distributed.all_sum(local_norm_sq, group=group) - grad_norm = mx.sqrt(global_norm_sq) - normalizer = mx.minimum(max_norm / (grad_norm + 1e-6), 1.0) - grads_slice = tree_map(lambda g: g * normalizer, grads_slice) - - return grads_slice, grad_norm - - -def fsdp_apply_gradients( - gradients, - parameters, - optimizer, - fsdp_group=None, - dp_group=None, - communication_size=32 * 1024**2, - communication_stream=None, - max_norm=None, +def clip_grad_norm_sharded( + gradients: Any, + max_norm: float, + group: Optional[mx.distributed.Group] = None, ): - """Perform a distributed optimizer step by sharding gradients and optimizer states across ranks. - - This helper function performs the following steps: - 1. Reduce-scatter the gradients across ranks so each rank gets a shard of the averaged gradients. - 2. Optionally clip the sharded gradients by global norm. - 3. Apply the optimizer update on the local parameter slice using the sharded gradients. - 4. All-gather the updated parameter slices from all ranks to reconstruct the full parameters tree. + """Clip the global norm of gradients that are sharded across a group. - This is similar to PyTorch's FSDP with `reshard_after_forward=False`. + This is the sharded equivalent of + :func:`mlx.optimizers.clip_grad_norm`. Each member of the group holds only + a shard of the gradients, so the global norm is computed by summing the + local squared norms across the group before rescaling. It is useful for + clipping the gradients of a module wrapped with :func:`mlx.nn.fully_shard`. Args: - gradients (Any): The Python tree containing the full gradients (it should - have the same structure as ``parameters``). Each gradient's first - dimension must be divisible by ``fsdp_group.size()``. - parameters (Any): The Python tree containing the full parameters (it should - have the same structure across processes). Each parameter's first - dimension must be divisible by ``fsdp_group.size()``. - optimizer: Optimizer with an ``apply_gradients`` method. - fsdp_group (Optional[mlx.core.distributed.Group]): The group of processes - for FSDP sharding. If ``None``, the global group is used. - dp_group (Optional[mlx.core.distributed.Group]): The group of processes - for data-parallel gradient averaging. Required when ``fsdp_group`` is - smaller than the world (e.g. FSDP intra-node, DDP inter-node). - Default: ``None``. - communication_size (int): Group arrays until their size in bytes exceeds - this number. Perform one communication step per group of arrays. If - less or equal to 0 array grouping is disabled. Default: ``32MiB``. - communication_stream (Optional[mlx.core.Stream]): The stream to use - for the communication. If unspecified the default communication - stream is used which can vary by back-end. Default: ``None``. - max_norm (Optional[float]): If provided, clip gradients to this - maximum global norm before applying the optimizer update. - Default: ``None``. + gradients (Any): A Python tree containing the local shard of the + gradient arrays. + max_norm (float): The maximum allowed global norm of the gradients. + group (Optional[mlx.core.distributed.Group]): The group across which + the gradients are sharded. If set to ``None`` the global group is + used. Default: ``None``. Returns: - If ``max_norm`` is ``None``, returns the updated full-parameter tree. - Otherwise returns ``(parameters, grad_norm)``, where ``grad_norm`` is - the global gradient norm before clipping. - - Example: - - >>> optimizer = optim.SGD(learning_rate=0.01) - >>> # Without gradient clipping - >>> updated_params = fsdp_apply_gradients(grads, params, optimizer) - >>> model.update(updated_params) - >>> - >>> # With gradient clipping - >>> updated_params, grad_norm = fsdp_apply_gradients( - ... grads, params, optimizer, max_norm=1.0 - ... ) - >>> model.update(updated_params) + (Any, mlx.core.array): The possibly rescaled local shard of the + gradients and the global gradient norm. """ - fsdp_group = fsdp_group or mx.distributed.init() - N = fsdp_group.size() * (dp_group.size() if dp_group is not None else 1) - - if N == 1: - if max_norm is not None: - gradients, grad_norm = _clip_grads_fsdp(gradients, max_norm) - return optimizer.apply_gradients(gradients, parameters), grad_norm - return optimizer.apply_gradients(gradients, parameters) - - flat_grads = tree_flatten(gradients) - flat_params = tree_flatten(parameters) - - keys, shapes, sizes, dtypes = _extract_info(flat_grads) - itemsize = dtypes[0].size - - groups = _group_by_size(keys, sizes, itemsize, communication_size) - - S = fsdp_group.size() - fsdp_rank = fsdp_group.rank() - # reduce-scatter gradients, shard parameters - grad_slices = {} - param_slices = {} - for group_idx, arr_group in enumerate(groups): - big_grad = mx.concatenate( - [flat_grads[i][1].reshape(S, -1) for i in arr_group], axis=1 - ) - grad_slices[group_idx] = ( - mx.distributed.sum_scatter( - big_grad, group=fsdp_group, stream=communication_stream - ) - / N - ) - if dp_group is not None: - grad_slices[group_idx] = mx.distributed.all_sum( - grad_slices[group_idx], group=dp_group, stream=communication_stream - ) - big_param = mx.concatenate( - [flat_params[i][1].reshape(S, -1) for i in arr_group], axis=1 - ) - param_slices[group_idx] = big_param[fsdp_rank] - - # clip gradients if needed - grad_norm = None - if max_norm is not None: - grad_slices, grad_norm = _clip_grads_fsdp( - grad_slices, max_norm, group=fsdp_group - ) - - # optimizer step - updated_param_slices = optimizer.apply_gradients(grad_slices, param_slices) - - # all-gather and reconstruct - new_flat = [] - for group_idx, arr_group in enumerate(groups): - big_gathered = mx.distributed.all_gather( - updated_param_slices[group_idx], - group=fsdp_group, - stream=communication_stream, - ) - split_sizes = [sizes[i] // S for i in arr_group] - split_indices = [] - acc = 0 - for s in split_sizes: - acc += s - split_indices.append(acc) - - parts = mx.split(big_gathered, split_indices[:-1], axis=1) - for idx_in_group, i in enumerate(arr_group): - new_flat.append((keys[i], parts[idx_in_group].reshape(shapes[i]))) - - result = tree_unflatten(new_flat) - if max_norm is not None: - return result, grad_norm - return result + local_norm_squared = tree_reduce( + lambda acc, g: acc + g.square().sum(), gradients, 0.0 + ) + global_norm_squared = mx.distributed.all_sum(local_norm_squared, group=group) + grad_norm = mx.sqrt(global_norm_squared) + normalizer = mx.minimum(max_norm / (grad_norm + 1e-6), 1.0) + clipped_gradients = tree_map(lambda g: g * normalizer, gradients) + return clipped_gradients, grad_norm diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 00d3182bc0..77bdea6a1f 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -1,10 +1,12 @@ # Copyright © 2025 Apple Inc. +import math + import mlx.core as mx import mlx.nn as nn import mlx_tests from mlx.nn.layers.distributed import shard_inplace, shard_linear -from mlx.nn.utils import average_gradients +from mlx.nn.utils import average_gradients, clip_grad_norm_sharded class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase): @@ -322,3 +324,39 @@ def test_all_gather(self): y = mx.distributed.all_gather(x) self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + + def test_clip_grad_norm_sharded(self): + world = mx.distributed.init() + N = world.size() + + value = 3.0 + grads_slice = {"a": mx.ones((4, 3)) * value, "b": mx.ones((5,)) * value} + local_numel = 4 * 3 + 5 + expected_norm = math.sqrt(N * local_numel) * value + + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=1e9, group=world + ) + mx.eval(clipped, grad_norm) + self.assertTrue( + mx.allclose( + grad_norm, mx.array(expected_norm), atol=self.atol, rtol=self.rtol + ) + ) + for k in grads_slice: + self.assertTrue( + mx.allclose(clipped[k], grads_slice[k], atol=self.atol, rtol=self.rtol) + ) + + max_norm = 1.0 + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=max_norm, group=world + ) + mx.eval(clipped, grad_norm) + scale = max_norm / (expected_norm + 1e-6) + for k in grads_slice: + self.assertTrue( + mx.allclose( + clipped[k], grads_slice[k] * scale, atol=self.atol, rtol=self.rtol + ) + ) diff --git a/python/tests/nccl_test_distributed.py b/python/tests/nccl_test_distributed.py index 7a6bcce04b..db9bbb2e62 100644 --- a/python/tests/nccl_test_distributed.py +++ b/python/tests/nccl_test_distributed.py @@ -1,10 +1,11 @@ # Copyright © 2024 Apple Inc. import mlx.core as mx -import mlx.optimizers as optim +import mlx.nn as nn import mlx_distributed_tests import mlx_tests -from mlx.nn.utils import average_gradients, fsdp_apply_gradients +from mlx.nn.utils import average_gradients +from mlx.utils import tree_flatten, tree_map class TestNCCLDistributed(mlx_distributed_tests.MLXDistributedCommonTestCase): @@ -116,230 +117,66 @@ def test_all_gather_split(self): self.assertEqual(y.shape, (sub.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) - def test_fsdp_apply_gradients(self): - world = mx.distributed.init() - N = world.size() - - params = { - "w1": mx.ones((N * 10, 8)), - "w2": mx.ones((N * 20,)), - } - grads = { - "w1": mx.ones((N * 10, 8)) * 0.1, - "w2": mx.ones((N * 20,)) * 0.1, - } - - optimizer = optim.SGD(learning_rate=0.1) - updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer) - mx.eval(updated_params_fsdp) - - self.assertEqual(updated_params_fsdp["w1"].shape, (N * 10, 8)) - self.assertEqual(updated_params_fsdp["w2"].shape, (N * 20,)) - - self.assertTrue( - mx.allclose( - updated_params_fsdp["w1"], mx.ones((N * 10, 8)) * 0.99, atol=1e-6 - ) - ) - self.assertTrue( - mx.allclose(updated_params_fsdp["w2"], mx.ones((N * 20,)) * 0.99, atol=1e-6) - ) - - grads = { - "w1": mx.ones((N * 10, 8)) * 10.0, - "w2": mx.ones((N * 20,)) * 10.0, - } - - new_params_clipped, grad_norm = fsdp_apply_gradients( - grads, params, optimizer, max_norm=1.0 - ) - mx.eval(new_params_clipped, grad_norm) - - self.assertIsNotNone(grad_norm) - expected_norm = mx.sqrt((N * 10 * 8 + N * 20) * 100.0) - self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4)) - self.assertEqual(new_params_clipped["w1"].shape, (N * 10, 8)) - self.assertEqual(new_params_clipped["w2"].shape, (N * 20,)) - - scale = 1.0 / expected_norm - expected_update = 1.0 - 0.1 * 10.0 * scale - self.assertTrue( - mx.allclose( - new_params_clipped["w1"], - mx.ones((N * 10, 8)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - self.assertTrue( - mx.allclose( - new_params_clipped["w2"], - mx.ones((N * 20,)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - params = {"w": mx.ones((N * 4,))} - grads = {"w": mx.ones((N * 4,)) * 0.5} - - optimizer_fsdp = optim.SGD(learning_rate=0.1) - updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer_fsdp) - - optimizer_ddp = optim.SGD(learning_rate=0.1) - avg_grads = average_gradients(grads) - updated_params_ddp = optimizer_ddp.apply_gradients(avg_grads, params) - mx.eval(updated_params_ddp, updated_params_fsdp) - - self.assertTrue( - mx.allclose( - updated_params_fsdp["w"], updated_params_ddp["w"], atol=1e-6, rtol=1e-4 - ), - ) + def test_fully_shard_grads(self): + dtypes = [ + (mx.float32, 1e-6, 1e-6), + (mx.bfloat16, 1e-3, 1e-3), + ] - def test_fsdp_ddp_apply_gradients(self): - world = mx.distributed.init() - N = world.size() - S = 4 - fsdp_group = world.split(world.rank() // S) - dp_group = world.split(world.rank() % S) - - self.assertEqual(fsdp_group.size(), S) - self.assertEqual(dp_group.size(), N // S) - - params = { - "w1": mx.ones((S * 10, 8)), - "w2": mx.ones((S * 20,)), - } - grads = { - "w1": mx.ones((S * 10, 8)) * 0.1, - "w2": mx.ones((S * 20,)) * 0.1, - } - - optimizer = optim.SGD(learning_rate=0.1) - updated = fsdp_apply_gradients( - grads, - params, - optimizer, - fsdp_group=fsdp_group, - dp_group=dp_group, - ) - mx.eval(updated) - - self.assertEqual(updated["w1"].shape, (S * 10, 8)) - self.assertEqual(updated["w2"].shape, (S * 20,)) - - self.assertTrue( - mx.allclose(updated["w1"], mx.ones((S * 10, 8)) * 0.99, atol=1e-6) - ) - self.assertTrue( - mx.allclose(updated["w2"], mx.ones((S * 20,)) * 0.99, atol=1e-6) - ) - - grads_big = { - "w1": mx.ones((S * 10, 8)) * 10.0, - "w2": mx.ones((S * 20,)) * 10.0, - } - - optimizer2 = optim.SGD(learning_rate=0.1) - clipped, grad_norm = fsdp_apply_gradients( - grads_big, - params, - optimizer2, - fsdp_group=fsdp_group, - dp_group=dp_group, - max_norm=1.0, - ) - mx.eval(clipped, grad_norm) - - self.assertIsNotNone(grad_norm) - expected_norm = mx.sqrt((S * 10 * 8 + S * 20) * 100.0) - self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4)) - self.assertEqual(clipped["w1"].shape, (S * 10, 8)) - self.assertEqual(clipped["w2"].shape, (S * 20,)) - - scale = 1.0 / expected_norm - expected_update = 1.0 - 0.1 * 10.0 * scale - self.assertTrue( - mx.allclose( - clipped["w1"], - mx.ones((S * 10, 8)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - self.assertTrue( - mx.allclose( - clipped["w2"], - mx.ones((S * 20,)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - - params_eq = {"w": mx.ones((S * 4,))} - grads_eq = {"w": mx.ones((S * 4,)) * 0.5} - - optimizer_hybrid = optim.SGD(learning_rate=0.1) - updated_hybrid = fsdp_apply_gradients( - grads_eq, - params_eq, - optimizer_hybrid, - fsdp_group=fsdp_group, - dp_group=dp_group, - ) - - optimizer_ddp = optim.SGD(learning_rate=0.1) - avg_grads = average_gradients(grads_eq) - updated_ddp = optimizer_ddp.apply_gradients(avg_grads, params_eq) - mx.eval(updated_hybrid, updated_ddp) - - self.assertTrue( - mx.allclose(updated_hybrid["w"], updated_ddp["w"], atol=1e-6, rtol=1e-4), - ) - - def test_fsdp_peak_memory(self): world = mx.distributed.init() N = world.size() - mx.random.seed(42) - params = { - "w1": mx.random.normal((N * 1024, 1024)), - "w2": mx.random.normal((N * 2048, 512)), - } - grads = { - "w1": mx.random.normal((N * 1024, 1024)), - "w2": mx.random.normal((N * 2048, 512)), - } - mx.eval(params, grads) - optimizer_ddp = optim.Adam(learning_rate=0.01) - optimizer_fsdp = optim.Adam(learning_rate=0.01) - - def pseudo_step_ddp(grads, params, optimizer): - grads = average_gradients(grads) - grads, grad_norm = optim.clip_grad_norm(grads, max_norm=1.0) - params = optimizer.apply_gradients(grads, params) - return grad_norm, params - - def pseudo_step_fsdp(grads, params, optimizer): - params, grad_norm = fsdp_apply_gradients( - grads, params, optimizer, max_norm=1.0 + rank = world.rank() + dims = 8 * N + part = slice(rank * dims // N, (rank + 1) * dims // N) + + class MLP(nn.Module): + def __init__(self, dims): + super().__init__() + self.l1 = nn.Linear(dims, dims) + self.l2 = nn.Linear(dims, dims) + + def __call__(self, x): + return self.l2(nn.relu(self.l1(x))) + + def loss_fn(model, x, y): + logits = model(x).astype(mx.float32) + return ((logits - y) ** 2).mean() + + for dtype, atol, rtol in dtypes: + + mx.random.seed(0xF0F0F0F0) + + kx, ky = mx.random.split(mx.random.key(rank)) + x = mx.random.normal((4, dims), dtype=dtype, key=kx) + y = mx.random.normal((4, dims), key=ky) + + # DDP reference: replicated params, gradients averaged across ranks + model = MLP(dims) + params = model.trainable_parameters() + + # Cast parameters to dtype for forward + model.update(tree_map(lambda p: p.astype(dtype), params)) + loss, grads = mx.value_and_grad(loss_fn)(model, x, y) + grads = average_gradients(grads, group=world) + loss = mx.distributed.all_sum(loss, group=world) / N + mx.eval(loss, grads) + + # Shard the model + model_sharded = MLP(dims) + model_sharded.update(params) + model_sharded = nn.fully_shard(model_sharded, compute_dtype=dtype) + loss_sharded, grads_sharded = mx.value_and_grad(loss_fn)( + model_sharded, x, y ) - return grad_norm, params - - mx.reset_peak_memory() - - for i in range(10): - grad_norm, params = pseudo_step_ddp(grads, params, optimizer_ddp) - mx.eval(grad_norm, params) - - ddp_peak_memory = mx.get_peak_memory() - mx.reset_peak_memory() - - for i in range(10): - grad_norm, params = pseudo_step_fsdp(grads, params, optimizer_fsdp) - mx.eval(grad_norm, params) - fsdp_peak_memory = mx.get_peak_memory() - self.assertTrue(fsdp_peak_memory < ddp_peak_memory) + loss_sharded = mx.distributed.all_sum(loss_sharded, group=world) / N + mx.eval(loss_sharded, grads_sharded) + grads_ref = dict(tree_flatten(grads)) + self.assertTrue(mx.allclose(loss, loss_sharded, atol=1e-4, rtol=1e-4)) + for k, gs in tree_flatten(grads_sharded["module"]): + self.assertTrue( + mx.allclose(gs, grads_ref[k][part], atol=atol, rtol=rtol) + ) if __name__ == "__main__": From de82b176073e5d3a4824bf9a89f37b421f1522c6 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 21 Jul 2026 14:32:47 -0700 Subject: [PATCH 006/222] Fix triplet_loss docstring to document the reduced output shape (#3884) --- python/mlx/nn/losses.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 61b6adb047..85b11255bb 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -416,8 +416,9 @@ def triplet_loss( ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: - array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input; - if reduction is "mean" or "sum", returns a scalar tensor. + array: Computed triplet loss. If reduction is ``"none"``, returns a tensor with the + same shape as the inputs but with the ``axis`` dimension removed; if reduction + is ``"mean"`` or ``"sum"``, returns a scalar tensor. """ pos_dist = mx.power( mx.power(mx.abs(anchors - positives), p).sum(axis) + eps, 1.0 / p From 0c537a413a39616c1d19b96689705719e67be4c1 Mon Sep 17 00:00:00 2001 From: Hao Xu Date: Tue, 21 Jul 2026 14:46:36 -0700 Subject: [PATCH 007/222] Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory (#3872) Co-authored-by: Angelos Katharopoulos --- python/src/convert.cpp | 103 ++++++++++++++++++++++++++------- python/tests/test_array.py | 33 +++++++---- python/tests/test_zero_copy.py | 88 ++++++++++++++++++++++++++++ 3 files changed, 192 insertions(+), 32 deletions(-) create mode 100644 python/tests/test_zero_copy.py diff --git a/python/src/convert.cpp b/python/src/convert.cpp index 1b8145094f..3b161b99e6 100644 --- a/python/src/convert.cpp +++ b/python/src/convert.cpp @@ -172,6 +172,52 @@ mx::array cpu_nd_array_to_mlx( return out; } +// Try to adopt a CPU host buffer as an mlx array without copying. On unified +// memory the host pointer is GPU-addressable, so we wrap it directly via the +// allocator instead of copying. The bytes are reinterpreted rather than +// converted, so the source element width must already match the destination +// dtype. The source ndarray is kept alive for the lifetime of the returned +// array. +// +// Returns std::nullopt when the buffer cannot be adopted (no Metal backend, +// dtype width mismatch, or a pointer the platform will not wrap), so the caller +// can fall back to a copy or raise. +std::optional cpu_nd_array_to_mlx_no_copy( + nb::ndarray nd_array, + const mx::Shape& shape, + mx::Dtype dst_dtype) { + if (!mx::metal::is_available() || + nd_array.itemsize() != mx::size_of(dst_dtype)) { + return std::nullopt; + } + + auto [storage_size, strides, flags] = get_strided_layout(nd_array, shape); + auto buf = mx::allocator::make_buffer( + const_cast(nd_array.data()), + storage_size * mx::size_of(dst_dtype)); + // make_buffer returns a null buffer when the pointer cannot be wrapped, e.g. + // when its alignment is not accepted by the platform. + if (buf.ptr() == nullptr) { + return std::nullopt; + } + + mx::array out(shape, dst_dtype, nullptr, {}); + out.set_data( + buf, + storage_size, + std::move(strides), + flags, + nd_array.byte_offset(), + // The buffer wraps caller-owned memory, so release the wrapper rather + // than returning it to the allocator's reuse pool, which must only + // recycle buffers it allocated itself. + [owner = std::move(nd_array)](mx::allocator::Buffer b) { + mx::allocator::release(b); + }); + out.set_status(mx::array::Status::available); + return out; +} + mx::array metal_nd_array_to_mlx( nb::ndarray nd_array, mx::Dtype src_dtype, @@ -211,48 +257,63 @@ mx::array nd_array_to_mlx( std::optional src_dlpack_dtype_override, std::optional copy) { auto src_dlpack_dtype = src_dlpack_dtype_override.value_or(nd_array.dtype()); - auto src_mlx_dtype = - mlx_dtype_from_dlpack(src_dlpack_dtype, "Cannot convert array to mlx."); + auto src_mlx_dtype = mlx_dtype_from_dlpack( + src_dlpack_dtype, "[convert] Cannot convert array to mlx."); auto dst_dtype = requested_dtype.value_or(src_mlx_dtype); auto device_type = nd_array.device_type(); - // CPU ndarrays are copied below, and their data_handle() is a host pointer, - // not a GPU buffer handle that can be queried by the active allocator. - bool can_reuse_buffer = device_type == nb::device::cpu::value - ? true - : mx::allocator::can_reuse_alien_buffer(nd_array.data_handle()); - bool should_copy = - copy.value_or(false) || dst_dtype != src_mlx_dtype || !can_reuse_buffer; - if (copy.has_value() && copy.value() == false && dst_dtype != src_mlx_dtype) { + + // A dtype change requires converting the elements, which cannot be done + // without a copy. + bool no_copy = copy.has_value() && !copy.value(); + if (no_copy && dst_dtype != src_mlx_dtype) { throw std::invalid_argument( - "Cannot convert DLPack array to requested dtype without a copy."); + "[convert] Cannot convert array to the requested dtype without a " + "copy."); } + switch (device_type) { case nb::device::cpu::value: { - if (copy.has_value() && copy.value() == false) { - throw std::invalid_argument( - "Cannot import a CPU DLPack array without a copy."); - } auto shape = get_shape(nd_array); + // For copy=None (try to share) and copy=False (must share), attempt a + // zero-copy adoption of the host buffer first. A copy is passed by value + // so the source is preserved for the fallback below. + if (!copy.value_or(false)) { + if (auto out = + cpu_nd_array_to_mlx_no_copy(nd_array, shape, dst_dtype)) { + return *out; + } + if (no_copy) { + throw std::invalid_argument( + "[convert] Cannot import a CPU array without a copy."); + } + } + // copy=True, or copy=None where adoption was not possible: copy. return dispatch_dlpack_dtype( src_dlpack_dtype, - [&](mx::Dtype src_dtype) { + [&](mx::Dtype) { return cpu_nd_array_to_mlx(nd_array, shape, dst_dtype); }, - "Cannot convert numpy array to mlx array."); + "[convert] Cannot convert array to mlx."); } case nb::device::metal::value: { - if (copy.has_value() && copy.value() == false && !can_reuse_buffer) { + // A Metal buffer can be adopted without a copy only if the active + // allocator recognizes it. + bool can_reuse_buffer = + mx::allocator::can_reuse_alien_buffer(nd_array.data_handle()); + if (no_copy && !can_reuse_buffer) { throw std::invalid_argument( - "Cannot import a private Metal DLPack buffer without a copy."); + "[convert] Cannot import a private Metal buffer without a copy."); } + bool should_copy = copy.value_or(false) || dst_dtype != src_mlx_dtype || + !can_reuse_buffer; return metal_nd_array_to_mlx( nd_array, src_mlx_dtype, dst_dtype, should_copy); } case nb::device::cuda::value: case nb::device::cuda_managed::value: - throw std::invalid_argument("CUDA DLPack import is not supported."); + throw std::invalid_argument("[convert] CUDA import is not supported."); default: - throw std::invalid_argument("Unsupported DLPack device."); + throw std::invalid_argument("[convert] Unsupported device."); } } diff --git a/python/tests/test_array.py b/python/tests/test_array.py index dfab616716..04e17f8710 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -2129,16 +2129,26 @@ def __dlpack__(self, *args, **kwargs): def test_from_dlpack_cpu(self): x = np.arange(3, dtype=np.float32) + # copy=None may adopt the buffer or copy; either way the values match + # the source at import time. y = mx.from_dlpack(x) - x += 10 self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + # copy=True always copies, so later mutations of the source are not seen. y = mx.from_dlpack(x, copy=True) x += 10 - self.assertEqual(y.tolist(), [10.0, 11.0, 12.0]) + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) - with self.assertRaises(ValueError): - mx.from_dlpack(x, copy=False) + # copy=False adopts the buffer when possible and raises otherwise; it + # must never silently copy. + x = np.arange(3, dtype=np.float32) + try: + y = mx.from_dlpack(x, copy=False) + x += 10 + except ValueError: + pass + else: + self.assertEqual(y.tolist(), [10.0, 11.0, 12.0]) def test_dlpack_cpu_dtype_mapping(self): class CpuDLPack: @@ -2214,17 +2224,12 @@ def test_from_dlpack_cpu_strided(self): self.assertEqual(y.tolist(), view.tolist()) self.assertFalse(memoryview(y).c_contiguous) self.assertEqual(memoryview(y).strides, view.strides) - expected = view.copy().tolist() - x[0, 0] = 99 - self.assertEqual(y.tolist(), expected) stepped = np.arange(20, dtype=np.int32)[2:10:2] y = mx.from_dlpack(stepped) self.assertEqual(y.tolist(), [2, 4, 6, 8]) self.assertFalse(memoryview(y).c_contiguous) self.assertEqual(memoryview(y).strides, stepped.strides) - stepped[0] = 99 - self.assertEqual(y.tolist(), [2, 4, 6, 8]) broadcast = np.broadcast_to(np.array([7], dtype=np.int32), (3,)) y = mx.from_dlpack(broadcast) @@ -2649,8 +2654,14 @@ def test_asarray(self): mx_arr = mx.asarray(np_arr) self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0]) self.assertEqual(mx_arr.dtype, mx.float32) - with self.assertRaises(ValueError): - mx.asarray(np_arr, copy=False) + # copy=False adopts the buffer when possible and raises otherwise; it + # must never silently copy. + try: + mx_arr = mx.asarray(np_arr, copy=False) + except ValueError: + pass + else: + self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0]) with self.assertRaises(ValueError): mx.asarray([1, 2, 3], copy=False) diff --git a/python/tests/test_zero_copy.py b/python/tests/test_zero_copy.py new file mode 100644 index 0000000000..231e2e45ee --- /dev/null +++ b/python/tests/test_zero_copy.py @@ -0,0 +1,88 @@ +# Copyright © 2024 Apple Inc. + +import gc +import unittest + +import mlx.core as mx +import mlx_tests +import numpy as np + + +class TestZeroCopy(mlx_tests.MLXTestCase): + """Tests for zero-copy CPU import: mx.asarray(host_buffer, copy=False). + + On unified memory (Metal) a page-aligned CPU buffer is adopted instead of + copied. On backends without Metal, or for a non-page-aligned buffer or a + dtype conversion, copy=False raises. + """ + + def test_default_shares_or_copies(self): + # With copy=None MLX adopts the buffer when it can and copies otherwise, + # so either way the values must match the source at import time. + a = np.arange(1_000_000, dtype=np.int32) + x = mx.asarray(a) + mx.eval(x) + self.assertTrue(np.array_equal(np.array(x), a)) + + def test_copy_true_copies(self): + a = np.arange(1_000_000, dtype=np.int32) + x = mx.asarray(a, copy=True) + a[0] = 12345 + mx.eval(x) + self.assertNotEqual(int(x[0]), 12345) + + def test_copy_false(self): + a = np.arange(1_000_000, dtype=np.int32) + if not mx.metal.is_available(): + with self.assertRaises(Exception): + mx.asarray(a, copy=False) + return + if a.ctypes.data % 16384 != 0: + self.skipTest("source buffer not page-aligned; adopt path not taken") + x = mx.asarray(a, copy=False) + self.assertTrue(np.array_equal(np.array(x), a)) + # Zero-copy adoption: a mutation of the source is visible in the array. + a[1] = 999 + mx.eval(x) + self.assertEqual(int(x[1]), 999) + + def test_copy_false_dtype_conversion_raises(self): + a = np.arange(16, dtype=np.float64) + with self.assertRaises(Exception): + mx.asarray(a, dtype=mx.float32, copy=False) + + def test_source_lifetime(self): + if not mx.metal.is_available(): + self.skipTest("copy=False requires Metal") + + def make(): + a = np.arange(1_000_000, dtype=np.float32) + 0.5 + if a.ctypes.data % 16384 != 0: + return None + return mx.asarray(a, copy=False) + + x = make() + if x is None: + self.skipTest("source buffer not page-aligned") + gc.collect() + mx.eval(x + 1) + self.assertAlmostEqual(float(x[10]), 10.5, places=5) + + def test_adopt_in_loop_not_recycled(self): + # Regression: an adopted buffer must be released (not recycled into the + # allocator's reuse pool) when freed. Otherwise, over many iterations the + # pool hands a caller-owned buffer to an unrelated array and corrupts / + # crashes. Adopt fresh buffers in a loop and compute after each. + if not mx.metal.is_available(): + self.skipTest("copy=False requires Metal") + w = mx.random.normal((64, 64)) + for i in range(200): + a = np.random.rand(256, 64).astype(np.float32) + x = mx.asarray(a, copy=False) # adopt; x (and a) freed next iteration + r = mx.sum(x @ w) + mx.eval(r) + self.assertTrue(True) # reaching here without crashing is the assertion + + +if __name__ == "__main__": + unittest.main() From 8462ad9fd210518341f36f42e997cfafb7c528db Mon Sep 17 00:00:00 2001 From: pierre427 Date: Wed, 22 Jul 2026 00:20:39 -0400 Subject: [PATCH 008/222] Round MLX_SDPA_BLOCKS up to a multiple of 32 (#3875) --- .../metal/scaled_dot_product_attention.cpp | 4 +++- python/tests/test_fast_sdpa.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index d6f74bc2a9..ee5e09c870 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -475,7 +475,9 @@ void sdpa_vector_2pass( } } if (int blocks_env = env::get_var("MLX_SDPA_BLOCKS", 0); blocks_env > 0) { - blocks = blocks_env; + // The 2-pass reduction consumes the partials in simd-width (32) chunks + // and silently drops the tail otherwise, so round up to a multiple of 32. + blocks = ((blocks_env + 31) / 32) * 32; } size_t k_head_stride = k.shape(1) == 1 ? k.strides(0) : k.strides(1); size_t k_seq_stride = k.strides()[2]; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 7bd867084e..f7440a0f8a 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -360,6 +360,24 @@ def test_sdpa_vector_batched(self): ref = mlx_ref_attn(q, k, v, mask=mask) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_blocks_env_override(self): + # MLX_SDPA_BLOCKS used to be applied as-is, and values that are not + # a multiple of 32 silently corrupted the 2-pass vector output. The + # override is now rounded up to a multiple of 32. + D = 128 + q = mx.random.normal(shape=(1, 32, 1, D), dtype=mx.float16) + k = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) + v = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) + ref = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) + try: + for blocks in (16, 33, 48, 100): + os.environ["MLX_SDPA_BLOCKS"] = str(blocks) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + finally: + del os.environ["MLX_SDPA_BLOCKS"] + @unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU") def test_sdpa(self): # fmt: off From 291e909f443d17509385ee88fded90e460016a86 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 22 Jul 2026 00:48:53 -0400 Subject: [PATCH 009/222] Reuse Metal WAR tracking hash tables (#3882) --- mlx/backend/metal/device.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 29eecf5562..197b919b1d 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -392,8 +392,9 @@ void CommandEncoder::maybeInsertBarrier() { if (needs_barrier_) { get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); needs_barrier_ = false; - prev_inputs_ = std::move(next_inputs_); - prev_outputs_ = std::move(next_outputs_); + // Preserve the hash tables' buckets for reuse across barrier epochs. + prev_inputs_.swap(next_inputs_); + prev_outputs_.swap(next_outputs_); } else { prev_inputs_.insert(next_inputs_.begin(), next_inputs_.end()); prev_outputs_.insert(next_outputs_.begin(), next_outputs_.end()); From 3541c66b9d5829264ccfc3b0323998e316df69ac Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Tue, 21 Jul 2026 22:12:50 -0700 Subject: [PATCH 010/222] Use unroll_count(4) for the NAX attention Q@K.T loop (#3843) --- .../metal/kernels/steel/attn/kernels/steel_attention_nax.h | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index adc9a42798..6c85da64f6 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -207,7 +207,11 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { - STEEL_PRAGMA_UNROLL + // Unrolling the head-dim loop by 4, rather than fully, potentially + // lets the compiler interleave the next K-tile loads with the running + // mma chain instead of hoisting all TD loads up front and is faster + // for head dim 128. +#pragma clang loop unroll_count(4) for (short id = 0; id < TD; id++) { NAXTile Qtile; NAXTile Ktile; From 0ebcee8db75f6482b6e673f4d850287dbdfe359a Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 22 Jul 2026 08:22:46 -0700 Subject: [PATCH 011/222] metal: add gemv_wide for fp16/bf16 matmuls of a few rows (#3888) --- mlx/backend/metal/jit_kernels.cpp | 46 ++++ mlx/backend/metal/kernels.h | 18 ++ mlx/backend/metal/kernels/gemv.h | 278 +++++++++++++++++++++ mlx/backend/metal/kernels/gemv.metal | 36 +++ mlx/backend/metal/matmul.cpp | 354 +++++++++++++++++++++++++++ mlx/backend/metal/nojit_kernels.cpp | 22 ++ python/tests/test_blas.py | 204 +++++++++++++++ 7 files changed, 958 insertions(+) diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 9ed9ce43e0..d639719d81 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -803,6 +803,52 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + "gemv_wide", + get_type_string(out.dtype()), + vecs_per_tg, + k_lanes); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + "gemv_wide_gather", + get_type_string(out.dtype()), + vecs_per_tg, + k_lanes); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + MTL::ComputePipelineState* get_steel_conv_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index cb289fc1e7..973041932a 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -253,6 +253,24 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( int tn, bool contiguous); +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes); + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes); + MTL::ComputePipelineState* get_steel_conv_general_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels/gemv.h b/mlx/backend/metal/kernels/gemv.h index 84579516ec..e5535bf016 100644 --- a/mlx/backend/metal/kernels/gemv.h +++ b/mlx/backend/metal/kernels/gemv.h @@ -764,3 +764,281 @@ template < simd_gid, simd_lid); } + +/////////////////////////////////////////////////////////////////////////////// +/// Multi-vector matrix-vector multiplication (wide gemv) +/////////////////////////////////////////////////////////////////////////////// + +constant bool gemv_wide_has_batch [[function_constant(0)]]; +constant bool gemv_wide_do_axpby [[function_constant(1)]]; + +// out[M, N] = x[M, K] @ mat[N, K]^T for small M: each threadgroup streams a +// block of matrix rows once and applies it to vecs_per_tg input vectors, so +// the matrix is read ceil(M / vecs_per_tg) times instead of once per padded +// GEMM tile. k_lanes lanes reduce K for one row, 32 / k_lanes rows per +// simdgroup, k_lanes / 8 simdgroups per threadgroup. +template < + typename T, + const int vecs_per_tg, + const int k_lanes, + typename AccT = typename DefaultAccT::type> +struct GemvWide { + static constexpr constant int unroll = 8; + static constexpr constant int num_simdgroups = k_lanes / 8; + + static METAL_FUNC void run( + const device T* mat, + const device T* in_vec, + const device T* bias, + device T* out_vec, + int in_vec_size, + int out_vec_size, + int M, + int matrix_ld, + int vector_ld, + float alpha, + float beta, + int bias_ld, + int bias_fd, + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int rows_per_simdgroup = 32 / k_lanes; + constexpr int rows_per_tg = rows_per_simdgroup * num_simdgroups; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = + simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = + tid.y * rows_per_tg + rows_per_simdgroup * simd_gid + sg_row; + + // Clamped tail rows/vectors read valid memory; the guarded writes + // below never store them. + const int row = min(out_row, out_vec_size - 1); + const device T* wrow = mat + int64_t(row) * matrix_ld; + const device vec* w4 = (const device vec*)wrow; + const int n_v4 = in_vec_size / 4; + const int n_main = n_v4 - n_v4 % (k_lanes * unroll); + + // Vector chunks beyond grid.x round-robin onto the same threadgroups: + // re-walking the row block hits cache where an extra grid column would + // re-stream it from DRAM. + const device vec* x4 = (const device vec*)in_vec; + const int x_ld4 = vector_ld / 4; + const int n_chunks = (M + vecs_per_tg - 1) / vecs_per_tg; + for (int chunk = tid.x; chunk < n_chunks; chunk += tgpg.x) { + const int vec0 = chunk * vecs_per_tg; + + int x_off4[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + x_off4[v] = min(vec0 + v, M - 1) * x_ld4; + } + + AccT result[vecs_per_tg] = {0}; + + // Adjacent lanes read adjacent blocks so every transaction lands on + // whole cachelines; the unroll keeps loads in flight on short rows. + for (int base = 0; base < n_main; base += k_lanes * unroll) { + vec wf[unroll]; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + wf[i] = vec(w4[base + i * k_lanes + k_lane]); + } + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + AccT acc = 0; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + acc += + dot(wf[i], + vec(x4[x_off4[v] + base + i * k_lanes + k_lane])); + } + result[v] += acc; + } + } + for (int idx = n_main + k_lane; idx < n_v4; idx += k_lanes) { + const vec wf = vec(w4[idx]); + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + result[v] += dot(wf, vec(x4[x_off4[v] + idx])); + } + } + + // The halving shuffles reduce each row's k_lanes while rows sharing + // the simdgroup stay separate. + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + MLX_MTL_PRAGMA_UNROLL + for (ushort off = k_lanes / 2; off >= 1; off >>= 1) { + result[v] += simd_shuffle_down(result[v], off); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + int64_t out_idx = int64_t(vec0 + v) * out_vec_size + out_row; + if (gemv_wide_do_axpby) { + AccT bias_val = static_cast( + bias[int64_t(vec0 + v) * bias_ld + out_row * bias_fd]); + out_vec[out_idx] = + static_cast(alpha * result[v] + beta * bias_val); + } else { + out_vec[out_idx] = static_cast(result[v]); + } + } + } + } + } + } +}; + +template +[[kernel]] void gemv_wide( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2), function_constant(gemv_wide_do_axpby)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant float& alpha + [[buffer(9), function_constant(gemv_wide_do_axpby)]], + const constant float& beta + [[buffer(10), function_constant(gemv_wide_do_axpby)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* vector_batch_stride [[buffer(13)]], + const constant int64_t* matrix_batch_stride [[buffer(14)]], + const constant int64_t* bias_batch_stride + [[buffer(15), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_ld + [[buffer(16), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_fd + [[buffer(17), function_constant(gemv_wide_do_axpby)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (gemv_wide_has_batch) { + in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim); + mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim); + + if (gemv_wide_do_axpby) { + bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim); + } + } else { + in_vec += tid.z * vector_batch_stride[0]; + mat += tid.z * matrix_batch_stride[0]; + + if (gemv_wide_do_axpby) { + bias += tid.z * bias_batch_stride[0]; + } + } + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + alpha, + beta, + bias_ld, + bias_fd, + tid, + tgpg, + simd_gid, + simd_lid); +} + +template +[[kernel]] void gemv_wide_gather( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* index_batch_strides [[buffer(13)]], + const constant int& vector_batch_ndim [[buffer(14)]], + const constant int* vector_batch_shape [[buffer(15)]], + const constant int64_t* vector_batch_stride [[buffer(16)]], + const constant int& matrix_batch_ndim [[buffer(17)]], + const constant int* matrix_batch_shape [[buffer(18)]], + const constant int64_t* matrix_batch_stride [[buffer(19)]], + const constant uint32_t* vec_indices [[buffer(20)]], + const constant uint32_t* mat_indices [[buffer(21)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + uint32_t indx_vec; + uint32_t indx_mat; + + // Update batch offsets + if (batch_ndim > 1) { + const constant auto* veci_bstrides = index_batch_strides; + const constant auto* mati_bstrides = index_batch_strides + batch_ndim; + + ulong2 batch_offsets = elem_to_loc_broadcast( + tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim); + + indx_vec = vec_indices[batch_offsets.x]; + indx_mat = mat_indices[batch_offsets.y]; + + } else { + indx_vec = vec_indices[index_batch_strides[0] * tid.z]; + indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z]; + } + + if (vector_batch_ndim > 1) { + in_vec += elem_to_loc( + indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim); + } else { + in_vec += indx_vec * vector_batch_stride[0]; + } + + if (matrix_batch_ndim > 1) { + mat += elem_to_loc( + indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim); + } else { + mat += indx_mat * matrix_batch_stride[0]; + } + + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + 1.0f, + 0.0f, + 0, + 0, + tid, + tgpg, + simd_gid, + simd_lid); +} diff --git a/mlx/backend/metal/kernels/gemv.metal b/mlx/backend/metal/kernels/gemv.metal index b306671559..5637ecf732 100644 --- a/mlx/backend/metal/kernels/gemv.metal +++ b/mlx/backend/metal/kernels/gemv.metal @@ -110,3 +110,39 @@ instantiate_gemv_t_bs_blocks(float32, float); instantiate_gemv_t_bs_blocks(float16, half); instantiate_gemv_t_bs_blocks(bfloat16, bfloat16_t); instantiate_gemv_t_bs_blocks(complex64, complex64_t); // clang-format on + +// Batch addressing and axpby are function constants. M == 1 stays with +// gemv, so tiles start at two vectors. +// clang-format off +#define instantiate_gemv_wide_helper(name, itype, nv, kl) \ + instantiate_kernel( \ + "gemv_wide_" #name "_nv" #nv "_kl" #kl, \ + gemv_wide, itype, nv, kl) + +#define instantiate_gemv_wide(name, itype, kl) \ + instantiate_gemv_wide_helper(name, itype, 2, kl) \ + instantiate_gemv_wide_helper(name, itype, 3, kl) \ + instantiate_gemv_wide_helper(name, itype, 4, kl) \ + instantiate_gemv_wide_helper(name, itype, 5, kl) // clang-format on + +instantiate_gemv_wide(float16, half, 16); +instantiate_gemv_wide(bfloat16, bfloat16_t, 16); +instantiate_gemv_wide(float16, half, 32); +instantiate_gemv_wide(bfloat16, bfloat16_t, 32); + +// clang-format off +#define instantiate_gemv_wide_gather_helper(name, itype, nv, kl) \ + instantiate_kernel( \ + "gemv_wide_gather_" #name "_nv" #nv "_kl" #kl, \ + gemv_wide_gather, itype, nv, kl) + +#define instantiate_gemv_wide_gather(name, itype, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 2, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 3, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 4, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 5, kl) // clang-format on + +instantiate_gemv_wide_gather(float16, half, 16); +instantiate_gemv_wide_gather(bfloat16, bfloat16_t, 16); +instantiate_gemv_wide_gather(float16, half, 32); +instantiate_gemv_wide_gather(bfloat16, bfloat16_t, 32); diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 87d2bf52d1..ec3cb10c74 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "mlx/backend/common/broadcasting.h" @@ -1230,6 +1231,170 @@ inline void gemv( /* Strides B_batch_stride = */ B_batch_stride); } +struct GemvWideConfig { + int vecs_per_tg; + int k_lanes; + int grid_x; +}; + +std::string gemv_wide_kernel_name( + const char* prefix, + const array& out, + const GemvWideConfig& config) { + std::ostringstream kname; + kname << prefix << type_to_name(out) << "_nv" << config.vecs_per_tg << "_kl" + << config.k_lanes; + return kname.str(); +} + +inline std::optional gemv_wide_config( + metal::Device& d, + const array& out, + int M, + int N, + int K, + int mat_ld, + int vec_ld, + int64_t mat_offset, + int64_t vec_offset, + const Strides& mat_batch_stride, + const Strides& vec_batch_stride) { + // Pre-M3 generations are limited by load issue rate rather than + // bandwidth and do not profit from the amortized stream; they keep the + // existing kernels. + if (d.get_architecture_gen() < 15) { + return std::nullopt; + } + + auto all_aligned = [](const Strides& strides, int alignment) { + return std::all_of( + strides.begin(), strides.end(), [alignment](int64_t stride) { + return stride % alignment == 0; + }); + }; + + bool vec4_aligned = mat_offset % 8 == 0 && vec_offset % 8 == 0 && + mat_ld % 4 == 0 && vec_ld % 4 == 0 && all_aligned(mat_batch_stride, 4) && + all_aligned(vec_batch_stride, 4); + + // Wide amortizes one matrix stream over several vectors; M == 1 has + // nothing to amortize and stays with gemv. + if (M <= 1 || N <= 1 || K % 4 != 0 || + (out.dtype() != float16 && out.dtype() != bfloat16) || !vec4_aligned) { + return std::nullopt; + } + + // A register tile holds at most five vectors (wider collapses occupancy); + // M balances across passes that each re-stream the matrix, and a fourth + // pass no longer pays. + int passes = (M + 4) / 5; + if (passes > 3) { + return std::nullopt; + } + + // Wide enough that the row grid alone saturates the GPU: N/4 threadgroups, + // orders of magnitude beyond concurrent execution capacity. + bool vocab_wide = N >= 65536; + + // Full rows double the threads per group and halve each lane's K share + // — more in flight where the grid is thinnest: one-pass tiles have a + // single grid column, and tiny outputs starve every column. + bool full_simd = passes == 1 || N <= 64; + int grid_x = vocab_wide ? 1 : passes; + return GemvWideConfig{(M + passes - 1) / passes, full_simd ? 32 : 16, grid_x}; +} + +bool gemv_wide( + const Stream& s, + metal::Device& d, + const array& mat, + const array& in_vec, + const array* bias, + array& out, + int M, + int N, + int K, + int mat_ld, + int vec_ld, + int batch_size_out, + const Shape& batch_shape, + const Strides& mat_batch_stride, + const Strides& vec_batch_stride, + std::vector& copies, + const Strides* bias_batch_stride, + float alpha, + float beta) { + auto plan = gemv_wide_config( + d, + out, + M, + N, + K, + mat_ld, + vec_ld, + mat.offset(), + in_vec.offset(), + mat_batch_stride, + vec_batch_stride); + if (!plan) { + return false; + } + + const bool has_batch = batch_shape.size() != 1; + const bool do_axpby = bias != nullptr; + int batch_ndim = batch_shape.size(); + + MTL::Size group_dims(32, plan->k_lanes / 8, 1); + MTL::Size grid_dims(plan->grid_x, (N + 3) / 4, batch_size_out); + + std::string base_name = gemv_wide_kernel_name("gemv_wide_", out, *plan); + std::ostringstream hash_name; + hash_name << base_name << "_nc" << has_batch << "_axpby" << do_axpby; + metal::MTLFCList func_consts = { + {&has_batch, MTL::DataType::DataTypeBool, 0}, + {&do_axpby, MTL::DataType::DataTypeBool, 1}, + }; + + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gemv_wide_kernel( + d, + base_name, + hash_name.str(), + func_consts, + out, + plan->vecs_per_tg, + plan->k_lanes); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(mat, 0); + compute_encoder.set_input_array(in_vec, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(N, 5); + compute_encoder.set_bytes(M, 6); + compute_encoder.set_bytes(mat_ld, 7); + compute_encoder.set_bytes(vec_ld, 8); + compute_encoder.set_bytes(batch_ndim, 11); + compute_encoder.set_vector_bytes(batch_shape, 12); + compute_encoder.set_vector_bytes(vec_batch_stride, 13); + compute_encoder.set_vector_bytes(mat_batch_stride, 14); + + if (do_axpby) { + compute_encoder.set_input_array(*bias, 2); + compute_encoder.set_bytes(alpha, 9); + compute_encoder.set_bytes(beta, 10); + compute_encoder.set_vector_bytes(*bias_batch_stride, 15); + int bias_ld = bias->strides()[bias->ndim() - 2]; + int bias_fd = bias->strides()[bias->ndim() - 1]; + compute_encoder.set_bytes(bias_ld, 16); + compute_encoder.set_bytes(bias_fd, 17); + } + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + compute_encoder.add_temporaries(std::move(copies)); + return true; +} + /////////////////////////////////////////////////////////////////////////////// // Matmul implementation /////////////////////////////////////////////////////////////////////////////// @@ -1290,6 +1455,32 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { ///////////////////////////////////////////////////////////////////////////// // Gemv specialization + // The wide gemv route streams the weight matrix once per <= 5 input + // vectors instead of running a row-padded GEMM tile. + if (!a_transposed && b_transposed && + gemv_wide( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& mat = */ b, + /* const array& in_vec = */ a, + /* const array* bias = */ nullptr, + /* array& out = */ out, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* int mat_ld = */ b_cols, + /* int vec_ld = */ a_cols, + /* int batch_size_out = */ static_cast(batch_size_out), + /* const Shape& batch_shape = */ batch_shape, + /* const Strides& mat_batch_stride = */ B_batch_stride, + /* const Strides& vec_batch_stride = */ A_batch_stride, + /* std::vector& copies = */ copies, + /* const Strides* bias_batch_stride = */ nullptr, + /* float alpha = */ 1.0f, + /* float beta = */ 0.0f)) { + return; + } + // Route to gemv if needed if (std::min(M, N) == 1) { return gemv( @@ -1422,6 +1613,32 @@ void AddMM::eval_gpu(const std::vector& inputs, array& out) { ///////////////////////////////////////////////////////////////////////////// // Gemv specialization + // The wide gemv route streams the weight matrix once per <= 5 input + // vectors instead of running a row-padded GEMM tile. + if (!transpose_a && transpose_b && + gemv_wide( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& mat = */ b, + /* const array& in_vec = */ a, + /* const array* bias = */ &c, + /* array& out = */ out, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* int mat_ld = */ ldb, + /* int vec_ld = */ lda, + /* int batch_size_out = */ static_cast(batch_size_out), + /* const Shape& batch_shape = */ batch_shape, + /* const Strides& mat_batch_stride = */ B_batch_stride, + /* const Strides& vec_batch_stride = */ A_batch_stride, + /* std::vector& copies = */ copies, + /* const Strides* bias_batch_stride = */ &C_batch_stride, + /* float alpha = */ alpha_, + /* float beta = */ beta_)) { + return; + } + // Route to gemv if needed if (std::min(M, N) == 1) { return gemv_axbpy( @@ -2260,6 +2477,138 @@ void gather_mv( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +// Dispatches gathers whose b is a transposed view (x @ W.T layouts) onto +// gemv_wide_gather; returns false when the layout or shape wants the existing +// routes instead. The route is decided from raw strides so no copy is encoded +// unless it will be dispatched. +bool gather_mm_wide( + const array& a_, + const array& b_, + const array& vec_indices, + const array& mat_indices, + array& out, + int M, + int N, + int K, + metal::Device& d, + const Stream& s) { + // b must already be a transposed view (second-to-last dim contiguous, + // inner stride != 1). check_transpose can't decide this: it copies any + // other layout to row-contiguous and reports it as non-transposed. + if (b_.strides()[b_.ndim() - 1] == 1 || b_.strides()[b_.ndim() - 2] != 1) { + return false; + } + int ldb = b_.strides()[b_.ndim() - 1]; + // a is served directly when row-contiguous, by a contiguous copy when + // neither layout matches, and declines when transposed. + bool a_direct = a_.strides()[a_.ndim() - 1] == 1; + if (!a_direct && a_.strides()[a_.ndim() - 2] == 1) { + return false; + } + int lda; + int64_t vec_offset; + Strides vec_batch_stride; + if (a_direct) { + lda = a_.strides()[a_.ndim() - 2]; + vec_offset = a_.offset(); + vec_batch_stride = Strides(a_.strides().begin(), a_.strides().end() - 2); + } else { + lda = a_.shape(-1); + vec_offset = 0; + vec_batch_stride = Strides(a_.ndim() - 2); + int64_t stride = int64_t(a_.shape(-2)) * a_.shape(-1); + for (int i = a_.ndim() - 3; i >= 0; --i) { + vec_batch_stride[i] = stride; + stride *= a_.shape(i); + } + } + Strides mat_batch_stride(b_.strides().begin(), b_.strides().end() - 2); + + auto cfg = gemv_wide_config( + d, + out, + M, + N, + K, + ldb, + lda, + b_.offset(), + vec_offset, + mat_batch_stride, + vec_batch_stride); + if (!cfg) { + return false; + } + + std::vector copies; + array a = a_; + if (!a_direct) { + a = contiguous_copy_gpu(a_, s); + copies.push_back(a); + } + const array& b = b_; + + int batch_size_out = out.size() / M / N; + int batch_ndim = out.ndim() - 2; + int batch_ndim_vec = a.ndim() - 2; + int batch_ndim_mat = b.ndim() - 2; + Strides index_strides = vec_indices.strides(); + index_strides.insert( + index_strides.end(), + mat_indices.strides().begin(), + mat_indices.strides().end()); + if (index_strides.empty()) { + // Scalar indices: the kernel still loads a stride (scaled by tid.z == 0). + index_strides = {0, 0}; + } + + MTL::Size group_dims(32, cfg->k_lanes / 8, 1); + MTL::Size grid_dims(cfg->grid_x, (N + 3) / 4, batch_size_out); + + const bool do_axpby = false; + std::string base_name = gemv_wide_kernel_name("gemv_wide_gather_", out, *cfg); + std::ostringstream hash_name; + hash_name << base_name << "_axpby0"; + metal::MTLFCList func_consts = { + {&do_axpby, MTL::DataType::DataTypeBool, 1}, + }; + + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gemv_wide_gather_kernel( + d, + base_name, + hash_name.str(), + func_consts, + out, + cfg->vecs_per_tg, + cfg->k_lanes); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(b, 0); + compute_encoder.set_input_array(a, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(N, 5); + compute_encoder.set_bytes(M, 6); + compute_encoder.set_bytes(ldb, 7); + compute_encoder.set_bytes(lda, 8); + compute_encoder.set_bytes(batch_ndim, 11); + compute_encoder.set_vector_bytes(out.shape(), 12); + compute_encoder.set_vector_bytes(index_strides, 13); + compute_encoder.set_bytes(batch_ndim_vec, 14); + compute_encoder.set_vector_bytes(a.shape(), 15); + compute_encoder.set_vector_bytes(a.strides(), 16); + compute_encoder.set_bytes(batch_ndim_mat, 17); + compute_encoder.set_vector_bytes(b.shape(), 18); + compute_encoder.set_vector_bytes(b.strides(), 19); + compute_encoder.set_input_array(vec_indices, 20); + compute_encoder.set_input_array(mat_indices, 21); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + compute_encoder.add_temporaries(std::move(copies)); + return true; +} + void gather_mm( const array& a_, const array& b_, @@ -2445,6 +2794,11 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { return; } + // The wide gather route streams each gathered matrix once per <= 5 rows. + if (gather_mm_wide(a, b, lhs_indices, rhs_indices, out, M, N, K, d, s)) { + return; + } + // Route to non specialized gather mm gather_mm(a, b, lhs_indices, rhs_indices, out, M, N, K, d, s); } diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 64cfa39ed5..9f6f8782f5 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -273,6 +273,28 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_steel_conv_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 0db6257f93..3a8183ba8f 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -538,6 +538,94 @@ def test_matrix_vector_edgecases(self): ) self.assertTrue(np.array_equal(c_mlx, c_npy)) + def test_wide_matmul(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible a @ b.T products of a few rows take the wide gemv route on metal; + # cover numerical correctness across the routing boundaries: row tails, + # K not divisible by 4, offset and sliced views, and batching. + # Inputs scale as K**-0.5 so outputs stay O(K**-0.5); atol scales + # with them, above output rounding for every dtype but below a + # dropped reduction block. + def run_test(dtype, shape_a, shape_b, f_np_b, f_mx_b): + with self.subTest(dtype=str(dtype), shape_a=shape_a, shape_b=shape_b): + np.random.seed(7) + scale = shape_a[-1] ** -0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, shape_a).astype(np.float32) + ).astype(dtype) + b_mx = mx.array( + np.random.normal(0.0, scale, shape_b).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + b_np = np.array(b_mx.astype(mx.float32)) + + out_np = a_np @ f_np_b(b_np) + out_mx = (a_mx @ f_mx_b(b_mx)).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + + nt_np = lambda b: b.swapaxes(-1, -2) + nt_mx = lambda b: mx.swapaxes(b, -1, -2) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (1, 2, 3, 5, 8, 11, 16): + for K, N in ( + (64, 128), + (512, 128), + (2048, 256), + (2048, 32), + (2052, 1000), + ): + run_test(dtype, (M, K), (N, K), nt_np, nt_mx) + + # K % 4 != 0 falls back to the general kernels + run_test(dtype, (5, 514), (333, 514), nt_np, nt_mx) + + # sliced weights: leading dimension != K, plus offset views at 16- + # and 8-byte alignment, and a slice with an odd vec4 tail + run_test( + dtype, + (5, 512), + (333, 576), + lambda b: b[:, :512].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, :512], -1, -2), + ) + run_test( + dtype, + (5, 512), + (333, 512), + lambda b: b[7:, :].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[7:, :], -1, -2), + ) + run_test( + dtype, + (3, 512), + (333, 520), + lambda b: b[:, 4:516].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, 4:516], -1, -2), + ) + run_test( + dtype, + (3, 2052), + (129, 2056), + lambda b: b[:, :2052].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, :2052], -1, -2), + ) + + # batched: regular, broadcast weights, and multi-dim batch + run_test(dtype, (4, 3, 512), (4, 257, 512), nt_np, nt_mx) + run_test( + dtype, + (4, 3, 512), + (1, 257, 512), + lambda b: np.broadcast_to(b, (4, 257, 512)).swapaxes(-1, -2), + lambda b: mx.swapaxes(mx.broadcast_to(b, (4, 257, 512)), -1, -2), + ) + run_test(dtype, (2, 3, 5, 512), (2, 3, 129, 512), nt_np, nt_mx) + def test_mismatch_stride_mm(self): np.random.seed(0) a_npy = np.random.normal(0.0, 1.0 / 128, (4, 16, 16)).astype(np.float32) @@ -733,6 +821,61 @@ def test_addmm(self): expected = 0.5 * (a @ b) + 2.0 * c self.assertTrue(mx.allclose(out, expected, rtol=tol, atol=tol)) + def test_wide_addmm(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible few-row addmm shapes take the wide gemv route on metal; cover + # the axpby epilogue against bias shapes, scales, and batching. + def run_test(dtype, B, M, K, N, c_shape, alpha, beta): + with self.subTest(dtype=str(dtype), c_shape=c_shape, alpha=alpha): + np.random.seed(3) + shape_a = (M, K) if B is None else (B, M, K) + shape_b = (N, K) if B is None else (B, N, K) + scale = K**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, shape_a).astype(np.float32) + ).astype(dtype) + b_mx = mx.array( + np.random.normal(0.0, scale, shape_b).astype(np.float32) + ).astype(dtype) + c_mx = mx.array( + np.random.normal(0.0, scale, c_shape).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + b_np = np.array(b_mx.astype(mx.float32)) + c_np = np.array(c_mx.astype(mx.float32)) + + out_np = alpha * (a_np @ b_np.swapaxes(-1, -2)) + beta * c_np + out_mx = mx.addmm( + c_mx, + a_mx, + mx.swapaxes(b_mx, -1, -2), + alpha, + beta, + ).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + atol = 0.05 * (abs(alpha) + abs(beta)) * scale + self.assertTrue(np.allclose(out_mx, out_np, atol=atol)) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (2, 5, 12): + for c_shape in ((250,), (1, 250), (M, 250)): + for alpha, beta in ((1.0, 1.0), (2.5, 0.5), (1.0, 0.0)): + run_test(dtype, None, M, 512, 250, c_shape, alpha, beta) + run_test(dtype, 3, 4, 512, 250, (3, 4, 250), 1.0, 1.0) + + # the epilogue must scale the accumulator before narrowing: a + # product past the fp16 max rescued by alpha stays finite (M = 4 + # keeps the 2-byte shape routed on every supported generation) + with self.subTest(dtype=str(dtype), case="alpha rescue"): + a = mx.full((4, 512), 2.0, dtype=dtype) + b = mx.full((250, 512), 100.0, dtype=dtype) + c = mx.ones((4, 250), dtype=dtype) + out = mx.addmm(c, a, mx.swapaxes(b, -1, -2), 0.125, 0.0) + self.assertTrue(np.allclose(out.astype(mx.float32), 12800.0)) + def test_addmm_grad(self): def make_ref_addmm(alpha, beta): return lambda c, a, b: alpha * (a @ b) + beta * c @@ -1193,6 +1336,67 @@ def test_shape( self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5)) + def test_gather_mm_blocks(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible gathered products with few-row blocks route to the wide + # gather on metal; check block indexing against a per-entry reference. + def run_test(dtype, G, M, K, N, E, idx_shape): + with self.subTest(dtype=str(dtype), G=G, M=M, idx=idx_shape): + np.random.seed(11) + scale = K**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, (G, M, K)).astype(np.float32) + ).astype(dtype) + w_mx = mx.array( + np.random.normal(0.0, scale, (E, N, K)).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + w_np = np.array(w_mx.astype(mx.float32)) + rhs = np.random.randint(0, E, size=idx_shape).astype(np.uint32) + + out_np = np.stack( + [a_np[i % G] @ w_np[r].T for i, r in enumerate(rhs.reshape(-1))] + ).reshape(*idx_shape, M, N) + out_mx = mx.gather_mm( + a_mx.reshape(*idx_shape, M, K), + mx.swapaxes(w_mx, -1, -2), + None, + mx.array(rhs), + ).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (2, 4, 5, 11): + run_test(dtype, 6, M, 2048, 1024, 8, (6,)) + run_test(dtype, 6, 3, 2048, 1024, 8, (2, 3)) + + # scalar indices: batch_ndim == 0 must still bind index strides + with self.subTest(dtype=str(dtype), idx="scalar"): + np.random.seed(11) + scale = 512**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, (4, 512)).astype(np.float32) + ).astype(dtype) + w_mx = mx.array( + np.random.normal(0.0, scale, (8, 129, 512)).astype(np.float32) + ).astype(dtype) + out_np = ( + np.array(a_mx.astype(mx.float32)) + @ np.array(w_mx[3].astype(mx.float32)).T + ) + out_mx = mx.gather_mm( + a_mx, + mx.swapaxes(w_mx, -1, -2), + None, + mx.array(3, dtype=mx.uint32), + ).astype(mx.float32) + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + def test_gather_matmul_grad(self): lhs_indices = mx.array([[7, 6], [4, 1], [0, 2]], dtype=mx.uint32) rhs_indices = mx.array([[2], [0], [1]], dtype=mx.uint32) From 6cae670e9eb3dec1d443e988682fcfc5b25a33fe Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 22 Jul 2026 09:06:25 -0700 Subject: [PATCH 012/222] Fix broken docstring rendering in Linear and RNN (#3890) --- python/mlx/nn/layers/linear.py | 1 - python/mlx/nn/layers/recurrent.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/python/mlx/nn/layers/linear.py b/python/mlx/nn/layers/linear.py index c53ab9f122..a1da07c465 100644 --- a/python/mlx/nn/layers/linear.py +++ b/python/mlx/nn/layers/linear.py @@ -32,7 +32,6 @@ class Linear(Module): y = x W^\top + b - where: where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``. The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, diff --git a/python/mlx/nn/layers/recurrent.py b/python/mlx/nn/layers/recurrent.py index a5d31fd8c2..92a75e8d9b 100644 --- a/python/mlx/nn/layers/recurrent.py +++ b/python/mlx/nn/layers/recurrent.py @@ -33,7 +33,7 @@ class RNN(Module): hidden_size (int): Dimension of the hidden state, ``H``. bias (bool, optional): Whether to use a bias. Default: ``True``. nonlinearity (callable, optional): Non-linearity to use. If ``None``, - then func:`tanh` is used. Default: ``None``. + then :func:`tanh` is used. Default: ``None``. """ def __init__( From 33c03c486c34a7dadab5339563612c9933c4a406 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 22 Jul 2026 09:07:26 -0700 Subject: [PATCH 013/222] Fix Adamax betas docstring and MultiOptimizer filters type (#3889) --- python/mlx/optimizers/optimizers.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index f1ef46773b..ca3fbc395a 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/python/mlx/optimizers/optimizers.py @@ -165,7 +165,7 @@ class MultiOptimizer(Optimizer): Args: optimizers (list[Optimizer]): A list of optimizers to delegate to - filters (list[Callable[[str, array], bool]): A list of predicates that + filters (list[Callable[[str, array], bool]]): A list of predicates that should be one less than the provided optimizers. """ @@ -606,8 +606,9 @@ class Adamax(Adam): Args: learning_rate (float or callable): The learning rate :math:`\lambda`. betas (Tuple[float, float], optional): The coefficients - :math:`(\beta_1, \beta_2)` used for computing running averages of the - gradient and its square. Default: ``(0.9, 0.999)`` + :math:`(\beta_1, \beta_2)` used for computing the running average of + the gradient and the exponentially weighted infinity norm. + Default: ``(0.9, 0.999)`` eps (float, optional): The term :math:`\epsilon` added to the denominator to improve numerical stability. Default: ``1e-8`` """ From 6c0ea7fb012be4ac55ea7c1c73334e3c9d247f56 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:34:20 -0700 Subject: [PATCH 014/222] Fix incorrect nvfp4 quantized_matmul through the split-K path (#3854) --- mlx/backend/metal/quantized.cpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 62d48714e0..94c563070a 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -884,11 +884,15 @@ void qmm_splitk( int current_tgs = n_tiles * m_tiles; int split_k = std::max(1, 512 / current_tgs); - // Cap split_k by the number of quantization groups - split_k = std::min(split_k, K / group_size); - - // Ensure K divides evenly by split_k * group_size - while (split_k > 1 && (K % (split_k * group_size) != 0)) { + // Each K partition must be a whole number of BK-wide (32) K-tiles as well as + // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do + // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's + // group_size=16) would over-read into the next group's weights/scales. + int k_align = group_size > 32 ? group_size : 32; + split_k = std::min(split_k, K / k_align); + + // Ensure K divides evenly by split_k * k_align + while (split_k > 1 && (K % (split_k * k_align) != 0)) { split_k--; } if (split_k <= 1) { From 7c92ce1268b62a43c7eddcb2abc0cab506c7949b Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Thu, 23 Jul 2026 18:27:32 +0200 Subject: [PATCH 015/222] [Metal] Avoid regex in custom kernel name generation (#3869) --- mlx/backend/common/metal_kernel.cpp | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index b1e722b83f..4795f08778 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -1,7 +1,6 @@ // Copyright © 2024 Apple Inc. #include -#include #include #include "mlx/backend/common/compiled.h" @@ -197,6 +196,25 @@ std::string write_template( return template_def.str(); } +std::string make_template_hash(const std::string& template_def) { + std::string template_hash; + template_hash.reserve(template_def.size()); + for (size_t i = 0; i < template_def.size(); ++i) { + auto c = template_def[i]; + if (c == '<' || c == '>') { + template_hash += '_'; + } else if ( + c == ',' && i + 1 < template_def.size() && template_def[i + 1] == ' ') { + template_hash += '_'; + ++i; + } else { + template_hash += c; + } + } + template_hash.pop_back(); + return template_hash; +} + } // namespace CustomKernelFunction metal_kernel( @@ -290,11 +308,8 @@ CustomKernelFunction metal_kernel( std::string kernel_name = "custom_kernel_" + name; std::string template_def = ""; if (!template_args.empty()) { - std::regex disallowed_chars("\\<|\\>|(, )"); template_def = write_template(template_args); - auto template_hash = - std::regex_replace(template_def, disallowed_chars, "_"); - template_hash.pop_back(); + auto template_hash = make_template_hash(template_def); kernel_name += "_"; kernel_name += template_hash; } From 9b40c9d1b8d2b46138431974e785f6abd148bda1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Fri, 24 Jul 2026 00:29:54 +0300 Subject: [PATCH 016/222] Fix prod dtype promotion when reducing a size-1 axis (#3898) --- mlx/ops.cpp | 2 +- python/tests/test_ops.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 849b8081dd..da1901678d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2397,7 +2397,7 @@ array prod( out_type = int32; } auto out = (is_noop) - ? a + ? astype(a, out_type, s) : array( std::move(out_shape), out_type, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 028d9a4573..5d16c6e96b 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -724,6 +724,19 @@ def test_prod(self): self.assertEqual(mx.prod(x, axis=0).tolist(), [3, 6]) self.assertEqual(mx.prod(x, axis=1).tolist(), [2, 9]) + # dtype should not depend on the length of the reduced axis + for dtype, expected in ( + (mx.bool_, mx.int32), + (mx.int8, mx.int32), + (mx.uint8, mx.uint32), + ): + self.assertEqual( + mx.prod(mx.array([2], dtype=dtype), axis=0).dtype, expected + ) + self.assertEqual( + mx.prod(mx.array([2, 2], dtype=dtype), axis=0).dtype, expected + ) + def test_min_and_max(self): x = mx.array( [ From 973e27f82ffe68dbd626cda31ba34997045d1eb7 Mon Sep 17 00:00:00 2001 From: Adam Luz Date: Fri, 24 Jul 2026 03:28:29 -0600 Subject: [PATCH 017/222] [CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions (#3893) --- mlx/backend/cuda/conv/gemm_conv.cu | 17 +++++++++++------ mlx/backend/cuda/conv/gemm_grouped_conv.cu | 17 +++++++++++------ 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/mlx/backend/cuda/conv/gemm_conv.cu b/mlx/backend/cuda/conv/gemm_conv.cu index 6fc9528289..af15ace1ab 100644 --- a/mlx/backend/cuda/conv/gemm_conv.cu +++ b/mlx/backend/cuda/conv/gemm_conv.cu @@ -25,17 +25,19 @@ __global__ void naive_unfold_nd( auto tid = block.group_index(); auto lid = block.thread_index(); - int index_batch = tid.z / out_pixels; // [0, N) - int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out) + // The matrix-row index (N * out_pixels values) rides grid.x, which allows + // up to 2^31-1 blocks; grid.y and grid.z are limited to 65,535. + int index_batch = tid.x / out_pixels; // [0, N) + int index_out_spatial = tid.x % out_pixels; // [0, H_out * W_out) int index_wt_spatial = - tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) + tid.z * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) if (index_wt_spatial >= filter_size / params.C) { return; } in += tid.y; // [0, C) - out += tid.z * filter_size + index_wt_spatial * params.C + tid.y; + out += tid.x * filter_size + index_wt_spatial * params.C + tid.y; bool valid = index_batch < params.N; @@ -106,9 +108,12 @@ array unfold_inputs_nd( dim3 block_dims; block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024); dim3 num_blocks; - num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x); + // mat_M = N * out_pixels can exceed 65,535 (a 256x256 output already has + // 65,536 positions), so it must ride grid.x; the handful of filter-spatial + // blocks go on grid.z. grid.y and grid.z are limited to 65,535. + num_blocks.x = mat_M; num_blocks.y = params.C; - num_blocks.z = mat_M; + num_blocks.z = cuda::ceil_div(wt_spatial_size, block_dims.x); encoder.set_input_array(in); encoder.set_output_array(unfolded); diff --git a/mlx/backend/cuda/conv/gemm_grouped_conv.cu b/mlx/backend/cuda/conv/gemm_grouped_conv.cu index 4060d744d7..82d32aae72 100644 --- a/mlx/backend/cuda/conv/gemm_grouped_conv.cu +++ b/mlx/backend/cuda/conv/gemm_grouped_conv.cu @@ -25,17 +25,19 @@ __global__ void naive_grouped_unfold_transpose_nd( auto tid = block.group_index(); auto lid = block.thread_index(); - int index_batch = tid.z / out_pixels; // [0, N) - int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out) + // The matrix-row index (N * out_pixels values) rides grid.x, which allows + // up to 2^31-1 blocks; grid.y and grid.z are limited to 65,535. + int index_batch = tid.x / out_pixels; // [0, N) + int index_out_spatial = tid.x % out_pixels; // [0, H_out * W_out) int index_wt_spatial = - tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) + tid.z * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) if (index_wt_spatial >= filter_size / params.C) { return; } in += tid.y; // [0, C) - out += tid.z * filter_size + tid.y * (filter_size / params.C); + out += tid.x * filter_size + tid.y * (filter_size / params.C); bool valid = index_batch < params.N; @@ -109,9 +111,12 @@ array grouped_unfold_transpose_inputs_nd( dim3 block_dims; block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024); dim3 num_blocks; - num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x); + // mat_M = N * out_pixels can exceed 65,535 (a 256x256 output already has + // 65,536 positions), so it must ride grid.x; the handful of filter-spatial + // blocks go on grid.z. grid.y and grid.z are limited to 65,535. + num_blocks.x = mat_M; num_blocks.y = params.C; - num_blocks.z = mat_M; + num_blocks.z = cuda::ceil_div(wt_spatial_size, block_dims.x); encoder.set_input_array(in); encoder.set_output_array(unfolded); From e6134a801750b5322940d6786e1cdd03e178bdd0 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 29 Jul 2026 20:26:35 +0200 Subject: [PATCH 018/222] [CUDA] columnwise quantize with tma (#3157) --- examples/python/qqmm.py | 99 +-- mlx/backend/cuda/ptx.cuh | 127 ++++ mlx/backend/cuda/quantized/fp_quantize.cu | 681 ++++++++------------- mlx/backend/cuda/quantized/fp_quantize.cuh | 601 ++++++++++++++++++ mlx/backend/cuda/quantized/qqmm_utils.cu | 19 +- 5 files changed, 1053 insertions(+), 474 deletions(-) create mode 100644 mlx/backend/cuda/ptx.cuh create mode 100644 mlx/backend/cuda/quantized/fp_quantize.cuh diff --git a/examples/python/qqmm.py b/examples/python/qqmm.py index 5be7eae2f3..dc46fa29a0 100644 --- a/examples/python/qqmm.py +++ b/examples/python/qqmm.py @@ -33,48 +33,71 @@ def test_qqmm(): [64, 128, 256, 1024, 1024 * 8], # N [64, 128, 256, 1024, 1024 * 8], # K ) + layouts = ["TN", "NT", "TT", "NN"] for group_size, mode, bits in tests: for M, N, K in product(*shapes): for dtype in dtypes: - x = mx.random.normal(shape=(M, K), key=k1, dtype=dtype) - w = mx.random.normal(shape=(N, K), key=k2, dtype=dtype) - w_q, scales_w = mx.quantize(w, group_size, bits, mode=mode) - w_dq = mx.dequantize( - w_q, - scales_w, - group_size=group_size, - bits=bits, - mode=mode, - dtype=dtype, - ) - y_q = mx.qqmm( - x, - w_q, - scales_w, - group_size=group_size, - bits=bits, - mode=mode, - ) - x_q, scales_x = mx.quantize( - x, group_size=group_size, bits=bits, mode=mode - ) - x_dq = mx.dequantize( - x_q, - scales_x, - group_size=group_size, - bits=bits, - mode=mode, - dtype=dtype, - ) - y_hat = mx.matmul(x_dq, mx.transpose(w_dq)) - ulp = ulp_bf16_at(y_hat) - error = (y_q - y_hat).abs() - if not (mx.logical_or(error < 1e-3, error <= ulp).all()): - raise AssertionError( - f"qqmm test failed for shape {(M, N, K)}, " - f"group_size={group_size}, bits={bits}, " - f"mode={mode}, dtype={dtype}" + for layout in layouts: + if layout == "NT": + x_shape = (M, K) + w_shape = (N, K) + elif layout == "TN": + x_shape = (K, M) + w_shape = (K, N) + elif layout == "TT": + x_shape = (K, M) + w_shape = (N, K) + else: # "NN" + x_shape = (M, K) + w_shape = (K, N) + + x = mx.random.normal(shape=x_shape, key=k1, dtype=dtype) + w = mx.random.normal(shape=w_shape, key=k2, dtype=dtype) + + if layout == "TT": + x = mx.transpose(x) + elif layout == "TN": + w = mx.transpose(w) + x = mx.transpose(x) + elif layout == "NN": + w = mx.transpose(w) + + y_q = mx.qqmm( + x, + w, + group_size=group_size, + bits=bits, + mode=mode, + ) + w_q, scales_w = mx.quantize(w, group_size, bits, mode=mode) + w_dq = mx.dequantize( + w_q, + scales_w, + group_size=group_size, + bits=bits, + mode=mode, + dtype=dtype, + ) + x_q, scales_x = mx.quantize( + x, group_size=group_size, bits=bits, mode=mode + ) + x_dq = mx.dequantize( + x_q, + scales_x, + group_size=group_size, + bits=bits, + mode=mode, + dtype=dtype, ) + y_hat = mx.matmul(x_dq, mx.transpose(w_dq)) + ulp = ulp_bf16_at(y_hat) + error = (y_q - y_hat).abs() + if not (mx.logical_or(error < 1e-3, error <= ulp).all()): + raise AssertionError( + f"qqmm test failed for shape {(M, N, K)}, " + f"group_size={group_size}, bits={bits}, " + f"mode={mode}, dtype={dtype}, layout={layout}" + ) def test_qqmm_vjp(): diff --git a/mlx/backend/cuda/ptx.cuh b/mlx/backend/cuda/ptx.cuh new file mode 100644 index 0000000000..6ec7caedd6 --- /dev/null +++ b/mlx/backend/cuda/ptx.cuh @@ -0,0 +1,127 @@ +#pragma once + +#include +#include + +namespace mlx::core { + +namespace ptx { + +#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ + defined(__CUDA_ARCH_SPECIFIC__) + +__device__ __forceinline__ void mbarrier_init(uint64_t* mbar, uint32_t count) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + : + : "r"(mbar_ptr), "r"(count) + : "memory"); +} + +__device__ __forceinline__ void mbarrier_invalidate(uint64_t* mbar) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.inval.shared.b64 [%0];" : : "r"(mbar_ptr) : "memory"); +} + +// Arrive at barrier (non-master threads) +__device__ __forceinline__ void mbarrier_arrive(uint64_t* mbar) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.arrive.shared.b64 _, [%0];" + : + : "r"(mbar_ptr) + : "memory"); +} + +// Arrive at barrier and set expected transaction count (master thread) +__device__ __forceinline__ void mbarrier_arrive_expect_tx( + uint64_t* mbar, + uint32_t tx_count) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" + : + : "r"(mbar_ptr), "r"(tx_count) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-asynchronous-copy-completion-mechanisms-mbarrier +__device__ __forceinline__ void mbarrier_wait_parity( + uint64_t* mbar, + uint32_t parity) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile( + "{\n\t" + ".reg .pred P;\n\t" + "WAIT_LOOP:\n\t" + "mbarrier.try_wait.parity.shared.b64 P, [%0], %1;\n\t" + "@!P bra WAIT_LOOP;\n\t" + "}\n\t" + : + : "r"(mbar_ptr), "r"(parity) + : "memory"); +} + +// Async bulk tensor copy: global -> shared (2D) +__device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( + void* dst_shmem, + const CUtensorMap* tensor_map, + uint32_t tile_x, + uint32_t tile_y, + uint64_t* mbar) { + uint32_t dst_ptr = __cvta_generic_to_shared(dst_shmem); + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" + : + : "r"(dst_ptr), "l"(tensor_map), "r"(tile_x), "r"(tile_y), "r"(mbar_ptr) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor +// shared::cta -> global +__device__ __forceinline__ void cp_async_bulk_tensor_2d_shared_to_global( + const uint64_t* tensor_map_ptr, + const uint32_t offset_x, + const uint32_t offset_y, + uint64_t* src_shmem) { + uint32_t src_shmem_ptr = __cvta_generic_to_shared(src_shmem); + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%1, %2}], [%3];" :: + "l"(tensor_map_ptr), + "r"(offset_x), + "r"(offset_y), + "r"(src_shmem_ptr) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group +template +__device__ __forceinline__ void cp_async_bulk_wait_group_read() { + if constexpr (N == 0) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + } else if constexpr (N == 1) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } else if constexpr (N == 2) { + asm volatile("cp.async.bulk.wait_group.read 2;"); + } else if constexpr (N == 4) { + asm volatile("cp.async.bulk.wait_group.read 4;"); + } +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group +__device__ __forceinline__ void cp_async_bulk_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + +// Ccreates a memory ordering barrier between generic and async proxies +// details: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#proxies +__device__ __forceinline__ void fence_proxy_async_shared_cta() { + asm volatile("fence.proxy.async.shared::cta;"); +} + +#endif // (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && + // (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000) +} // namespace ptx +} // namespace mlx::core \ No newline at end of file diff --git a/mlx/backend/cuda/quantized/fp_quantize.cu b/mlx/backend/cuda/quantized/fp_quantize.cu index 7ebdc03a2b..b5bfd72f2f 100644 --- a/mlx/backend/cuda/quantized/fp_quantize.cu +++ b/mlx/backend/cuda/quantized/fp_quantize.cu @@ -3,8 +3,7 @@ #include "mlx/backend/common/quantized.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh" -#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh" +#include "mlx/backend/cuda/quantized/fp_quantize.cuh" #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/vector_types.cuh" #include "mlx/dtype_utils.h" @@ -14,384 +13,75 @@ #include #include -constexpr float F8E4M3_MAX = 448.0f; -constexpr float F4E2M1_MAX = 6.0f; - namespace mlx::core { namespace cu { -template -struct Dequantize { - __device__ float operator()(uint8_t x) { - if constexpr (bits == 8) { - return float(*(cutlass::float_e4m3_t*)(&x)); - } else { - return float(*(cutlass::float_e2m1_t*)(&x)); - } - } -}; - -template -__device__ __forceinline__ void absmax_x2(T& out, const T& x1, const T& x2) { - if constexpr ( - (std::is_same::value) || - (std::is_same::value)) { - T a = x1; - T b = x2; - out = __hmax2(__habs2(a), __habs2(b)); - } else if constexpr (std::is_same::value) { - float2 a = x1; - float2 b = x2; - out.x = fmaxf(fabsf(a.x), fabsf(b.x)); - out.y = fmaxf(fabsf(a.y), fabsf(b.y)); - } +inline void create_2D_tensor_map( + CUtensorMap* tensorMap, + void* input_ptr, + CUtensorMapDataType dtype, + uint64_t rows, + uint64_t cols, + uint32_t tile_y, + uint32_t tile_x, + uint64_t stride_bytes, + CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE) { + constexpr uint32_t rank = 2; // 2D + uint64_t global_dim[rank] = {cols, rows}; + // For row-major layout + uint64_t strides[rank - 1] = {stride_bytes}; + uint32_t tile_dim[rank] = {tile_x, tile_y}; + uint32_t elem_stride[rank] = {1, 1}; + + CHECK_CUDA_ERROR(cuTensorMapEncodeTiled( + tensorMap, + dtype, + rank, + input_ptr, + global_dim, + strides, + tile_dim, + elem_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); } -namespace cg = cooperative_groups; - -template -__global__ void fp_quantize_dequantize( - T* w, - T* out, - size_t size, - float* global_scale = nullptr) { - const bool use_global_scale = global_scale != nullptr; - const float scale_enc = - use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; - const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f; - - using Tx2 = Vector2_t; - uint32_t rbits = 0; // reserved bits for future use - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - size_t thread_idx = tidx + grid_dim_x * size_t(tidy); - size_t base_idx = thread_idx * group_size; - - if (base_idx >= size) { - return; - } - - auto w_tile = load_vector(w, thread_idx); - float scale_dec_b = 0.0f; - - Tx2 amax_2x = Tx2{0.0f, 0.0f}; - -#pragma unroll - for (int i = 0; i < group_size; i += 2) { - auto pair = Tx2{w_tile[i], w_tile[i + 1]}; - absmax_x2(amax_2x, amax_2x, pair); - } +inline std::tuple get_columnwise_quantize_mxfp8_launch_args( + size_t grid_dim_x_size, // rows + size_t grid_dim_y_size, // cols + size_t block_size_x, // ROWS_PER_BLOCK + size_t block_size_y, // COL_PER_BLOCK + int threads_per_block) { + dim3 grid; + grid.x = cuda::ceil_div(grid_dim_x_size, block_size_x); + grid.y = cuda::ceil_div(grid_dim_y_size, block_size_y); + grid.z = 1; - scale_dec_b = static_cast( - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y)))); - - scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; - scale_dec_b *= scale_enc; - // Convert to mx scale or nv scale - using ScaleType = std::conditional_t< - use_mx_scale, - cutlass::float_ue8m0_t, - cutlass::float_e4m3_t>; - auto s = ScaleType(scale_dec_b); - float scale_enc_b = scale_enc / float(s); - float scale_dec = float(s) * inv_scale_enc; - AlignedVector w_hat; - -#pragma unroll - for (int i = 0; i < group_size / 8; i++) { - auto& w = *reinterpret_cast*>(&w_tile[i * 8]); - cutlass::NumericArrayConverter fp32_t; - auto scaled = fp32_t(w) * scale_enc_b; - cutlass::Array dq; - if constexpr (bits == 8) { - cutlass::NumericArrayConverter fp8_fp32; - auto quant = fp8_fp32(scaled); - cutlass::NumericArrayConverter fp32_fp8; - dq = fp32_fp8(quant); - } else { - cutlass::NumericArrayConverter fp4_fp32; - auto quant = fp4_fp32(scaled); - cutlass::NumericArrayConverter fp32_fp4; - dq = fp32_fp4(quant); - } - cutlass::NumericArrayConverter t_fp32; - *reinterpret_cast*>(&w_hat[i * 8]) = - t_fp32(dq * scale_dec); - } - store_vector(out, thread_idx, w_hat); + dim3 block(threads_per_block, 1, 1); + return std::make_tuple(grid, block); } -template -__global__ void fp_quantize_rowwise( - T* w, - uint8_t* out, - uint8_t* scales, - size_t size, - float* global_scale = nullptr) { - // NVFP4 conversion: - // Global encode scale: (448 × 6) / *global_scale - // Per-block decode scale: S_dec_b = (block_amax / 6) × S_enc → stored as FP8 - // E4M3 Per-block encode scale: S_enc_b = S_enc / S_dec_b - const bool use_global_scale = global_scale != nullptr; - const float scale_enc = - use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; - - using Tx2 = Vector2_t; - using Tx4 = Vector4_t; - uint32_t rbits = 0; // reserved bits for future use - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - size_t thread_idx = tidx + grid_dim_x * size_t(tidy); - size_t base_idx = thread_idx * group_size; - - if (base_idx >= size) { - return; - } - - auto w_tile = load_vector(w, thread_idx); - float scale_dec_b = 0.0f; - - Tx2 amax_2x = Tx2{0.0f, 0.0f}; - -#pragma unroll - for (int i = 0; i < group_size; i += 2) { - auto pair = Tx2{w_tile[i], w_tile[i + 1]}; - absmax_x2(amax_2x, amax_2x, pair); - } - - scale_dec_b = static_cast( - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y)))); - - scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; - scale_dec_b *= scale_enc; - // Convert to mx scale or nv scale - using ScaleType = std::conditional_t< - use_mx_scale, - cutlass::float_ue8m0_t, - cutlass::float_e4m3_t>; - auto s = ScaleType(scale_dec_b); - uint8_t q_scale = s.storage; - float scale_enc_b = scale_enc / float(s); - - scales[thread_idx] = q_scale; - constexpr int elem_per_byte = bits == 8 ? 1 : 2; - AlignedVector quantized; - -#pragma unroll - for (int i = 0; i < group_size / 4; i++) { - Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); - if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); - *reinterpret_cast(&quantized[i * 4]) = quantized_val; - } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); - *reinterpret_cast(&quantized[i * 2]) = quantized_val; - } +inline CUtensorMapDataType get_tma_dtype(Dtype dtype) { + switch (dtype) { + case float16: + return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + case bfloat16: + return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + case float32: + return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + default: + throw std::runtime_error( + "[fp_quantize_columnwise_tma] Unsupported dtype for TMA"); } - store_vector(out, thread_idx, quantized); } -template -__global__ void fp_quantize_columnwise( - T* w, - uint8_t* out, - uint8_t* scales, +inline std::tuple get_columnwise_quantize_fallback_launch_args( size_t size, + int group_size, int M, - int K, - float* global_scale = nullptr) { - // Input: [M, K] with strides [1, M] (M-major) - // Quantized output: [M, K/elem_per_byte] row-major (K-major) - // Scales: [M, K/group_size] row-major (K-major) - // Quantize along K (last dimension, groups of group_size elements) - const bool use_global_scale = global_scale != nullptr; - const float scale_enc = - use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; - - using Tx2 = Vector2_t; - using Tx4 = Vector4_t; - uint32_t rbits = 0; - - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - - constexpr int BLOCK_X = 16; - constexpr int BLOCK_Y = 32; - constexpr int elem_per_byte = (bits == 8) ? 1 : 2; - constexpr int bytes_per_group = group_size / elem_per_byte; - - constexpr int rows_per_block = BLOCK_X; - constexpr int cols_per_block = BLOCK_Y * group_size; - constexpr int local_cols = cols_per_block / elem_per_byte; - constexpr int bytes_per_block = rows_per_block * local_cols; - - constexpr int SMEM_PAD = 4; - constexpr int padded_local_cols = local_cols + SMEM_PAD; - - auto tidx = idx_in_block.x; - auto tidy = idx_in_block.y; - - int num_col_blocks = (K + cols_per_block - 1) / cols_per_block; - auto bidx = block_idx.x % num_col_blocks; - auto bidy = block_idx.x / num_col_blocks; - - T thread_data[group_size]; - - __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols]; - __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD]; - - int row_base = bidy * rows_per_block + tidx; - int col_base = bidx * cols_per_block + tidy * group_size; - - bool valid = (row_base < M) && (col_base + group_size <= K); - if (valid) { -#pragma unroll - for (int i = 0; i < group_size; i++) { - auto index = row_base + (col_base + i) * M; - thread_data[i] = w[index]; - } - - // Compute scale - Tx2 amax_2x = Tx2{0.0f, 0.0f}; -#pragma unroll - for (int r = 0; r < group_size; r += 2) { - auto pair = Tx2{thread_data[r], thread_data[r + 1]}; - absmax_x2(amax_2x, amax_2x, pair); - } - float scale_dec_b = - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y))); - scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; - scale_dec_b *= scale_enc; - // Convert to mx scale or nv scale - using ScaleType = std::conditional_t< - use_mx_scale, - cutlass::float_ue8m0_t, - cutlass::float_e4m3_t>; - auto s = ScaleType(scale_dec_b); - float scale_enc_b = scale_enc / float(s); - scales_smem[tidx][tidy] = s.storage; - - int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group; - -#pragma unroll - for (int j = 0; j < group_size / 4; j++) { - Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); - if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); - *reinterpret_cast(&quantized_smem[shared_idx + j * 4]) = - quantized_val; - } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); - *reinterpret_cast(&quantized_smem[shared_idx + j * 2]) = - quantized_val; - } - } - } - __syncthreads(); - - int output_cols = K / elem_per_byte; - int num_groups_per_row = K / group_size; - int linear_tid = tidx + tidy * BLOCK_X; - // Write back quantized values -#pragma unroll - for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) { - int local_row = i / local_cols; - int local_col = i % local_cols; - - int global_row = bidy * rows_per_block + local_row; - int global_col = bidx * local_cols + local_col; - - if (global_row < M && global_col < output_cols) { - int physical_idx = local_row * padded_local_cols + local_col; - out[global_row * output_cols + global_col] = quantized_smem[physical_idx]; - } - } - // Write back scales - constexpr int num_scales = BLOCK_X * BLOCK_Y; -#pragma unroll - for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) { - int local_row = i / BLOCK_Y; - int local_col = i % BLOCK_Y; - - int global_row = bidy * BLOCK_X + local_row; - int global_col = bidx * BLOCK_Y + local_col; - - if (global_row < M && global_col < num_groups_per_row) { - scales[global_row * num_groups_per_row + global_col] = - scales_smem[local_row][local_col]; - } - } -} - -template -__global__ void fp_dequantize( - const uint8_t* w, - const uint8_t* scales, - T* out, - size_t size, - float* global_scale = nullptr) { - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - constexpr int pack_factor = bits == 8 ? 1 : 2; - const bool use_global_scale = global_scale != nullptr; - const float inv_scale_enc = use_mx_scale - ? 1.0f - : (use_global_scale ? (*global_scale) / (F8E4M3_MAX * F4E2M1_MAX) : 1.0f); - size_t offset = tidx + grid_dim_x * size_t(tidy); - size_t oindex = offset * pack_factor; - - if (oindex >= size) { - return; - } - - size_t gindex = oindex / group_size; - using ScaleType = std::conditional_t< - use_mx_scale, - cutlass::float_ue8m0_t, - cutlass::float_e4m3_t>; - auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc; - - out += oindex; - - uint32_t val = w[offset]; -#pragma clang loop unroll(full) - for (int i = 0; i < pack_factor; i++) { - uint8_t d; - if (bits == 4) { - d = (val >> (bits * i)) & 0x0f; - } else if (bits == 8) { - d = val; - } - out[i] = static_cast(scale * Dequantize{}(d)); - } -} - -inline std::tuple -get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) { + int K) { constexpr int BLOCK_X = 16; constexpr int BLOCK_Y = 32; int rows_per_block = BLOCK_X; @@ -449,7 +139,7 @@ void fp_quantize_dequantize( }); } -void fp_quantize( +void fp_quantize_rowwise( const array& w, array& wq, array& scales, @@ -464,66 +154,205 @@ void fp_quantize( } enc.set_output_array(wq); enc.set_output_array(scales); - if (w.strides().back() != 1) { - dispatch_float_types(w.dtype(), "fp_quantize_columnwise", [&](auto type_tag) { - using T = cuda_type_t; - if constexpr (!std::is_same_v) { - auto M = w.shape(-2); - auto K = w.shape(-1); - auto kernel = cu::fp_quantize_columnwise; - if (bits == 8) { - kernel = cu::fp_quantize_columnwise; - } else if (group_size == 16) { - kernel = cu::fp_quantize_columnwise; - } - auto [num_blocks, block_dims] = - cu::get_columnwise_quantize_launch_args(w.size(), group_size, M, K); - enc.add_kernel_node( - kernel, - num_blocks, - block_dims, - gpu_ptr(w), - gpu_ptr(wq), - gpu_ptr(scales), - w.size(), - M, - K, - global_scale.has_value() ? gpu_ptr(global_scale.value()) - : nullptr); - } else { - throw std::runtime_error( - "[Quantize::eval_gpu] Can not quantize input with type float64."); + dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + auto kernel = cu::fp_quantize_rowwise; + if (bits == 8) { + kernel = cu::fp_quantize_rowwise; + } else if (group_size == 16) { + kernel = cu::fp_quantize_rowwise; } - }); - } else { - dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) { - using T = cuda_type_t; - if constexpr (!std::is_same_v) { - auto kernel = cu::fp_quantize_rowwise; - if (bits == 8) { - kernel = cu::fp_quantize_rowwise; - } else if (group_size == 16) { - kernel = cu::fp_quantize_rowwise; + bool large = w.size() > UINT_MAX; + auto [num_blocks, block_dims] = + get_launch_args(w.size(), w.shape(), w.strides(), large, group_size); + + enc.add_kernel_node( + kernel, + num_blocks, + block_dims, + gpu_ptr(w), + gpu_ptr(wq), + gpu_ptr(scales), + w.size(), + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); + } else { + throw std::runtime_error( + "[Quantize::eval_gpu] Can not quantize input with type float64."); + } + }); +} + +void fp_quantize_columnwise_fallback( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + enc.set_input_array(w); + if (global_scale.has_value()) { + enc.set_input_array(global_scale.value()); + } + enc.set_output_array(wq); + enc.set_output_array(scales); + dispatch_float_types( + w.dtype(), "fp_quantize_columnwise_fallback", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + auto M = w.shape(-2); + auto K = w.shape(-1); + auto kernel = + cu::fp_quantize_columnwise_fallback; + if (bits == 8) { + kernel = cu::fp_quantize_columnwise_fallback; + } else if (group_size == 16) { + kernel = + cu::fp_quantize_columnwise_fallback; + } + auto [num_blocks, block_dims] = + cu::get_columnwise_quantize_fallback_launch_args( + w.size(), group_size, M, K); + enc.add_kernel_node( + kernel, + num_blocks, + block_dims, + gpu_ptr(w), + gpu_ptr(wq), + gpu_ptr(scales), + w.size(), + M, + K, + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); + } else { + throw std::runtime_error( + "[Quantize::eval_gpu] Can not quantize input with type float64."); } - bool large = w.size() > UINT_MAX; - auto [num_blocks, block_dims] = get_launch_args( - w.size(), w.shape(), w.strides(), large, group_size); - - enc.add_kernel_node( - kernel, - num_blocks, - block_dims, - gpu_ptr(w), - gpu_ptr(wq), - gpu_ptr(scales), - w.size(), - global_scale.has_value() ? gpu_ptr(global_scale.value()) - : nullptr); - } else { - throw std::runtime_error( - "[Quantize::eval_gpu] Can not quantize input with type float64."); - } - }); + }); +} + +void fp_quantize_columnwise_mxfp8( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + enc.set_input_array(w); + enc.set_output_array(wq); + enc.set_output_array(scales); + + size_t rows = w.shape(-1); + size_t cols = w.size() / rows; + size_t stride_bytes = w.strides(-1) * w.itemsize(); + + dispatch_float_types( + w.dtype(), "fp_quantize_columnwise_mxfp8", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + constexpr int THREADS_PER_BLOCK = 64; + constexpr int ROWS_PER_BLOCK = 64; + constexpr int COLS_PER_BLOCK = 64; + constexpr size_t TILE_M = 32; + constexpr size_t TILE_K = COLS_PER_BLOCK; + constexpr size_t STAGES = ROWS_PER_BLOCK / TILE_M; + + // For columnwise: grid.x = cols, grid.y = rows + // scales_per_stage = TILE_K (one scale per column per stage) + auto [grid, block] = cu::get_columnwise_quantize_mxfp8_launch_args( + cols, rows, COLS_PER_BLOCK, ROWS_PER_BLOCK, THREADS_PER_BLOCK); + + CUtensorMap tensor_map_input; + CUtensorMap tensor_map_output; + + cu::create_2D_tensor_map( + &tensor_map_input, + gpu_ptr(w), + cu::get_tma_dtype(w.dtype()), + rows, + cols, + TILE_M, + TILE_K, + stride_bytes); + + cu::create_2D_tensor_map( + &tensor_map_output, + gpu_ptr(wq), + CU_TENSOR_MAP_DATA_TYPE_UINT8, + cols, + rows, + TILE_K, + TILE_M, + rows); + + auto kernel = cu::fp_quantize_columnwise_mxfp8< + T, + false, + THREADS_PER_BLOCK, + COLS_PER_BLOCK, + ROWS_PER_BLOCK>; + enc.add_kernel_node( + kernel, + grid, + block, + tensor_map_input, + tensor_map_output, + gpu_ptr(scales), + rows, + cols); + } else { + throw std::runtime_error( + "[fp_quantize_columnwise_tma] Cannot quantize input with type float64."); + } + }); +} + +void fp_quantize_columnwise( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + // Use TMA version for SM100+ with MXFP8 (bits=8, group_size=32) + // NVFP4 todo + const size_t rows = w.shape(-1); + const size_t cols = w.size() / rows; + const bool has_full_tma_tiles = ((rows % 128) == 0) && ((cols % 128) == 0); + bool use_tma = + (enc.device().compute_capability_major() >= 10 && bits == 8 && + group_size == 32 && has_full_tma_tiles); + if (use_tma) { + fp_quantize_columnwise_mxfp8( + w, wq, scales, group_size, bits, global_scale, enc, s); + } else { + fp_quantize_columnwise_fallback( + w, wq, scales, group_size, bits, global_scale, enc, s); + } +} + +void fp_quantize( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + if (w.strides(-1) == 1) { + fp_quantize_rowwise(w, wq, scales, group_size, bits, global_scale, enc, s); + } else { + fp_quantize_columnwise( + w, wq, scales, group_size, bits, global_scale, enc, s); } } diff --git a/mlx/backend/cuda/quantized/fp_quantize.cuh b/mlx/backend/cuda/quantized/fp_quantize.cuh new file mode 100644 index 0000000000..769794129c --- /dev/null +++ b/mlx/backend/cuda/quantized/fp_quantize.cuh @@ -0,0 +1,601 @@ +// Copyright © 2025 Apple Inc. +#pragma once + +#include "mlx/backend/cuda/ptx.cuh" +#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh" +#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh" +#include "mlx/backend/cuda/vector_types.cuh" +#include "mlx/dtype_utils.h" + +#include +#include +#include +#include + +constexpr float F8E4M3_MAX = 448.0f; +constexpr float F4E2M1_MAX = 6.0f; +constexpr size_t TMA_SHMEM_ALIGNMENT = 128; +constexpr size_t BUFFS_NUM = 2; + +namespace mlx::core { +namespace cu { + +template +struct Dequantize { + __device__ float operator()(uint8_t x) { + if constexpr (bits == 8) { + return float(*(cutlass::float_e4m3_t*)(&x)); + } else { + return float(*(cutlass::float_e2m1_t*)(&x)); + } + } +}; + +template +__device__ __forceinline__ void absmax_x2(T& out, const T& x1, const T& x2) { + if constexpr ( + (std::is_same::value) || + (std::is_same::value)) { + T a = x1; + T b = x2; + out = __hmax2(__habs2(a), __habs2(b)); + } else if constexpr (std::is_same::value) { + float2 a = x1; + float2 b = x2; + out.x = fmaxf(fabsf(a.x), fabsf(b.x)); + out.y = fmaxf(fabsf(a.y), fabsf(b.y)); + } +} + +__device__ __forceinline__ void copy_2d_to_shared( + void* dst, + const CUtensorMap* tensor_map, + uint32_t tile_x, + uint32_t tile_y, + uint32_t num_bytes, + uint64_t* barrier, + const bool is_master_thread) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (is_master_thread) { + // Arrive and tell how many bytes are expected + ptx::mbarrier_arrive_expect_tx(barrier, num_bytes); + // Initiate bulk tensor copy + ptx::cp_async_bulk_tensor_2d_global_to_shared( + dst, tensor_map, tile_x, tile_y, barrier); + } else { + // Other threads just arrive + ptx::mbarrier_arrive(barrier); + } +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +namespace cg = cooperative_groups; + +template +__global__ void fp_quantize_dequantize( + T* w, + T* out, + size_t size, + float* global_scale = nullptr) { + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f; + + using Tx2 = Vector2_t; + uint32_t rbits = 0; // reserved bits for future use + + size_t thread_idx = cg::this_grid().thread_rank(); + size_t base_idx = thread_idx * group_size; + + if (base_idx >= size) { + return; + } + + auto w_tile = load_vector(w, thread_idx); + float scale_dec_b = 0.0f; + + Tx2 amax_2x = Tx2{0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < group_size; i += 2) { + auto pair = Tx2{w_tile[i], w_tile[i + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + scale_dec_b = static_cast( + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y)))); + + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + float scale_enc_b = scale_enc / float(s); + float scale_dec = float(s) * inv_scale_enc; + AlignedVector w_hat; + +#pragma unroll + for (int i = 0; i < group_size / 8; i++) { + auto& w = *reinterpret_cast*>(&w_tile[i * 8]); + cutlass::NumericArrayConverter fp32_t; + auto scaled = fp32_t(w) * scale_enc_b; + cutlass::Array dq; + if constexpr (bits == 8) { + cutlass::NumericArrayConverter fp8_fp32; + auto quant = fp8_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp8; + dq = fp32_fp8(quant); + } else { + cutlass::NumericArrayConverter fp4_fp32; + auto quant = fp4_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp4; + dq = fp32_fp4(quant); + } + cutlass::NumericArrayConverter t_fp32; + *reinterpret_cast*>(&w_hat[i * 8]) = + t_fp32(dq * scale_dec); + } + store_vector(out, thread_idx, w_hat); +} + +template +__global__ void fp_quantize_rowwise( + T* w, + uint8_t* out, + uint8_t* scales, + size_t size, + float* global_scale = nullptr) { + // NVFP4 conversion: + // Global encode scale: (448 × 6) / *global_scale + // Per-block decode scale: S_dec_b = (block_amax / 6) × S_enc → stored as FP8 + // E4M3 Per-block encode scale: S_enc_b = S_enc / S_dec_b + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + uint32_t rbits = 0; // reserved bits for future use + + size_t thread_idx = cg::this_grid().thread_rank(); + size_t base_idx = thread_idx * group_size; + + if (base_idx >= size) { + return; + } + + auto w_tile = load_vector(w, thread_idx); + float scale_dec_b = 0.0f; + + Tx2 amax_2x = Tx2{0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < group_size; i += 2) { + auto pair = Tx2{w_tile[i], w_tile[i + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + scale_dec_b = static_cast( + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y)))); + + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + uint8_t q_scale = s.storage; + float scale_enc_b = scale_enc / float(s); + + scales[thread_idx] = q_scale; + constexpr int elem_per_byte = bits == 8 ? 1 : 2; + AlignedVector quantized; + +#pragma unroll + for (int i = 0; i < group_size / 4; i++) { + Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); + if constexpr (bits == 8) { + uint32_t quantized_val = + scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized[i * 4]) = quantized_val; + } else { + uint16_t quantized_val = + scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized[i * 2]) = quantized_val; + } + } + store_vector(out, thread_idx, quantized); +} + +template < + typename T, + bool USE_SR, + int THREADS_PER_BLOCK, + int COLS_PER_BLOCK, + int ROWS_PER_BLOCK> +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + fp_quantize_columnwise_mxfp8( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_output, + uint8_t* __restrict__ scales, + const size_t rows, + const size_t cols) { +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + + constexpr size_t TILE_M = 32; + constexpr size_t TILE_K = COLS_PER_BLOCK; + constexpr size_t STEPS = ROWS_PER_BLOCK / TILE_M; + constexpr int elem_per_byte = 1; + + const auto block_idx = cg::this_thread_block().group_index(); + const auto idx_in_block = cg::this_thread_block().thread_index(); + const int tidx = idx_in_block.x; // Thread handles column tidx + const bool is_master = (tidx == 0); + + const size_t block_offset_col = block_idx.x * COLS_PER_BLOCK; + const size_t block_offset_row = block_idx.y * ROWS_PER_BLOCK; + + constexpr size_t BUFF_ELEMS = TILE_M * TILE_K; + constexpr uint32_t tile_bytes = static_cast(BUFF_ELEMS * sizeof(T)); + + // Note: 128 byte alignment is a requirement for shared memory buffers + // if only we use swizzling while tma copy + // if swizzling is not used, pointers must be 16 byte aligned + __shared__ alignas(16) T in_smem[BUFF_ELEMS * BUFFS_NUM]; + __shared__ alignas(16) uint8_t out_smem[BUFF_ELEMS * BUFFS_NUM]; + + __shared__ alignas(8) uint64_t mbar[STEPS]; + + T thread_data[TILE_M]; + uint32_t rbits = 0; // Reserved for stochastic rounding + const size_t scale_stride = rows / TILE_M; + + // Master thread init memory barriers for all steps + // fence for tma, synchronize threads so all see mbarrier + if (is_master) { +#pragma unroll + for (int iter = 0; iter < STEPS; ++iter) { + ptx::mbarrier_init(&mbar[iter], THREADS_PER_BLOCK); + } + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + // Launch first async copy before entering the loop + copy_2d_to_shared( + &in_smem[0], + &tensor_map_input, + static_cast(block_offset_col), + static_cast(block_offset_row), + tile_bytes, + &mbar[0], + is_master); + +#pragma unroll + for (size_t step = 0; step < STEPS; ++step) { + // buffer memory offset in shared memory (we use double buffering for + // pipelining) + const size_t buff = step % BUFFS_NUM; + const size_t next_step = step + 1; + const size_t step_row_offset = step * TILE_M; + + if (next_step < STEPS) { + // before launching another async copy, check that there is less than 2 + // (to ensure that shared -> global synch is finished and buffer can be + // reused) + ptx::cp_async_bulk_wait_group_read<1>(); + const size_t next_buff = next_step % BUFFS_NUM; + const size_t next_row_offset = block_offset_row + next_step * TILE_M; + const size_t next_buff_elem_offset = next_buff * BUFF_ELEMS; + + copy_2d_to_shared( + &in_smem[next_buff_elem_offset], + &tensor_map_input, + static_cast(block_offset_col), + static_cast(next_row_offset), + tile_bytes, + &mbar[next_step], + is_master); + } + + ptx::fence_proxy_async_shared_cta(); + // Wait until the data is ready, parity is always 0 because for simplicity + // we dont reuse barriers between steps + ptx::mbarrier_wait_parity(&mbar[step], 0); + const size_t buff_offset = buff * BUFF_ELEMS; + // Read the data from shared to registers +#pragma unroll + for (int row = 0; row < TILE_M; ++row) { + thread_data[row] = in_smem[buff_offset + row * TILE_K + tidx]; + } + Tx2 amax_2x = Tx2{T(0.0f), T(0.0f)}; +#pragma unroll + for (int row = 0; row < TILE_M; row += 2) { + auto pair = Tx2{thread_data[row], thread_data[row + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + float scale = + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y))); + + scale /= F8E4M3_MAX; + + auto s = cutlass::float_ue8m0_t(scale); + scale = float(s); + // Write scale directly to global memory + const size_t global_col = block_offset_col + tidx; + const size_t global_row_group = + (block_offset_row + step_row_offset) / TILE_M; + if (global_col < cols && (block_offset_row + step_row_offset) < rows) { + scales[global_col * scale_stride + global_row_group] = s.storage; + } + const size_t out_buff_offset = buff * BUFF_ELEMS / elem_per_byte; + // Quantize to registers first + constexpr int GROUPS = TILE_M / 4; + uint32_t quantized_regs[GROUPS]; +#pragma unroll + for (int j = 0; j < GROUPS; ++j) { + Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); + quantized_regs[j] = + cu::scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); + } + // Write output to shared memory with swapped store order to reduce bank + // conflicts. Without swap: stride between threads is TILE_M=32 bytes + // = 8 banks, every 4th thread hits the same bank -> 8-way conflict. + const int lane = tidx % 32; + const int group = (lane / 4) % 2; + const size_t base = out_buff_offset + tidx * TILE_M; + switch (group) { + case 0: + *reinterpret_cast(&out_smem[base + 0]) = { + quantized_regs[0], + quantized_regs[1], + quantized_regs[2], + quantized_regs[3]}; + *reinterpret_cast(&out_smem[base + 16]) = { + quantized_regs[4], + quantized_regs[5], + quantized_regs[6], + quantized_regs[7]}; + break; + case 1: + *reinterpret_cast(&out_smem[base + 16]) = { + quantized_regs[4], + quantized_regs[5], + quantized_regs[6], + quantized_regs[7]}; + *reinterpret_cast(&out_smem[base + 0]) = { + quantized_regs[0], + quantized_regs[1], + quantized_regs[2], + quantized_regs[3]}; + break; + } + __syncthreads(); + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + if (is_master) { + const size_t global_row = block_offset_row + step_row_offset; + const uint32_t out_x = static_cast(global_row); + const uint32_t out_y = static_cast(block_offset_col); + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), + out_x, + out_y, + reinterpret_cast(&out_smem[out_buff_offset])); + ptx::cp_async_bulk_commit_group(); + } + } + // Wait for all TMA stores to complete + ptx::cp_async_bulk_wait_group_read<0>(); + + __syncthreads(); + if (is_master) { +#pragma unroll + for (int iter = 0; iter < STEPS; ++iter) { + ptx::mbarrier_invalidate(&mbar[iter]); + } + } +#endif // __CUDA_ARCH__ >= 1000 +} + +// TODO: add kernel with tma instructions +template +__global__ void fp_quantize_columnwise_fallback( + T* w, + uint8_t* out, + uint8_t* scales, + size_t size, + int M, + int K, + float* global_scale = nullptr) { + // Input: [M, K] with strides [1, M] (M-major) + // Quantized output: [M, K/elem_per_byte] row-major (K-major) + // Scales: [M, K/group_size] row-major (K-major) + // Quantize along K (last dimension, groups of group_size elements) + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + uint32_t rbits = 0; + + auto block_idx = cg::this_thread_block().group_index(); + auto idx_in_block = cg::this_thread_block().thread_index(); + + constexpr int BLOCK_X = 16; + constexpr int BLOCK_Y = 32; + constexpr int elem_per_byte = (bits == 8) ? 1 : 2; + constexpr int bytes_per_group = group_size / elem_per_byte; + + constexpr int rows_per_block = BLOCK_X; + constexpr int cols_per_block = BLOCK_Y * group_size; + constexpr int local_cols = cols_per_block / elem_per_byte; + constexpr int bytes_per_block = rows_per_block * local_cols; + + constexpr int SMEM_PAD = 4; + constexpr int padded_local_cols = local_cols + SMEM_PAD; + + auto tidx = idx_in_block.x; + auto tidy = idx_in_block.y; + + int num_col_blocks = (K + cols_per_block - 1) / cols_per_block; + auto bidx = block_idx.x % num_col_blocks; + auto bidy = block_idx.x / num_col_blocks; + + T thread_data[group_size]; + + __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols]; + __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD]; + + int row_base = bidy * rows_per_block + tidx; + int col_base = bidx * cols_per_block + tidy * group_size; + + bool valid = (row_base < M) && (col_base + group_size <= K); + if (valid) { +#pragma unroll + for (int i = 0; i < group_size; i++) { + auto index = row_base + (col_base + i) * M; + thread_data[i] = w[index]; + } + + // Compute scale + Tx2 amax_2x = Tx2{0.0f, 0.0f}; +#pragma unroll + for (int r = 0; r < group_size; r += 2) { + auto pair = Tx2{thread_data[r], thread_data[r + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + float scale_dec_b = + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y))); + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + float scale_enc_b = scale_enc / float(s); + scales_smem[tidx][tidy] = s.storage; + + int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group; + +#pragma unroll + for (int j = 0; j < group_size / 4; j++) { + Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); + if constexpr (bits == 8) { + uint32_t quantized_val = + scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized_smem[shared_idx + j * 4]) = + quantized_val; + } else { + uint16_t quantized_val = + scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized_smem[shared_idx + j * 2]) = + quantized_val; + } + } + } + __syncthreads(); + + int output_cols = K / elem_per_byte; + int num_groups_per_row = K / group_size; + int linear_tid = tidx + tidy * BLOCK_X; + // Write back quantized values +#pragma unroll + for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) { + int local_row = i / local_cols; + int local_col = i % local_cols; + + int global_row = bidy * rows_per_block + local_row; + int global_col = bidx * local_cols + local_col; + + if (global_row < M && global_col < output_cols) { + int physical_idx = local_row * padded_local_cols + local_col; + out[global_row * output_cols + global_col] = quantized_smem[physical_idx]; + } + } + // Write back scales + constexpr int num_scales = BLOCK_X * BLOCK_Y; +#pragma unroll + for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) { + int local_row = i / BLOCK_Y; + int local_col = i % BLOCK_Y; + + int global_row = bidy * BLOCK_X + local_row; + int global_col = bidx * BLOCK_Y + local_col; + + if (global_row < M && global_col < num_groups_per_row) { + scales[global_row * num_groups_per_row + global_col] = + scales_smem[local_row][local_col]; + } + } +} + +template +__global__ void fp_dequantize( + const uint8_t* w, + const uint8_t* scales, + T* out, + size_t size, + float* global_scale = nullptr) { + auto block_size = cg::this_thread_block().dim_threads(); + auto block_idx = cg::this_thread_block().group_index(); + auto idx_in_block = cg::this_thread_block().thread_index(); + + auto tidx = block_idx.x * block_size.x + idx_in_block.x; + auto tidy = block_idx.y * block_size.y + idx_in_block.y; + + auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; + + constexpr int pack_factor = bits == 8 ? 1 : 2; + const bool use_global_scale = global_scale != nullptr; + const float inv_scale_enc = use_mx_scale + ? 1.0f + : (use_global_scale ? (*global_scale) / (F8E4M3_MAX * F4E2M1_MAX) : 1.0f); + size_t offset = tidx + grid_dim_x * size_t(tidy); + size_t oindex = offset * pack_factor; + + if (oindex >= size) { + return; + } + + size_t gindex = oindex / group_size; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc; + + out += oindex; + + uint32_t val = w[offset]; +#pragma clang loop unroll(full) + for (int i = 0; i < pack_factor; i++) { + uint8_t d; + if (bits == 4) { + d = (val >> (bits * i)) & 0x0f; + } else if (bits == 8) { + d = val; + } + out[i] = static_cast(scale * Dequantize{}(d)); + } +} + +} // namespace cu +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 96a1fff7bb..81d256be23 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -8,13 +8,6 @@ namespace mlx::core { -namespace cg = cooperative_groups; - -constexpr int TILE_ROWS = 128; -constexpr int TILE_COLS = 4; -constexpr int TILES_PER_LANE = 1; -constexpr int LANES_PER_BLOCK = 32; - // To pass scales to tensor cores, they need to be repacked into a tiled layout // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout // Tiled layout for scale factors is very well described in CUTLASS @@ -48,6 +41,14 @@ constexpr int LANES_PER_BLOCK = 32; // [252, 253, 254, 255], // [380, 381, 382, 383], // [508, 509, 510, 511]]]]], +namespace cu { + +constexpr int TILE_ROWS = 128; +constexpr int TILE_COLS = 4; +constexpr int TILES_PER_LANE = 1; +constexpr int LANES_PER_BLOCK = 32; + +namespace cg = cooperative_groups; inline std::tuple get_swizzle_launch_args( size_t M_swizzled, @@ -68,8 +69,6 @@ inline std::tuple get_swizzle_launch_args( return std::make_tuple(grid, block); } -namespace cu { - __global__ void compute_qqmm_pointers( float* alpha_out, float* beta_out, @@ -222,7 +221,7 @@ void swizzle_scales( size_t output_cols = scales_tiled.shape(-1); auto [num_blocks, block_dims] = - get_swizzle_launch_args(output_rows, output_cols); + cu::get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, From b400c6ced2209d133582066f1014cc72fd92f71a Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Thu, 30 Jul 2026 07:12:07 -0700 Subject: [PATCH 019/222] metal: reduce NVFP4 scales per 16-lane group (#3934) Co-authored-by: Anastasiia Filippova --- mlx/backend/metal/kernels/fp_quantized.h | 18 ++++++++++-------- python/tests/test_quantized.py | 9 +++++++++ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 5dd81b4f1e..677183fa92 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -2002,7 +2002,8 @@ template device uint8_t* out [[buffer(1)]], device uint8_t* scales [[buffer(2)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); @@ -2011,9 +2012,9 @@ template if (use_mx_scale) { scale = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale = simd_lid < 16 ? w_max_l : w_max_r; } scale /= bits == 4 ? 6.0f : 448.0f; @@ -2075,7 +2076,8 @@ template const device T* w [[buffer(0)]], device T* out [[buffer(1)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); @@ -2084,9 +2086,9 @@ template if (use_mx_scale) { scale = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale = simd_lid < 16 ? w_max_l : w_max_r; } scale /= bits == 4 ? 6.0f : 448.0f; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index f031d2aa92..b993d57bb7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -154,6 +154,15 @@ def test_nvfp4_quantize_dequantize(self): w_hat = mx.dequantize(w_q, scales, mode="nvfp4") self.assertTrue(mx.allclose(w, w_hat, rtol=1e-5, atol=1e-5)) + # A scale shared across a 32-value SIMD group instead of computed + # per 16-value group cannot represent the low-magnitude groups. + alternating = mx.zeros((64, 16), dtype=mx.bfloat16) + alternating[::2] = 6 * 2**-9 + alternating[1::2] = 6.0 + w_q, scales = mx.quantize(alternating, mode="nvfp4") + w_hat = mx.dequantize(w_q, scales, mode="nvfp4", dtype=mx.bfloat16) + self.assertTrue(mx.allclose(alternating, w_hat, rtol=1e-5, atol=1e-6)) + # test quantize/dequantize 0s a = mx.zeros((256, 512)) w_q, scales = mx.quantize(a, mode="nvfp4") From bfd6d0d47f562f0ad0436de7448355f99fcbf0c7 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 30 Jul 2026 08:41:54 -0700 Subject: [PATCH 020/222] Update homebrew in CI (#3946) Co-authored-by: Anastasiia Filippova --- .github/actions/setup/action.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 4f7b14105f..898bd0f9a4 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -52,6 +52,7 @@ runs: shell: bash run: | echo "::group::Install macOS dependencies" + brew update brew install openmpi xcodebuild -showComponent MetalToolchain echo "::endgroup::" From 8d9d90655acf201a35c1370e510eda0c67ee71c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Thu, 30 Jul 2026 18:42:31 +0300 Subject: [PATCH 021/222] Make index autodiff errors explicitly recommend stop_gradient (#3820) --- mlx/primitives.cpp | 39 +++++++++++------- python/tests/test_array.py | 2 +- python/tests/test_autograd.py | 75 +++++++++++++++++++++++++++++++++++ python/tests/test_blas.py | 18 +++++++++ 4 files changed, 118 insertions(+), 16 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 6ac88eef6c..ae434af94d 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -2481,9 +2481,9 @@ std::vector Gather::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); std::vector inds(primals.begin() + 1, primals.end()); @@ -2499,7 +2499,8 @@ std::vector Gather::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather] Cannot calculate JVP with respect to indices."); + "[gather] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } std::vector inds(primals.begin() + 1, primals.end()); return {gather(tangents[0], inds, axes_, slice_sizes_, stream())}; @@ -2546,9 +2547,9 @@ std::vector GatherAxis::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather_axis] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); vjps.push_back(array( @@ -2567,7 +2568,8 @@ std::vector GatherAxis::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather_axis] Cannot calculate JVP with respect to indices."); + "[gather_axis] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } return {take_along_axis(tangents[0], primals[1], axis_, stream())}; } @@ -4483,7 +4485,8 @@ std::vector Scatter::vjp( } } else { throw std::invalid_argument( - "[scatter] Cannot calculate VJP with respect to indices."); + "[scatter] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4595,7 +4598,8 @@ std::vector ScatterAxis::vjp( vjps.push_back(take_along_axis(cotangents[0], indices, axis_, stream())); } else { throw std::invalid_argument( - "[scatter_axis] Cannot calculate VJP with respect to indices."); + "[scatter_axis] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4608,7 +4612,8 @@ std::vector ScatterAxis::jvp( for (auto arg : argnums) { if (arg == 1) { throw std::invalid_argument( - "[scatter_axis] Cannot calculate JVP with respect to indices."); + "[scatter_axis] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } if (argnums.size() == 2) { @@ -4716,7 +4721,8 @@ std::vector MaskedScatter::vjp( vjps.push_back(reshape(gsrc_flat, src.shape(), s)); } else { throw std::invalid_argument( - "[masked_scatter] Cannot calculate VJP with respect to mask."); + "[masked_scatter] Cannot calculate VJP with respect to mask. " + "Use stop_gradient on mask to stop gradients from being computed."); } } return vjps; @@ -5748,7 +5754,8 @@ std::vector BlockMaskedMM::vjp( if ((needs_lhs_mask_vjp && primals[op_mask_idx].dtype() == bool_) || (needs_rhs_mask_vjp && primals[op_mask_idx + 1].dtype() == bool_)) { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } auto expand_mask = [&](array mask, int Y, int X) { @@ -5945,7 +5952,8 @@ std::vector BlockMaskedMM::vjp( } else { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } } return vjps; @@ -6010,7 +6018,8 @@ std::vector GatherMM::vjp( stream())); } else { throw std::invalid_argument( - "[GatherMM] Cannot calculate VJP with respect to indices."); + "[GatherMM] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 04e17f8710..7279fa6276 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1205,7 +1205,7 @@ def test_indexing_grad(self): ind = mx.array([0, 1, 0]).astype(mx.float32) def index_fn(x, ind): - return x[ind.astype(mx.int32)].sum() + return x[mx.stop_gradient(ind.astype(mx.int32))].sum() grad_x, grad_ind = mx.grad(index_fn, argnums=(0, 1))(x, ind) expected = mx.array([[2, 2], [1, 1]]) diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index fa6aed6f09..2f72fba113 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -444,6 +444,81 @@ def fun(x, idx): self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 1.0]))) self.assertEqual(dfdx.dtype, mx.float32) + def test_index_vjp_requires_stop_gradient(self): + msg = "stop_gradient" + x = mx.array([1.0, 2.0, 3.0, 4.0]) + idx = mx.array([1, 3]) + updates = mx.array([5.0, 6.0]) + x_axis = x[:, None] + idx_axis = idx[:, None] + updates_axis = updates[:, None] + + def gather_fun(x, idx): + return mx.take(x, idx) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_fun, [x, idx], [mx.ones((2,))]) + + def gather_axis_fun(x, idx): + return mx.take_along_axis(x, idx, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_axis_fun, [x_axis, idx_axis], [mx.ones((2, 1))]) + + def scatter_fun(x, idx, updates): + return x.at[idx].add(updates) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(scatter_fun, [x, idx, updates], [mx.ones((4,))]) + + def scatter_axis_fun(x, idx, updates): + return mx.put_along_axis(x, idx, updates, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp( + scatter_axis_fun, + [x_axis, idx_axis, updates_axis], + [mx.ones((4, 1))], + ) + + def test_stop_gradient_computed_indices(self): + def gather_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + return mx.take(w, idx).sum() + + grad = mx.grad(gather_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([0.0, 0.0, 1.0, 1.0]))) + + def gather_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + return mx.take_along_axis(w, idx, axis=1).sum() + + grad = mx.grad(gather_axis_fun)(mx.array([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0]])) + self.assertTrue( + mx.array_equal( + grad, + mx.array([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]), + ) + ) + + def scatter_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + out = mx.zeros((4,)) + out[idx] = w[:2] + return out.sum() + + grad = mx.grad(scatter_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([1.0, 1.0, 0.0, 0.0]))) + + def scatter_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + updates = w.sum(axis=1, keepdims=True) + out = mx.put_along_axis(mx.zeros((3, 3)), idx, updates, axis=1) + return out.sum() + + grad = mx.grad(scatter_axis_fun)(mx.ones((3, 3))) + self.assertTrue(mx.array_equal(grad, mx.ones((3, 3)))) + def test_scatter_add_vjp(self): def fun(src, updates): x = src.at[mx.array([1, 3])].add(updates) diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 3a8183ba8f..957c46aec9 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -1443,6 +1443,24 @@ def f_test(a, b): self.assertEqual(r.shape, t.shape) self.assertTrue(mx.allclose(r, t, atol=1e-4).item()) + def test_gather_matmul_index_vjp_requires_stop_gradient(self): + a = mx.ones((4, 1, 2, 2)) + b = mx.ones((4, 1, 2, 2)) + + def fun(w): + indices = mx.reshape(mx.argsort(w)[:2], (1, 2)) + return mx.gather_mm(a, b, indices, indices).sum() + + with self.assertRaisesRegex(ValueError, "stop_gradient"): + mx.grad(fun)(mx.array([3.0, 1.0, 2.0, 0.0])) + + def fun_stopped(w): + indices = mx.stop_gradient(mx.reshape(mx.argsort(w)[:2], (1, 2))) + return mx.gather_mm(a, b, indices, indices).sum() + + grad = mx.grad(fun_stopped)(mx.array([3.0, 1.0, 2.0, 0.0])) + self.assertTrue(mx.array_equal(grad, mx.zeros((4,)))) + def test_gather_mm_sorted(self): def gather_mm_ref(a, b, rhs): b = b[rhs] From 85e3fb3bb1a5e57caff90a522014ce455563ef22 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 30 Jul 2026 14:53:04 -0700 Subject: [PATCH 022/222] Making JACCL coordinator optional (#3899) --- docs/src/usage/distributed.rst | 27 +++++++ mlx/distributed/distributed.cpp | 53 +++++++++++-- mlx/distributed/distributed.h | 34 +++++++- mlx/distributed/jaccl/jaccl.cpp | 11 +++ mlx/distributed/jaccl/jaccl.h | 3 + mlx/distributed/jaccl/lib/jaccl/jaccl.cpp | 97 ++++++++++++++++++++--- mlx/distributed/jaccl/lib/jaccl/jaccl.h | 25 +++++- mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 4 +- mlx/distributed/jaccl/lib/jaccl/mesh.h | 2 +- mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 25 +++++- mlx/distributed/jaccl/lib/jaccl/rdma.h | 93 +++++++++++++--------- mlx/distributed/jaccl/lib/jaccl/ring.cpp | 4 +- mlx/distributed/jaccl/lib/jaccl/ring.h | 2 +- mlx/distributed/jaccl/no_jaccl.cpp | 7 ++ python/src/distributed.cpp | 64 ++++++++++++++- python/tests/mlx_distributed_tests.py | 12 +++ 16 files changed, 392 insertions(+), 71 deletions(-) diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 1f271835cc..4586d7e8ae 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -347,6 +347,33 @@ of a gigantic model using MLX LM. and GPU need to collaborate for some computation and is pretty critical for low-latency communication since the communication is done by the CPU. +Custom side channel +^^^^^^^^^^^^^^^^^^^ + +During initialization JACCL exchanges RDMA connection metadata between ranks +over a side channel. By default this is a simple TCP star all-gather, but you +can provide your own side-channel all-gather via the +``all_gather_factory`` argument of :func:`mlx.core.distributed.init` when using +``backend='jaccl'``. + +The factory is called once per rank with ``(rank, size)`` and must return a +callable with signature ``f(src: bytes, n_bytes: int) -> bytes``. The returned +bytes must have length ``size * n_bytes`` and contain the inputs from all ranks +concatenated in rank order. + +.. code-block:: python + + def make_side_channel(rank, size): + def all_gather(src: bytes, n_bytes: int) -> bytes: + # Exchange src with all ranks and return size * n_bytes bytes + ... + return all_gather + + world = mx.distributed.init( + backend="jaccl", + all_gather_factory=make_side_channel, + ) + .. _nccl_section: Getting Started with NCCL diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index 3cde6a263b..34be2ec600 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -138,9 +138,30 @@ Group Group::split(int color, int key /* = -1 */) const { return Group(group_->split(color, key)); } -Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { +namespace { + +std::unordered_map>& +get_backends() { static std::unordered_map> backends; + return backends; +} + +Group register_group(std::shared_ptr group, std::string bk) { + auto& backends = get_backends(); + if (group == nullptr) { + group = std::make_shared(); + } else { + backends.insert({"any", group}); + } + backends.insert({std::move(bk), group}); + return Group(group); +} + +} // namespace + +Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { + auto& backends = get_backends(); // Already initialized so return the group. if (auto g = backends.find(bk); g != backends.end()) { @@ -185,13 +206,31 @@ Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { throw std::invalid_argument(msg.str()); } - if (group == nullptr) { - group = std::make_shared(); - } else { - backends.insert({"any", group}); + return register_group(std::move(group), std::move(bk_)); +} + +Group init(bool strict, const std::string& bk, AllGatherFactory factory) { + if (bk != "jaccl") { + std::ostringstream msg; + msg << "[distributed] An all gather factory is only supported with the " + << "'jaccl' backend but '" << bk << "' was provided."; + throw std::invalid_argument(msg.str()); } - backends.insert({std::move(bk_), group}); - return Group(group); + + auto& backends = get_backends(); + + // Already initialized so return the cached group. The factory is ignored in + // this case since the backend is already up. + if (auto g = backends.find(bk); g != backends.end()) { + return Group(g->second); + } + + auto group = jaccl::init(strict, std::move(factory)); + return register_group(std::move(group), bk); +} + +void clear_backends() { + get_backends().clear(); } } // namespace mlx::core::distributed diff --git a/mlx/distributed/distributed.h b/mlx/distributed/distributed.h index 00c7a80e86..0d6e4d41ef 100644 --- a/mlx/distributed/distributed.h +++ b/mlx/distributed/distributed.h @@ -2,6 +2,8 @@ #pragma once +#include +#include #include #include "mlx/api.h" @@ -13,7 +15,26 @@ namespace mlx::core::distributed { // Forward declaration of the base group implementation. namespace detail { class GroupImpl; -}; +} + +/** + * Byte-level all-gather function used by the JACCL side channel. + * + * Args: + * src: Pointer to this rank's data of size n_bytes. + * dst: Pointer to an output buffer of size size_ * n_bytes. After the call, + * dst[rank * n_bytes, (rank+1) * n_bytes] contains the data from rank. + * n_bytes: The number of bytes contributed by each rank. + */ +using AllGatherFn = std::function; + +/** + * Factory that produces a per-rank side-channel all-gather function. + * + * The factory receives the rank and size of the group and returns the + * AllGatherFn that will be used for the side channel. + */ +using AllGatherFactory = std::function; /* Check if a communication backend is available */ MLX_API bool is_available(); @@ -58,4 +79,15 @@ struct MLX_API Group { */ MLX_API Group init(bool strict = false, const std::string& bk = "any"); +/** + * Initialize the distributed backend using a custom JACCL side channel. + * + * This behaves like init(strict, bk) but requires bk to be "jaccl". + */ +MLX_API Group +init(bool strict, const std::string& bk, AllGatherFactory factory); + +/** Clear the distributed backend cache. */ +MLX_API void clear_backends(); + } // namespace mlx::core::distributed diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 01cc415a00..3d4b8e070e 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -8,6 +8,9 @@ #include #include +#include +#include + using GroupImpl = mlx::core::distributed::detail::GroupImpl; namespace mlx::core::distributed::jaccl { @@ -170,4 +173,12 @@ std::shared_ptr init(bool strict /* = false */) { return std::make_shared(std::move(group)); } +std::shared_ptr init(bool strict, AllGatherFactory factory) { + auto group = ::jaccl::init(strict, factory); + if (group == nullptr) { + return nullptr; + } + return std::make_shared(std::move(group)); +} + } // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/jaccl.h b/mlx/distributed/jaccl/jaccl.h index cc2326b61d..0e2210ac77 100644 --- a/mlx/distributed/jaccl/jaccl.h +++ b/mlx/distributed/jaccl/jaccl.h @@ -2,6 +2,8 @@ #pragma once +#include + #include "mlx/distributed/distributed.h" namespace mlx::core::distributed::jaccl { @@ -10,5 +12,6 @@ using GroupImpl = mlx::core::distributed::detail::GroupImpl; bool is_available(); std::shared_ptr init(bool strict = false); +std::shared_ptr init(bool strict, AllGatherFactory factory); } // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp index 42d665ed10..7715022fcc 100644 --- a/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp @@ -1,13 +1,13 @@ // Copyright © 2025 Apple Inc. #include +#include #include #include #include "jaccl/jaccl.h" #include "jaccl/mesh.h" -#include "jaccl/rdma.h" #include "jaccl/ring.h" using json = nlohmann::json; @@ -75,16 +75,37 @@ namespace jaccl { Config::Config() : rank_(0), size_(0) {} +Config& Config::set_rank(const char* rank_str) { + if (rank_str) { + rank_ = std::atoi(rank_str); + } + return *this; +} + Config& Config::set_rank(int rank) { rank_ = rank; return *this; } +Config& Config::set_coordinator(const char* coordinator) { + if (coordinator != nullptr) { + coordinator_ = coordinator; + } + return *this; +} + Config& Config::set_coordinator(std::string coordinator) { coordinator_ = std::move(coordinator); return *this; } +Config& Config::set_devices_from_file(const char* dev_file) { + if (dev_file) { + set_devices(parse_devices_json(dev_file)); + } + return *this; +} + Config& Config::set_devices( std::vector>> devices) { devices_ = std::move(devices); @@ -106,6 +127,17 @@ Config& Config::prefer_ring(bool prefer /* = true */) { return *this; } +Config& Config::set_all_gather(AllGatherFn agf) { + all_gather_fn_ = std::move(agf); + return *this; +} + +Config& Config::set_all_gather_factory( + std::function factory) { + all_gather_factory_ = std::move(factory); + return *this; +} + bool Config::is_valid_mesh() const { if (size_ < 2) { return false; @@ -142,6 +174,12 @@ bool Config::is_valid_ring() const { return true; } +bool Config::is_valid() const { + return size_ >= 1 && rank_ < size_ && rank_ >= 0 && + (!coordinator_.empty() || all_gather_factory_ || all_gather_fn_) && + (is_valid_mesh() || is_valid_ring()); +} + std::vector Config::get_mesh_connectivity() const { if (!is_valid_mesh()) { throw std::runtime_error("[jaccl] The devices do not form a valid mesh."); @@ -166,21 +204,36 @@ Config::get_ring_connectivity() const { return std::make_pair(devices_[rank_][left], devices_[rank_][right]); } -std::optional Config::from_env() { +SideChannel Config::get_side_channel() const { + if (all_gather_factory_) { + return SideChannel(rank_, size_, all_gather_factory_(rank_, size_)); + } + + if (all_gather_fn_) { + return SideChannel(rank_, size_, all_gather_fn_); + } + + auto tcp = + std::make_shared(rank_, size_, get_coordinator().c_str()); + return SideChannel( + rank_, + size_, + [tcp = std::move(tcp)](const char* src, char* dst, size_t n_bytes) { + (*tcp)(src, dst, n_bytes); + }); +} + +Config Config::from_env() { const char* dev_file = getenv("JACCL_IBV_DEVICES", "MLX_IBV_DEVICES"); const char* coordinator = getenv("JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR"); const char* rank_str = getenv("JACCL_RANK", "MLX_RANK"); const char* ring = getenv("JACCL_RING", "MLX_JACCL_RING"); - if (!dev_file || !coordinator || !rank_str) { - return std::nullopt; - } - return Config() - .set_rank(std::atoi(rank_str)) + .set_rank(rank_str) .set_coordinator(coordinator) - .set_devices(parse_devices_json(dev_file)) + .set_devices_from_file(dev_file) .prefer_ring(ring != nullptr); } @@ -190,7 +243,7 @@ bool is_available() { std::shared_ptr init(bool strict /* = false */) { auto cfg = Config::from_env(); - if (!cfg.has_value()) { + if (!cfg.is_valid()) { if (strict) { std::ostringstream msg; msg << "[jaccl] You need to provide via environment variables a rank " @@ -202,22 +255,40 @@ std::shared_ptr init(bool strict /* = false */) { return nullptr; } - return init(*cfg, strict); + return init(cfg, strict); +} + +std::shared_ptr init( + bool strict, + std::function factory) { + auto cfg = Config::from_env().set_all_gather_factory(factory); + if (!cfg.is_valid()) { + if (strict) { + std::ostringstream msg; + msg << "[jaccl] You need to provide via environment variables a rank " + << "(JACCL_RANK/MLX_RANK) and a device file " + << "(JACCL_IBV_DEVICES/MLX_IBV_DEVICES)"; + throw std::runtime_error(msg.str()); + } + return nullptr; + } + + return init(cfg, strict); } std::shared_ptr init(const Config& cfg, bool strict /* = false */) { if (cfg.get_prefer_ring() && cfg.is_valid_ring()) { auto [left, right] = cfg.get_ring_connectivity(); return std::make_shared( - cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator()); + cfg.get_rank(), cfg.get_size(), left, right, cfg.get_side_channel()); } else if (cfg.is_valid_mesh()) { auto mesh = cfg.get_mesh_connectivity(); return std::make_shared( - cfg.get_rank(), mesh, cfg.get_coordinator()); + cfg.get_rank(), mesh, cfg.get_side_channel()); } else if (cfg.is_valid_ring()) { auto [left, right] = cfg.get_ring_connectivity(); return std::make_shared( - cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator()); + cfg.get_rank(), cfg.get_size(), left, right, cfg.get_side_channel()); } else { if (!strict) { return nullptr; diff --git a/mlx/distributed/jaccl/lib/jaccl/jaccl.h b/mlx/distributed/jaccl/lib/jaccl/jaccl.h index a4df21ae0e..39feb34b47 100644 --- a/mlx/distributed/jaccl/lib/jaccl/jaccl.h +++ b/mlx/distributed/jaccl/lib/jaccl/jaccl.h @@ -2,10 +2,12 @@ #pragma once +#include #include #include #include "jaccl/group.h" +#include "jaccl/rdma.h" namespace jaccl { @@ -13,14 +15,20 @@ class Config { public: Config(); + Config& set_rank(const char* rank_str); Config& set_rank(int rank); + Config& set_coordinator(const char* coordinator); Config& set_coordinator(std::string coordinator); + Config& set_devices_from_file(const char* dev_file); Config& set_devices( std::vector>> devices); Config& prefer_ring(bool prefer = true); + Config& set_all_gather(AllGatherFn agf); + Config& set_all_gather_factory(std::function factory); bool is_valid_mesh() const; bool is_valid_ring() const; + bool is_valid() const; int get_rank() const { return rank_; @@ -38,7 +46,7 @@ class Config { return prefer_ring_; } - static std::optional from_env(); + static Config from_env(); friend std::shared_ptr init(const Config& cfg, bool strict); @@ -46,12 +54,15 @@ class Config { std::vector get_mesh_connectivity() const; std::pair, std::vector> get_ring_connectivity() const; + SideChannel get_side_channel() const; int rank_; int size_; std::string coordinator_; std::vector>> devices_; bool prefer_ring_; + AllGatherFn all_gather_fn_; + std::function all_gather_factory_; }; /** @@ -77,6 +88,18 @@ bool is_available(); */ std::shared_ptr init(bool strict = false); +/** + * Initialize a JACCL communication group from environment variables, using a + * custom all-gather factory for the side channel. + * + * The factory is called once per rank with the rank and group size, and must + * return an all-gather function that will be used to exchange RDMA connection + * metadata during setup. + */ +std::shared_ptr init( + bool strict, + std::function factory); + /** * Initialize a JACCL communication group from an explicit Config object. */ diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index d120114323..2409e12754 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -9,10 +9,10 @@ namespace jaccl { MeshGroup::MeshGroup( int rank, const std::vector& device_names, - const std::string& coordinator_addr) + SideChannel sc) : rank_(rank), size_(device_names.size()), - side_channel_(rank_, size_, coordinator_addr.c_str()), + side_channel_(std::move(sc)), connections_(create_connections(device_names)) { if (size_ > MESH_MAX_PEERS) { std::ostringstream msg; diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.h b/mlx/distributed/jaccl/lib/jaccl/mesh.h index 2ced24416f..14823db873 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.h @@ -23,7 +23,7 @@ class MeshGroup : public Group { MeshGroup( int rank, const std::vector& device_names, - const std::string& coordinator_addr); + SideChannel sc); int rank() override { return rank_; diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index 7dc23043d2..a0e7153c0d 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -294,7 +294,7 @@ std::vector create_connections( return connections; } -SideChannel::SideChannel(int rank, int size, const char* addr) +TCPAllGather::TCPAllGather(int rank, int size, const char* addr) : rank_(rank), size_(size) { auto address = parse_address(addr); @@ -329,8 +329,29 @@ SideChannel::SideChannel(int rank, int size, const char* addr) } } +void TCPAllGather::operator()(const char* src, char* dst, size_t n_bytes) { + std::lock_guard lock(mutex_); + if (rank_ == 0) { + std::copy(src, src + n_bytes, dst); + for (int i = 1; i < size_; i++) { + sockets_[i - 1].recv(IBV_TAG, dst + i * n_bytes, n_bytes); + } + for (int i = 1; i < size_; i++) { + sockets_[i - 1].send(IBV_TAG, dst, size_ * n_bytes); + } + } else { + sockets_[0].send(IBV_TAG, src, n_bytes); + sockets_[0].recv(IBV_TAG, dst, size_ * n_bytes); + } +} + +SideChannel::SideChannel(int rank, int size, AllGatherFn agf) + : rank_(rank), size_(size), all_gather_fn_(std::move(agf)) {} + SideChannel::SideChannel(SideChannel&& sc) - : rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) { + : rank_(sc.rank_), + size_(sc.size_), + all_gather_fn_(std::move(sc.all_gather_fn_)) { sc.rank_ = -1; sc.size_ = -1; } diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index b28c4ee10e..01d9ed37d6 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -4,6 +4,8 @@ #include +#include +#include #include #include #include @@ -257,6 +259,36 @@ inline int poll( return completions; } +/** + * A function that performs an all-gather across ranks. + * + * Args: + * src: Pointer to this rank's data of size n_bytes. + * dst: Pointer to an output buffer of size size_ * n_bytes. After the call, + * dst[r * n_bytes, (r+1) * n_bytes] contains the data from rank r. + * n_bytes: The number of bytes contributed by each rank. + */ +using AllGatherFn = + std::function; + +class TCPAllGather { + public: + TCPAllGather(int rank, int size, const char* addr); + + TCPAllGather(const TCPAllGather&) = delete; + TCPAllGather(TCPAllGather&&) = delete; + TCPAllGather& operator=(const TCPAllGather&) = delete; + TCPAllGather& operator=(TCPAllGather&&) = delete; + + void operator()(const char* src, char* dst, size_t n_bytes); + + private: + int rank_; + int size_; + std::vector sockets_; + std::mutex mutex_; +}; + /** * Implement a TCP side channel to exchange information about the RDMA * connections. @@ -266,7 +298,7 @@ inline int poll( */ class SideChannel { public: - SideChannel(int rank, int size, const char* addr); + SideChannel(int rank, int size, AllGatherFn agf); SideChannel(SideChannel&& sc); SideChannel(const SideChannel&) = delete; @@ -280,54 +312,37 @@ class SideChannel { if constexpr (is_container::value) { using U = typename T::value_type; - // Share the lengths first and set the communication size to be the + // Share the lengths first to set the communication size to be the // maximum length of the containers. auto lengths = all_gather(v.size()); auto max_len = *std::max_element(lengths.begin(), lengths.end()); - for (auto& s : result) { - s.resize(max_len); - } - // All gather of length max_len - if (rank_ == 0) { - std::copy(v.begin(), v.end(), result[rank_].begin()); - for (int i = 1; i < size_; i++) { - sockets_[i - 1].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); - } - for (int i = 1; i < size_; i++) { - for (int j = 0; j < size_; j++) { - sockets_[i - 1].send( - IBV_TAG, result[j].data(), sizeof(U) * max_len); - } - } - } else { - std::copy(v.begin(), v.end(), result[rank_].begin()); - sockets_[0].send(IBV_TAG, result[rank_].data(), sizeof(U) * max_len); - for (int i = 0; i < size_; i++) { - sockets_[0].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); - } - } + // Allocate flat memory for the all gather + std::vector buffer; + buffer.resize(max_len * (size_ + 1)); + + // Copy our value and share it + std::copy(v.begin(), v.end(), buffer.begin()); + all_gather_fn_( + reinterpret_cast(&buffer[0]), + reinterpret_cast(&buffer[max_len]), + max_len * sizeof(U)); - // Resize the outputs back to the original length + // Put the values into the individual containers for (int i = 0; i < size_; i++) { - result[i].resize(lengths[i]); + std::copy( + buffer.begin() + (i + 1) * max_len, + buffer.begin() + (i + 1) * max_len + lengths[i], + std::inserter(result[i], result[i].end())); } } // T is a scalar else { - if (rank_ == 0) { - result[rank_] = v; - for (int i = 1; i < size_; i++) { - sockets_[i - 1].recv(IBV_TAG, &result[i], sizeof(T)); - } - for (int i = 1; i < size_; i++) { - sockets_[i - 1].send(IBV_TAG, result.data(), size_ * sizeof(T)); - } - } else { - sockets_[0].send(IBV_TAG, &v, sizeof(T)); - sockets_[0].recv(IBV_TAG, result.data(), size_ * sizeof(T)); - } + all_gather_fn_( + reinterpret_cast(&v), + reinterpret_cast(result.data()), + sizeof(T)); } return result; @@ -342,7 +357,7 @@ class SideChannel { private: int rank_; int size_; - std::vector sockets_; + AllGatherFn all_gather_fn_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index 541e825b4b..c2099cb907 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -11,11 +11,11 @@ RingGroup::RingGroup( int size, const std::vector& left_devices, const std::vector& right_devices, - const std::string& coordinator_addr) + SideChannel sc) : rank_(rank), size_(size), n_conns_(left_devices.size()), - side_channel_(rank_, size_, coordinator_addr.c_str()), + side_channel_(std::move(sc)), left_(create_connections(left_devices)), right_(create_connections(right_devices)) { if (left_.size() > RING_MAX_CONNS || right_.size() > RING_MAX_CONNS) { diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h index 016c6af645..a8b427ba80 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring.h @@ -24,7 +24,7 @@ class RingGroup : public Group { int size, const std::vector& left_devices, const std::vector& right_devices, - const std::string& coordinator_addr); + SideChannel sc); int rank() override { return rank_; diff --git a/mlx/distributed/jaccl/no_jaccl.cpp b/mlx/distributed/jaccl/no_jaccl.cpp index 12fd6ab2f1..cab81eae30 100644 --- a/mlx/distributed/jaccl/no_jaccl.cpp +++ b/mlx/distributed/jaccl/no_jaccl.cpp @@ -17,4 +17,11 @@ std::shared_ptr init(bool strict /* = false */) { return nullptr; } +std::shared_ptr init(bool strict, AllGatherFactory /* factory */) { + if (strict) { + throw std::runtime_error("Cannot initialize jaccl distributed backend."); + } + return nullptr; +} + } // namespace mlx::core::distributed::jaccl diff --git a/python/src/distributed.cpp b/python/src/distributed.cpp index 9f4a7cb59e..ed80001df5 100644 --- a/python/src/distributed.cpp +++ b/python/src/distributed.cpp @@ -12,6 +12,9 @@ #include "python/src/small_vector.h" #include "python/src/utils.h" +#include +#include + namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; @@ -75,10 +78,54 @@ void init_distributed(nb::module_& parent_module) { m.def( "init", - &mx::distributed::init, + [](bool strict, + const std::string& backend, + std::optional all_gather_factory) + -> mx::distributed::Group { + if (!all_gather_factory.has_value()) { + return mx::distributed::init(strict, backend); + } + + if (backend != "jaccl") { + throw std::invalid_argument( + "all_gather_factory is only supported with backend='jaccl'."); + } + + auto py_factory = std::move(*all_gather_factory); + auto cpp_factory = [py_factory = std::move(py_factory)]( + int rank, + int size) -> mx::distributed::AllGatherFn { + nb::gil_scoped_acquire gil; + nb::object py_inner = py_factory(rank, size); + if (!PyCallable_Check(py_inner.ptr())) { + throw std::invalid_argument( + "all_gather_factory must return a callable"); + } + return [py_inner = std::move(py_inner), size]( + const void* src, void* dst, size_t n_bytes) { + nb::gil_scoped_acquire gil; + nb::bytes src_bytes(src, n_bytes); + nb::object result_obj = py_inner(src_bytes, n_bytes); + nb::bytes result = nb::cast(result_obj); + size_t expected = static_cast(size) * n_bytes; + if (result.size() != expected) { + std::ostringstream msg; + msg << "Custom all-gather returned " << result.size() + << " bytes but expected " << expected; + throw std::runtime_error(msg.str()); + } + std::memcpy(dst, result.data(), expected); + }; + }; + + return mx::distributed::init(strict, backend, std::move(cpp_factory)); + }, "strict"_a = false, "backend"_a = "any", - nb::sig("def init(strict: bool = False, backend: str = 'any') -> Group"), + nb::kw_only(), + "all_gather_factory"_a = nb::none(), + nb::sig( + "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Optional[Callable[[int, int], Callable[[bytes, int], bytes]]] = None) -> Group"), R"pbdoc( Initialize the communication backend and create the global communication group. @@ -99,6 +146,13 @@ void init_distributed(nb::module_& parent_module) { set to ``any`` all available backends are tried and the first one that succeeds becomes the global group which will be returned in subsequent calls. Default: ``any`` + all_gather_factory (Callable, optional): A factory used only with the + ``jaccl`` backend. It is called once per rank with ``(rank, size)`` + and must return a callable with signature + ``f(src: bytes, n_bytes: int) -> bytes``. The returned callable + performs a byte-level all-gather used as the JACCL side channel + when exchanging RDMA connection metadata. The returned bytes must + have length ``size * n_bytes``. Returns: Group: The group representing all the launched processes. @@ -349,4 +403,10 @@ void init_distributed(nb::module_& parent_module) { Returns: array: The output array with shape ``[x.shape[0] // group.size(), *x.shape[1:]]``. )pbdoc"); + + // Ensure the distributed backend cache is cleared before the interpreter + // goes away, so that any Python objects held by cached groups are released + // while Python is still alive. + auto atexit = nb::module_::import_("atexit"); + atexit.attr("register")(nb::cpp_function(&mx::distributed::clear_backends)); } diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 77bdea6a1f..cbb9663046 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -360,3 +360,15 @@ def test_clip_grad_norm_sharded(self): clipped[k], grads_slice[k] * scale, atol=self.atol, rtol=self.rtol ) ) + + def test_jaccl_all_gather_factory_validation(self): + # A custom side-channel factory is only valid with the jaccl backend. + with self.assertRaises(ValueError): + mx.distributed.init( + backend="ring", + all_gather_factory=lambda rank, size: lambda src, n_bytes: b"", + ) + + # The factory must be callable. + with self.assertRaises(TypeError): + mx.distributed.init(backend="jaccl", all_gather_factory="not_callable") From 2263a6b4d00dae2235c97100eeb094f1719b929e Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 30 Jul 2026 14:58:40 -0700 Subject: [PATCH 023/222] Fix docstring mismatches in the Python bindings (#3948) --- python/src/export.cpp | 2 +- python/src/fft.cpp | 1 + python/src/ops.cpp | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/src/export.cpp b/python/src/export.cpp index 51c8a8965c..b4c2998a03 100644 --- a/python/src/export.cpp +++ b/python/src/export.cpp @@ -249,7 +249,7 @@ void init_export(nb::module_& m) { A context managing class for exporting multiple traces of the same function to a file. - Make an instance of this class by calling fun:`mx.exporter`. + Make an instance of this class by calling :func:`mx.exporter`. )pbdoc") .def("close", &PyFunctionExporter::close) .def("__enter__", [](PyFunctionExporter& exporter) { return &exporter; }) diff --git a/python/src/fft.cpp b/python/src/fft.cpp index eb6d531c92..99f8478bbb 100644 --- a/python/src/fft.cpp +++ b/python/src/fft.cpp @@ -494,6 +494,7 @@ void init_fft(nb::module_& parent_module) { Returns: array: The real DFT of the input along the given axes. The output + data type will be complex. )pbdoc"); m.def( "irfftn", diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 19b19bd78f..5658a78648 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3479,7 +3479,7 @@ void init_ops(nb::module_& m) { mode: Padding mode. One of the following strings: "constant" (default): Pads with a constant value. "edge": Pads with the edge values of array. - constant_value (array or scalar, optional): Optional constant value + constant_values (array or scalar, optional): Optional constant value to pad the edges of the array with. Returns: @@ -4724,13 +4724,13 @@ void init_ops(nb::module_& m) { bits (int, optional): The number of bits occupied by each element of ``w`` in the quantized array. See supported values and defaults in the :ref:`table of quantization modes `. Default: ``None``. + mode (str, optional): The quantization mode. Default: ``"affine"``. global_scale (array, optional): The per-input float32 scale used for ``"nvfp4"`` quantization if provided. Default: ``None``. dtype (Dtype, optional): The data type of the dequantized output. If ``None`` the return type is inferred from the scales and biases when possible and otherwise defaults to ``bfloat16``. Default: ``None``. - mode (str, optional): The quantization mode. Default: ``"affine"``. Returns: array: The dequantized version of ``w`` From 2ad0d4d311f54de855b06cd21ca85d3b628c1012 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 30 Jul 2026 17:42:17 -0700 Subject: [PATCH 024/222] Raise a clear error for an invalid quantization mode in nn layers (#3914) Co-authored-by: Anastasiia Filippova --- python/mlx/nn/layers/quantized.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py index 57e7c88898..7cebb5b2da 100644 --- a/python/mlx/nn/layers/quantized.py +++ b/python/mlx/nn/layers/quantized.py @@ -15,6 +15,11 @@ def _defaults_for_mode(mode, group_size, bits): "nvfp4": (16, 4), "mxfp8": (32, 8), } + if mode not in mode_defaults: + raise ValueError( + f"Invalid quantization mode '{mode}'. " + f"Valid modes are: {', '.join(mode_defaults)}." + ) default_group_size, default_bits = mode_defaults[mode] return group_size or default_group_size, bits or default_bits From fb5133e1049cc1482a330cd3a7135fda62a74b0d Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 31 Jul 2026 10:59:18 -0700 Subject: [PATCH 025/222] Fix incorrect examples and outputs in the usage docs (#3956) --- docs/src/usage/compile.rst | 8 ++++---- docs/src/usage/indexing.rst | 15 +++++++-------- docs/src/usage/quick_start.rst | 2 +- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/docs/src/usage/compile.rst b/docs/src/usage/compile.rst index bfe7f09661..662f69410c 100644 --- a/docs/src/usage/compile.rst +++ b/docs/src/usage/compile.rst @@ -207,9 +207,9 @@ You have two options to deal with this. The first option is to simply return state.append(z) return mx.exp(z), state - _, state = fun(mx.array(1.0), mx.array(2.0)) - # Prints [array(3, dtype=float32)] - print(state) + _, state = fun(mx.array(1.0), mx.array(2.0)) + # Prints [array(3, dtype=float32)] + print(state) In some cases returning updated state can be pretty inconvenient. Hence, :func:`compile` has a parameter to capture implicit outputs: @@ -464,7 +464,7 @@ recompiled. x = mx.array(1.0) y = mx.array(-2.0) - # Firt call compiles the function + # First call compiles the function print(compiled_fun(x, y)) # Second call with different shapes diff --git a/docs/src/usage/indexing.rst b/docs/src/usage/indexing.rst index 8dc01cfc8e..9bce59a6fb 100644 --- a/docs/src/usage/indexing.rst +++ b/docs/src/usage/indexing.rst @@ -28,12 +28,11 @@ For multi-dimensional arrays, the ``...`` or :obj:`Ellipsis` syntax works as in >>> arr = mx.arange(8).reshape(2, 2, 2) >>> arr[:, :, 0] - array(3, dtype=int32) array([[0, 2], - [4, 6]], dtype=int32 + [4, 6]], dtype=int32) >>> arr[..., 0] array([[0, 2], - [4, 6]], dtype=int32 + [4, 6]], dtype=int32) You can index with ``None`` to create a new axis: @@ -41,9 +40,9 @@ You can index with ``None`` to create a new axis: >>> arr = mx.arange(8) >>> arr.shape - [8] + (8,) >>> arr[None].shape - [1, 8] + (1, 8) You can also use an :obj:`array` to index another :obj:`array`: @@ -161,7 +160,7 @@ Other index types are routed through the standard scatter code. >>> updates = mx.array([5.0, 6.0]) >>> a[mask] = updates >>> a - array([5.0, 2.0, 6.0], dtype=float32) + array([5, 2, 6], dtype=float32) Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar assignments, ``updates`` must provide at least as many elements as there are @@ -174,8 +173,8 @@ assignments, ``updates`` must provide at least as many elements as there are [False, False, True]]) >>> a[mask] = 1.0 >>> a - array([[1.0, 0.0, 1.0], - [0.0, 0.0, 1.0]], dtype=float32) + array([[1, 0, 1], + [0, 0, 1]], dtype=float32) Boolean masks follow NumPy semantics: diff --git a/docs/src/usage/quick_start.rst b/docs/src/usage/quick_start.rst index bc1d92ad63..8c952a6f74 100644 --- a/docs/src/usage/quick_start.rst +++ b/docs/src/usage/quick_start.rst @@ -14,7 +14,7 @@ Import ``mlx.core`` and make an :class:`array`: >> import mlx.core as mx >> a = mx.array([1, 2, 3, 4]) >> a.shape - [4] + (4,) >> a.dtype int32 >> b = mx.array([1.0, 2.0, 3.0, 4.0]) From 121df0582eb721bbd6970fbc2febd9492b2a9a26 Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 4 Aug 2026 02:15:59 +0900 Subject: [PATCH 026/222] Skip test_gather_qmm_sorted cpu test on M1 mac (#3973) --- .github/actions/setup/action.yml | 1 + python/tests/test_quantized.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+) diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 898bd0f9a4..f0344a427e 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -55,6 +55,7 @@ runs: brew update brew install openmpi xcodebuild -showComponent MetalToolchain + sysctl -a | grep machdep.cpu echo "::endgroup::" - name: Setup Windows environment diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index b993d57bb7..7fbbe6938e 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,5 +1,7 @@ # Copyright © 2023 Apple Inc. +import platform +import subprocess import unittest from itertools import product @@ -7,6 +9,14 @@ import mlx_tests +def is_m1_mac(): + if platform.system() != "Darwin": + return False + cmd = "sysctl -n machdep.cpu.brand_string" + cpu = subprocess.check_output(cmd, shell=True).decode().strip() + return cpu.startswith("Apple M1") + + class TestQuantized(mlx_tests.MLXTestCase): def test_quantize_dequantize(self): w = mx.random.normal(shape=(128, 512)) @@ -1166,6 +1176,10 @@ def test_gather_qmm_matrix_path(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-1) + @unittest.skipIf( + is_m1_mac() and not mx.metal.is_available(), + "Accelerate bug https://github.com/ml-explore/mlx/pull/3563", + ) def test_gather_qmm_sorted(self): def quantize(w, transpose=True, group_size=None, mode="affine"): if mode == "affine": From 12b5eb99ce9d3bb21cf01b77a781a53c9ffd83df Mon Sep 17 00:00:00 2001 From: Daniel L <68606021+danlee2002@users.noreply.github.com> Date: Mon, 3 Aug 2026 18:59:16 -0400 Subject: [PATCH 027/222] Fixes an axis mismatch bug in matrix norm for case -1 and 1 (#3827) --- mlx/linalg.cpp | 10 ++-------- python/tests/test_linalg.py | 6 +++--- tests/linalg_tests.cpp | 15 +++++++++++++++ 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index bdf05fb80b..75a662b1e6 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -84,8 +84,8 @@ inline array matrix_norm( bool keepdims, StreamOrDevice s) { auto dtype = at_least_float(a.dtype()); - auto row_axis = axis[0]; - auto col_axis = axis[1]; + int row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; + int col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; if (ord == -1.0) { col_axis -= (!keepdims && col_axis > row_axis && col_axis > 0); return astype( @@ -99,24 +99,18 @@ inline array matrix_norm( dtype, s); } else if (ord == std::numeric_limits::infinity()) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0); return astype( max(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s), dtype, s); } else if (ord == -std::numeric_limits::infinity()) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0); return astype( min(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s), dtype, s); } else if (ord == 2.0 || ord == -2.0) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; auto a_matrix = (row_axis > col_axis) ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s) : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s); diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index afdf75d7c8..39f0a3b891 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -65,8 +65,8 @@ def test_norm(self): with self.subTest(shape=shape, keepdims=keepdims): self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)) - # neg/pos inf norm test - norms = [-float("inf"), float("inf")] + # tests for negative indexing: -1/1/inf/-inf/ + norms = [-1, 1, -float("inf"), float("inf")] for shape in [(3, 3), (2, 3, 3), (2, 3, 3, 3)]: x_mx = mx.arange(1, math.prod(shape) + 1, dtype=mx.float32).reshape(shape) x_np = np.arange(1, math.prod(shape) + 1, dtype=np.float32).reshape(shape) @@ -76,7 +76,7 @@ def test_norm(self): for axes in neg_axes: out_np = np.linalg.norm( x_np, - ord=np.inf if ord == float("inf") else -np.inf, + ord=ord, axis=tuple(axes), ) out_mx = mx.linalg.norm(x_mx, ord=ord, axis=axes) diff --git a/tests/linalg_tests.cpp b/tests/linalg_tests.cpp index 7c81062a38..9c2494df4c 100644 --- a/tests/linalg_tests.cpp +++ b/tests/linalg_tests.cpp @@ -189,6 +189,13 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") { Device::cpu) .item(), doctest::Approx(3.0)); + CHECK_EQ( + norm(x, -1, std::vector{-1, -2}, false, Device::cpu).item(), + doctest::Approx(3.0)); + CHECK_EQ( + norm(x, 1, std::vector{-1, -2}, false, Device::cpu).item(), + doctest::Approx(21.0)); + x = reshape(arange(18, float32), {2, 3, 3}); CHECK_THROWS(norm(x, 2.0, std::vector{0, 1, 2})); CHECK(allclose( @@ -283,6 +290,14 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") { Device::cpu), array({3.0, 30.0})) .item()); + CHECK(allclose( + norm(x, 1, std::vector{-1, -2}, false, Device::cpu), + array({21.0, 48.0})) + .item()); + CHECK(allclose( + norm(x, -1, std::vector{-1, -2}, false, Device::cpu), + array({3.0, 30.0})) + .item()); } TEST_CASE("[mlx.core.linalg.norm] string ord") { From 43c9e03912ec19e5b6f28b24a6a49d1fa2f923fe Mon Sep 17 00:00:00 2001 From: Ishti <45158028+ishtihoss@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:28:16 -0700 Subject: [PATCH 028/222] Fix BatchNorm running variance estimator (#3817) --- python/mlx/nn/layers/normalization.py | 17 +++- python/tests/test_nn.py | 133 +++++++++++++++++++------- 2 files changed, 113 insertions(+), 37 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index 56bcc3ccba..cbd9861a93 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -336,13 +336,14 @@ def _extra_repr(self): f"track_running_stats={self.track_running_stats}" ) - def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: + def _calc_stats(self, x: mx.array, ddof: int = 0) -> Tuple[mx.array, mx.array]: """ Calculate the mean and variance of the input tensor across the batch and spatial dimensions. Args: x (array): Input tensor. + ddof (int): Delta degrees of freedom for variance. Returns: tuple: Tuple containing mean and variance. @@ -350,7 +351,7 @@ def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: reduction_axes = tuple(range(0, x.ndim - 1)) mean = mx.mean(x, axis=reduction_axes) - var = mx.var(x, axis=reduction_axes) + var = mx.var(x, axis=reduction_axes, ddof=ddof) return mean, var @@ -369,13 +370,23 @@ def __call__(self, x: mx.array) -> mx.array: f"Expected input tensor to have 2, 3 or 4 dimensions, but got {x.ndim}" ) + if self.training: + stats_size = 1 + for size in x.shape[:-1]: + stats_size *= size + if stats_size == 1: + raise ValueError( + "BatchNorm training requires more than one value per channel." + ) + # Calculate the mean and variance used to normalize the input x. If we # are in training mode update the running stats if needed. mean, var = self._calc_stats(x) if self.training and self.track_running_stats: mu = self.momentum + _, running_var = self._calc_stats(x, ddof=1) self.running_mean = (1 - mu) * self.running_mean + mu * mean - self.running_var = (1 - mu) * self.running_var + mu * var + self.running_var = (1 - mu) * self.running_var + mu * running_var elif self.track_running_stats: mean = self.running_mean var = self.running_var diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 8ad72a323c..0df667a64d 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -10,6 +10,13 @@ import numpy as np from mlx.utils import tree_flatten, tree_map, tree_reduce +try: + import torch + + has_torch = True +except ImportError: + has_torch = False + class TestBase(mlx_tests.MLXTestCase): def test_module_utilities(self): @@ -672,7 +679,7 @@ def test_batch_norm(self): ], ) expected_mean = mx.array([0.008929, 0.005680, -0.016092, 0.027778]) - expected_var = mx.array([0.928435, 1.00455, 1.04117, 0.94258]) + expected_var = mx.array([0.935544, 1.030691, 1.076463, 0.953224]) self.assertTrue(x.shape == y.shape) self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) @@ -683,11 +690,11 @@ def test_batch_norm(self): y = bn(x) expected_y = mx.array( [ - [-0.15984, 1.73159, -1.25456, 1.57891], - [-0.872193, -1.4281, -0.414439, -0.228678], - [0.602743, -0.30566, -0.554687, 0.139639], - [0.252199, 0.29066, -0.599572, -0.0512532], - [0.594096, -0.0334829, 2.11359, -0.151081], + [-0.159232, 1.70949, -1.23382, 1.57007], + [-0.868873, -1.40987, -0.407588, -0.227397], + [0.600449, -0.301759, -0.545518, 0.138857], + [0.251239, 0.286951, -0.589661, -0.0509662], + [0.591834, -0.0330556, 2.07865, -0.150235], ] ) @@ -740,9 +747,9 @@ def test_batch_norm(self): ) self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) expected_mean = mx.array( - [[[0.00207845, -5.3259e-05, 0.04755, -0.0697296, 0.0236228]]] + [0.00207845, -5.3259e-05, 0.04755, -0.0697296, 0.0236228] ) - expected_var = mx.array([[[0.968415, 1.05322, 0.96913, 0.932305, 0.967224]]]) + expected_var = mx.array([0.978188, 1.07511, 0.979006, 0.93692, 0.976827]) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_var, expected_var, atol=1e-5)) @@ -780,46 +787,104 @@ def test_batch_norm(self): self.assertTrue(mx.allclose(y.mean(axis=(0, 1, 2)), mx.zeros((6,)), atol=1e-5)) self.assertTrue(mx.allclose(y.var(axis=(0, 1, 2)), mx.ones((6,)), atol=1e-2)) + @unittest.skipIf(not has_torch, "requires Torch") + def test_batch_norm_matches_torch(self): + rng = np.random.default_rng(0) + momentum = 0.1 + eps = 1e-5 + + def check_batch_norm(shape, torch_module, to_torch=None, from_torch=None): + features = shape[-1] + x_np = rng.normal(size=shape).astype(np.float32) + weight_np = rng.normal(size=(features,)).astype(np.float32) + bias_np = rng.normal(size=(features,)).astype(np.float32) + + mlx_bn = nn.BatchNorm(features, eps=eps, momentum=momentum) + mlx_bn.weight = mx.array(weight_np) + mlx_bn.bias = mx.array(bias_np) + mlx_y = mlx_bn(mx.array(x_np)) + mx.eval(mlx_y, mlx_bn.running_mean, mlx_bn.running_var) + + torch_bn = torch_module(features, eps=eps, momentum=momentum) + with torch.no_grad(): + torch_bn.weight.copy_(torch.from_numpy(weight_np)) + torch_bn.bias.copy_(torch.from_numpy(bias_np)) + x_torch_np = x_np.transpose(to_torch) if to_torch else x_np + torch_y = torch_bn(torch.from_numpy(x_torch_np)).detach().numpy() + if from_torch: + torch_y = torch_y.transpose(from_torch) + + self.assertTrue(mx.allclose(mlx_y, mx.array(torch_y), rtol=1e-4, atol=1e-4)) + self.assertTrue( + mx.allclose( + mlx_bn.running_mean, + mx.array(torch_bn.running_mean.detach().numpy()), + rtol=1e-5, + atol=1e-5, + ) + ) + self.assertTrue( + mx.allclose( + mlx_bn.running_var, + mx.array(torch_bn.running_var.detach().numpy()), + rtol=1e-4, + atol=1e-4, + ) + ) + + mlx_bn.eval() + torch_bn.eval() + mlx_y = mlx_bn(mx.array(x_np)) + mx.eval(mlx_y) + torch_y = torch_bn(torch.from_numpy(x_torch_np)).detach().numpy() + if from_torch: + torch_y = torch_y.transpose(from_torch) + self.assertTrue(mx.allclose(mlx_y, mx.array(torch_y), rtol=1e-4, atol=1e-4)) + + check_batch_norm((5, 4), torch.nn.BatchNorm1d) + check_batch_norm( + (2, 4, 5), + torch.nn.BatchNorm1d, + to_torch=(0, 2, 1), + from_torch=(0, 2, 1), + ) + check_batch_norm( + (2, 3, 3, 6), + torch.nn.BatchNorm2d, + to_torch=(0, 3, 1, 2), + from_torch=(0, 2, 3, 1), + ) + def test_batch_norm_stats(self): batch_size = 2 num_features = 4 h = 3 w = 3 - momentum = 0.1 batch_norm = nn.BatchNorm(num_features) - batch_norm.train() - running_mean = batch_norm.running_mean - running_var = batch_norm.running_var - - data = mx.random.normal((batch_size, num_features)) + data = mx.random.normal((batch_size, h, w, num_features)) normalized_data = batch_norm(data) - means = mx.mean(data, axis=0) - variances = mx.var(data, axis=0) - running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * variances - self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + self.assertTrue( + mx.allclose( + mx.mean(normalized_data, axis=(0, 1, 2)), mx.zeros((4,)), atol=1e-5 + ) + ) + self.assertTrue( + mx.allclose( + mx.var(normalized_data, axis=(0, 1, 2)), mx.ones((4,)), atol=1e-2 + ) + ) + self.assertEqual(batch_norm.running_mean.shape, (num_features,)) + self.assertEqual(batch_norm.running_var.shape, (num_features,)) batch_norm = nn.BatchNorm(num_features) - batch_norm.train() - running_mean = batch_norm.running_mean - running_var = batch_norm.running_var - data = mx.random.normal((batch_size, h, w, num_features)) + data = mx.random.normal((1, num_features)) - normalized_data = batch_norm(data) - means = mx.mean(data, axis=(0, 1, 2)) - variances = mx.var(data, axis=(0, 1, 2)) - running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * variances - self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) - - self.assertEqual(batch_norm.running_mean.shape, running_mean.shape) - self.assertEqual(batch_norm.running_var.shape, running_var.shape) + with self.assertRaises(ValueError): + batch_norm(data) def test_conv1d(self): N = 5 From b6faa911efceecd350d61f09aaf3d627aac5bd41 Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 4 Aug 2026 08:59:47 +0900 Subject: [PATCH 029/222] Fix build error caused by TMA macro guard (#3988) --- mlx/backend/cuda/ptx.cuh | 12 +++++------- mlx/backend/cuda/quantized/fp_quantize.cuh | 8 ++++---- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/mlx/backend/cuda/ptx.cuh b/mlx/backend/cuda/ptx.cuh index 6ec7caedd6..b7b87e671e 100644 --- a/mlx/backend/cuda/ptx.cuh +++ b/mlx/backend/cuda/ptx.cuh @@ -1,14 +1,12 @@ #pragma once -#include -#include +#include namespace mlx::core { namespace ptx { -#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ - defined(__CUDA_ARCH_SPECIFIC__) +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) __device__ __forceinline__ void mbarrier_init(uint64_t* mbar, uint32_t count) { uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); @@ -121,7 +119,7 @@ __device__ __forceinline__ void fence_proxy_async_shared_cta() { asm volatile("fence.proxy.async.shared::cta;"); } -#endif // (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && - // (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000) +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) + } // namespace ptx -} // namespace mlx::core \ No newline at end of file +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/fp_quantize.cuh b/mlx/backend/cuda/quantized/fp_quantize.cuh index 769794129c..bea4e21979 100644 --- a/mlx/backend/cuda/quantized/fp_quantize.cuh +++ b/mlx/backend/cuda/quantized/fp_quantize.cuh @@ -55,7 +55,7 @@ __device__ __forceinline__ void copy_2d_to_shared( uint32_t num_bytes, uint64_t* barrier, const bool is_master_thread) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) if (is_master_thread) { // Arrive and tell how many bytes are expected ptx::mbarrier_arrive_expect_tx(barrier, num_bytes); @@ -66,7 +66,7 @@ __device__ __forceinline__ void copy_2d_to_shared( // Other threads just arrive ptx::mbarrier_arrive(barrier); } -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) } namespace cg = cooperative_groups; @@ -228,7 +228,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) uint8_t* __restrict__ scales, const size_t rows, const size_t cols) { -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) using Tx2 = Vector2_t; using Tx4 = Vector4_t; @@ -408,7 +408,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) ptx::mbarrier_invalidate(&mbar[iter]); } } -#endif // __CUDA_ARCH__ >= 1000 +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) } // TODO: add kernel with tma instructions From ddb81627fa906a953861f26900028dd10ef0f347 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Mon, 3 Aug 2026 17:34:14 -0700 Subject: [PATCH 030/222] Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma (#3831) --- python/mlx/nn/init.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/python/mlx/nn/init.py b/python/mlx/nn/init.py index 3fdf42990b..f0bda14a66 100644 --- a/python/mlx/nn/init.py +++ b/python/mlx/nn/init.py @@ -194,12 +194,13 @@ def glorot_uniform( ) -> Callable[[mx.array, float], mx.array]: r"""A Glorot uniform initializer. - This initializer samples from a uniform distribution with a range - computed from the number of input (``fan_in``) and output (``fan_out``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is computed from the number of input (``fan_in``) and output (``fan_out``) units according to: .. math:: - \sigma = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}} + \text{limit} = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}} For more details see the original reference: `Understanding the difficulty of training deep feedforward neural networks @@ -295,13 +296,14 @@ def he_uniform( ) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]: r"""A He uniform (Kaiming uniform) initializer. - This initializer samples from a uniform distribution with a range - computed from the number of input (``fan_in``) or output (``fan_out``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is computed from the number of input (``fan_in``) or output (``fan_out``) units according to: .. math:: - \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}} + \text{limit} = \gamma \sqrt{\frac{3.0}{\text{fan}}} where :math:`\text{fan}` is either the number of input units when the ``mode`` is ``"fan_in"`` or output units when the ``mode`` is From 697828f24fabaa5c2cb50b98a70eac36157eb333 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Mon, 3 Aug 2026 20:45:06 -0400 Subject: [PATCH 031/222] Fix custom metal kernel cache collision for same name, different source (#3833) Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com> Co-authored-by: Cheng --- mlx/backend/metal/custom_kernel.cpp | 33 +++++------------------------ python/tests/test_fast.py | 32 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 28 deletions(-) diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index 0648ed221e..b73dd72def 100644 --- a/mlx/backend/metal/custom_kernel.cpp +++ b/mlx/backend/metal/custom_kernel.cpp @@ -6,17 +6,9 @@ #include "mlx/backend/metal/utils.h" #include "mlx/fast_primitives.h" -namespace mlx::core::fast { - -struct CustomKernelCache { - std::unordered_map> - libraries; -}; +#include -static CustomKernelCache& cache() { - static CustomKernelCache cache_; - return cache_; -}; +namespace mlx::core::fast { void CustomKernel::eval_gpu( const std::vector& inputs, @@ -55,25 +47,10 @@ void CustomKernel::eval_gpu( auto& d = metal::device(s.device); - { - // Clear kernels from the device library cache if needed - auto& kernel_cache = cache(); - if (auto it = kernel_cache.libraries.find(name_); - it != kernel_cache.libraries.end()) { - if (it->second.first != source_ || - it->second.second != compile_options_) { - auto& d = metal::device(s.device); - d.clear_library(name_); - it->second = std::make_tuple(source_, compile_options_); - } - } else { - kernel_cache.libraries.emplace( - name_, std::make_tuple(source_, compile_options_)); - } - } - + std::string lib_name = fmt::format( + "{}_{:x}_{}", name_, std::hash{}(source_), compile_options_); auto lib = d.get_library( - name_, compile_options_, [this] { return metal::utils() + source_; }); + lib_name, compile_options_, [this] { return metal::utils() + source_; }); auto kernel = d.get_kernel(name_, lib); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 8343c45941..5dacaa605c 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1026,6 +1026,38 @@ def call_kernel(a: mx.array, source): out = call_kernel(a, source) self.assertTrue(mx.array_equal(out, mx.ones_like(out))) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_kernel_same_name_different_source_one_eval(self): + # Regression test for #3832: two kernels sharing a name but with + # different sources, dispatched in a SINGLE eval batch, must each run + # their own compiled code instead of silently reusing the first's. + def call_kernel(a, source): + kernel = mx.fast.metal_kernel( + name="dup_name", + input_names=["inp"], + output_names=["out"], + source=source, + ) + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.arange(32, dtype=mx.float32) + out_a = call_kernel( + a, "uint e = thread_position_in_grid.x; out[e] = inp[e] * 2.0f;" + ) + out_b = call_kernel( + a, "uint e = thread_position_in_grid.x; out[e] = inp[e] + 100.0f;" + ) + mx.eval(out_a, out_b) # one batch — the reported failure case + self.assertTrue(mx.array_equal(out_a, a * 2.0)) + self.assertTrue(mx.array_equal(out_b, a + 100.0)) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") def test_custom_metal_kernel_math_mode(self): with self.assertRaises(ValueError): From c0d916b54f3ed0ec2a3a82a6521b650a3ce485f3 Mon Sep 17 00:00:00 2001 From: Madan kumar Date: Tue, 4 Aug 2026 11:43:34 +0530 Subject: [PATCH 032/222] Fix log_cosh_loss docstring to document the element-wise loss (#3846) --- python/mlx/nn/losses.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 85b11255bb..9218691e14 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -505,8 +505,7 @@ def log_cosh_loss( .. math:: \text{logcosh}(y_{\text{true}}, y_{\text{pred}}) = - \frac{1}{n} \sum_{i=1}^{n} - \log(\cosh(y_{\text{pred}}^{(i)} - y_{\text{true}}^{(i)})) + \log(\cosh(y_{\text{pred}} - y_{\text{true}})) Args: From a2fa5374e21de3c8245a58d54c6bafb1f5716a8c Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 3 Aug 2026 23:15:11 -0700 Subject: [PATCH 033/222] Fix step activation docstring to match >= threshold behavior (#3902) --- python/mlx/nn/layers/activations.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/mlx/nn/layers/activations.py b/python/mlx/nn/layers/activations.py index 4f36ae2b38..58fef50687 100644 --- a/python/mlx/nn/layers/activations.py +++ b/python/mlx/nn/layers/activations.py @@ -225,7 +225,8 @@ def step(x: mx.array, threshold: float = 0.0): r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set - to 1 if the input is greater than a specified threshold, and 0 otherwise. + to 1 if the input is greater than or equal to a specified threshold, and 0 + otherwise. .. math:: \text{step}(x) = \begin{cases} @@ -234,7 +235,7 @@ def step(x: mx.array, threshold: float = 0.0): \end{cases} Args: - threshold: The value to threshold at. + threshold: The value to threshold at. Default: ``0.0``. """ return mx.where(x >= threshold, 1, 0) @@ -606,7 +607,8 @@ class Step(Module): r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set - to 1 if the input is greater than a specified threshold, and 0 otherwise. + to 1 if the input is greater than or equal to a specified threshold, and 0 + otherwise. .. math:: \text{step}(x) = \begin{cases} @@ -615,7 +617,7 @@ class Step(Module): \end{cases} Args: - threshold: The value to threshold at. + threshold: The value to threshold at. Default: ``0.0``. """ def __init__(self, threshold: float = 0.0): From 0b5e91f7949ef34ef15596f54d51787ab70d4d5f Mon Sep 17 00:00:00 2001 From: stoyoda0012-cyber <253017259+stoyoda0012-cyber@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:16:13 +0900 Subject: [PATCH 034/222] docs: document the reduced-precision float32 default and MLX_ENABLE_TF32 (#3894) Co-authored-by: Satoshi Toyoda Co-authored-by: Claude Fable 5 --- docs/src/index.rst | 1 + docs/src/usage/precision.rst | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+) create mode 100644 docs/src/usage/precision.rst diff --git a/docs/src/index.rst b/docs/src/index.rst index 46d069929f..d9beb16753 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -43,6 +43,7 @@ are the CPU and GPU. usage/function_transforms usage/compile usage/numpy + usage/precision usage/distributed usage/using_streams usage/export diff --git a/docs/src/usage/precision.rst b/docs/src/usage/precision.rst new file mode 100644 index 0000000000..b5ad9a4f9e --- /dev/null +++ b/docs/src/usage/precision.rst @@ -0,0 +1,21 @@ +.. _precision: + +Numerical Precision +=================== + +By default, MLX may run ``float32`` matrix-multiplication family +operations (matmul, quantized matmul, grouped matmul, convolution and +attention) at reduced precision on hardware with dedicated +matrix-multiplication units. Inputs and outputs stay ``float32``, but +results can differ from a full-precision reference by several orders of +magnitude more than ``float32`` rounding alone would explain. + +To keep these operations in full ``float32``, set ``MLX_ENABLE_TF32=0`` +when launching the process: + +.. code-block:: shell + + MLX_ENABLE_TF32=0 python my_script.py + +Which operations take the reduced-precision path, and how large the +difference is, depends on the backend and the hardware. From f75ee500d879e7562ae128333317b27a46f5a301 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 3 Aug 2026 23:18:26 -0700 Subject: [PATCH 035/222] Fix InstanceNorm Shape docstring to require at least 3 dimensions (#3903) --- python/mlx/nn/layers/normalization.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index cbd9861a93..e79440dce3 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -25,7 +25,8 @@ class InstanceNorm(Module): affine (bool): Default: ``False``. Shape: - - Input: :math:`(..., C)` where :math:`C` is equal to :attr:`dims`. + - Input: :math:`(N, ..., C)` where :math:`C` is equal to :attr:`dims`. + The input must have at least 3 dimensions. - Output: Same shape as the input. Examples: From cff8e0f4216a969f4710780275d90170133106ab Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 3 Aug 2026 23:18:43 -0700 Subject: [PATCH 036/222] Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn (#3906) --- python/mlx/nn/layers/base.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 9802c6c311..c2bcf7ff36 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -255,13 +255,13 @@ def filter_and_map( but it can also be used to extract any subset of the module's parameters. Args: - filter_fn (Callable): Given a value, the key in which it is found - and the containing module, decide whether to keep the value or + filter_fn (Callable): Given the containing module, the key in which + it is found and the value, decide whether to keep the value or drop it. map_fn (Callable, optional): Optionally transform the value before returning it. - is_leaf_fn (Callable, optional): Given a value, the key in which it - is found and the containing module decide if it is a leaf. + is_leaf_fn (Callable, optional): Given the containing module, the + key in which it is found and the value decide if it is a leaf. Returns: A dictionary containing the contents of the module recursively filtered From 24482988e9cf08b6f3bbdd7b2af505d028d389e5 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Tue, 4 Aug 2026 00:17:53 -0700 Subject: [PATCH 037/222] Export C++20 requirement to CMake consumers (#3971) --- CMakeLists.txt | 1 + examples/cmake_project/CMakeLists.txt | 2 +- examples/export/CMakeLists.txt | 2 +- examples/extensions/CMakeLists.txt | 2 +- mlx.pc.in | 1 - 5 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f8ee65ad4b..ca689e13d1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -151,6 +151,7 @@ cmake_policy(SET CMP0135 NEW) add_library(mlx) +target_compile_features(mlx INTERFACE cxx_std_20) target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS}) target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS}) diff --git a/examples/cmake_project/CMakeLists.txt b/examples/cmake_project/CMakeLists.txt index 7a73ce1d9d..1d1b9d3db8 100644 --- a/examples/cmake_project/CMakeLists.txt +++ b/examples/cmake_project/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.27) project(example LANGUAGES CXX) -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) # Comment the following two commands only the MLX C++ library is installed and diff --git a/examples/export/CMakeLists.txt b/examples/export/CMakeLists.txt index aefcdf59fb..abe6e7b5d8 100644 --- a/examples/export/CMakeLists.txt +++ b/examples/export/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.27) project(import_mlx LANGUAGES CXX) -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) find_package( diff --git a/examples/extensions/CMakeLists.txt b/examples/extensions/CMakeLists.txt index 0f70187209..eedc3a7dab 100644 --- a/examples/extensions/CMakeLists.txt +++ b/examples/extensions/CMakeLists.txt @@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.27) project(_ext LANGUAGES CXX) # ----------------------------- Setup ----------------------------- -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) diff --git a/mlx.pc.in b/mlx.pc.in index c4e2515d71..2889149db6 100644 --- a/mlx.pc.in +++ b/mlx.pc.in @@ -44,7 +44,6 @@ if (@MLX_BUILD_METAL@) endif() set_target_properties(mlx PROPERTIES - CXX_STANDARD 17 INTERFACE_COMPILE_OPTIONS "${MLX_CXX_FLAGS}" ) From ef5fc0fab27c63e6976e96b07c234c9e2601d8a7 Mon Sep 17 00:00:00 2001 From: Valentin Roussellet Date: Tue, 4 Aug 2026 00:37:39 -0700 Subject: [PATCH 038/222] Fix implicit `thread` address space qualifier becoming explicit in metal 4.1 (#3963) --- cmake/extension.cmake | 3 +- docs/src/dev/extensions.rst | 26 +++ mlx/backend/metal/kernels/CMakeLists.txt | 3 +- mlx/backend/metal/kernels/arg_reduce.metal | 8 +- mlx/backend/metal/kernels/atomic.h | 12 +- mlx/backend/metal/kernels/binary_ops.h | 84 ++++----- mlx/backend/metal/kernels/complex.h | 5 +- mlx/backend/metal/kernels/fft/readwrite.h | 79 ++++----- mlx/backend/metal/kernels/fp4.h | 8 +- mlx/backend/metal/kernels/fp8.h | 14 +- mlx/backend/metal/kernels/fp_quantized.h | 18 +- mlx/backend/metal/kernels/fp_quantized_nax.h | 18 +- mlx/backend/metal/kernels/gemv_masked.h | 4 +- mlx/backend/metal/kernels/logging.h | 2 +- mlx/backend/metal/kernels/quantized.h | 14 +- mlx/backend/metal/kernels/quantized_nax.h | 28 +-- mlx/backend/metal/kernels/reduction/ops.h | 75 ++++---- mlx/backend/metal/kernels/scan.h | 44 ++--- mlx/backend/metal/kernels/sort.h | 2 +- .../steel/attn/kernels/steel_attention_nax.h | 4 +- mlx/backend/metal/kernels/steel/attn/loader.h | 28 +-- mlx/backend/metal/kernels/steel/attn/mma.h | 58 ++++--- mlx/backend/metal/kernels/steel/attn/nax.h | 71 ++++---- .../metal/kernels/steel/attn/transforms.h | 8 +- .../steel/conv/loaders/loader_channel_l.h | 48 ++--- .../steel/conv/loaders/loader_channel_n.h | 16 +- .../steel/conv/loaders/loader_general.h | 20 +-- .../steel/gemm/kernels/steel_gemm_masked.h | 4 +- mlx/backend/metal/kernels/steel/gemm/loader.h | 14 +- mlx/backend/metal/kernels/steel/gemm/mma.h | 87 ++++++---- mlx/backend/metal/kernels/steel/gemm/nax.h | 71 ++++---- .../metal/kernels/steel/gemm/transforms.h | 8 +- .../kernels/steel/utils/integral_constant.h | 5 +- mlx/backend/metal/kernels/ternary_ops.h | 2 +- mlx/backend/metal/kernels/unary_ops.h | 164 +++++++++--------- mlx/backend/metal/kernels/utils.h | 27 +-- 36 files changed, 584 insertions(+), 498 deletions(-) diff --git a/cmake/extension.cmake b/cmake/extension.cmake index 13db804a1c..2f244a14f2 100644 --- a/cmake/extension.cmake +++ b/cmake/extension.cmake @@ -26,7 +26,8 @@ macro(mlx_build_metallib) set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib") # Collect compile options - set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions) + set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions + -Wmetal-addr-spaces) if(MLX_METAL_DEBUG OR MTLLIB_DEBUG) set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only -frecord-sources) diff --git a/docs/src/dev/extensions.rst b/docs/src/dev/extensions.rst index f98f096a0d..89c8991e56 100644 --- a/docs/src/dev/extensions.rst +++ b/docs/src/dev/extensions.rst @@ -330,6 +330,32 @@ GPU kernels in MLX are written using Metal. * Documentation for metal shading language: `Metal Specification`_ * Using metal from C++: `Metal-cpp`_ +.. note:: + + As of Metal 4.1, the implicit address space of ``this`` in a member function + is ``__metal_generic`` rather than ``thread``. Member functions in Metal + shading code should therefore carry an explicit address space qualifier, + written after the parameter list: + + .. code-block:: C++ + + struct Accumulator { + float vals[4]; + thread float& at(short i) thread { + return vals[i]; + } + }; + + Without the trailing ``thread``, returning a ``thread``-qualified reference + to a member no longer compiles:: + + error: reference to type 'thread float' could not bind to an lvalue + of type '__metal_generic float' + + Metal libraries built with ``mlx_build_metallib`` are compiled with + ``-Wmetal-addr-spaces``, which reports a missing qualifier as a warning on + toolchains where it is not yet an error. + Let's keep the GPU kernel simple. We will launch exactly as many threads as there are elements in the output. Each thread will pick the element it needs from ``x`` and ``y``, do the point-wise operation, and update its assigned diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 2bd99d2375..6d9a0883f0 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -17,7 +17,8 @@ function(build_kernel_base TARGET SRCFILE DEPS) -Wextra -fno-fast-math -Wno-c++17-extensions - -Wno-c++20-extensions) + -Wno-c++20-extensions + -Wmetal-addr-spaces) if(MLX_METAL_DEBUG) set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources) endif() diff --git a/mlx/backend/metal/kernels/arg_reduce.metal b/mlx/backend/metal/kernels/arg_reduce.metal index 4a83d8e57e..6ce20555f7 100644 --- a/mlx/backend/metal/kernels/arg_reduce.metal +++ b/mlx/backend/metal/kernels/arg_reduce.metal @@ -16,7 +16,7 @@ template struct ArgMin { static constexpr constant U init = Limits::max; - IndexValPair reduce(IndexValPair best, IndexValPair current) { + IndexValPair reduce(IndexValPair best, IndexValPair current) thread { if (best.val > current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -27,7 +27,7 @@ struct ArgMin { template IndexValPair - reduce_many(IndexValPair best, thread U* vals, uint32_t offset) { + reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { if (vals[i] < best.val) { best.val = vals[i]; @@ -42,7 +42,7 @@ template struct ArgMax { static constexpr constant U init = Limits::min; - IndexValPair reduce(IndexValPair best, IndexValPair current) { + IndexValPair reduce(IndexValPair best, IndexValPair current) thread { if (best.val < current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -53,7 +53,7 @@ struct ArgMax { template IndexValPair - reduce_many(IndexValPair best, thread U* vals, uint32_t offset) { + reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { if (vals[i] > best.val) { best.val = vals[i]; diff --git a/mlx/backend/metal/kernels/atomic.h b/mlx/backend/metal/kernels/atomic.h index 76362f58ae..3835fff024 100644 --- a/mlx/backend/metal/kernels/atomic.h +++ b/mlx/backend/metal/kernels/atomic.h @@ -172,7 +172,7 @@ union uint_or_packed { template struct mlx_atomic_update_helper { - uint operator()(uint_or_packed init, T update, size_t elem_offset) { + uint operator()(uint_or_packed init, T update, size_t elem_offset) thread { Op op; init.val[elem_offset] = op(update, init.val[elem_offset]); return init.bits; @@ -209,7 +209,7 @@ struct __None { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { #pragma unused(b) return a; } @@ -223,7 +223,7 @@ struct __Add { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a + b; } }; @@ -235,7 +235,7 @@ struct __Mul { return b != 0; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a * b; } }; @@ -246,7 +246,7 @@ struct __Max { return a > b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return max(a, b); } }; @@ -257,7 +257,7 @@ struct __Min { return a < b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return min(a, b); } }; diff --git a/mlx/backend/metal/kernels/binary_ops.h b/mlx/backend/metal/kernels/binary_ops.h index 4e3d881fe0..863d6369e2 100644 --- a/mlx/backend/metal/kernels/binary_ops.h +++ b/mlx/backend/metal/kernels/binary_ops.h @@ -9,33 +9,33 @@ constant mlx::os_log logger("mlx", "binary_ops"); struct Add { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x + y; } }; struct FloorDivide { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x / y; } template <> - float operator()(float x, float y) { + float operator()(float x, float y) thread { return trunc(x / y); } template <> - half operator()(half x, half y) { + half operator()(half x, half y) thread { return trunc(x / y); } template <> - bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { + bfloat16_t operator()(bfloat16_t x, bfloat16_t y) thread { return trunc(x / y); } }; struct Divide { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x / y; } }; @@ -43,12 +43,12 @@ struct Divide { struct Remainder { template metal::enable_if_t & !metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { return x % y; } template metal::enable_if_t & metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { auto r = x % y; if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -56,7 +56,7 @@ struct Remainder { return r; } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { T r = fmod(x, y); if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -64,25 +64,25 @@ struct Remainder { return r; } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { return x % y; } }; struct Equal { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y; } }; struct NaNEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y || (metal::isnan(x) && metal::isnan(y)); } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x == y || (metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) && metal::isnan(y.imag)) || @@ -93,35 +93,35 @@ struct NaNEqual { struct Greater { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x > y; } }; struct GreaterEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x >= y; } }; struct Less { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x < y; } }; struct LessEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x <= y; } }; struct LogAddExp { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { if (metal::isnan(x) || metal::isnan(y)) { return metal::numeric_limits::quiet_NaN(); } @@ -133,7 +133,7 @@ struct LogAddExp { : (maxval + log1p(metal::exp(minval - maxval))); }; - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) || metal::isnan(y.imag)) { return metal::numeric_limits::quiet_NaN(); @@ -154,12 +154,12 @@ struct LogAddExp { struct Maximum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::max(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -167,7 +167,7 @@ struct Maximum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -177,12 +177,12 @@ struct Maximum { struct Minimum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::min(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -190,7 +190,7 @@ struct Minimum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -200,30 +200,32 @@ struct Minimum { struct Multiply { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x * y; } }; struct NotEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x != y; } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x.real != y.real || x.imag != y.imag; } }; struct Power { template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { return metal::pow(base, exp); } template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { T res = 1; // Undefined to raise integer to negative power if (exp < 0) { @@ -243,7 +245,7 @@ struct Power { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (x.real == 0 && x.imag == 0) { if (metal::isnan(y.real) || metal::isnan(y.imag)) { auto nan = metal::numeric_limits::quiet_NaN(); @@ -261,70 +263,70 @@ struct Power { struct Subtract { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x - y; } }; struct LogicalAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x && y; }; }; struct LogicalOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x || y; }; }; struct BitwiseAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x & y; }; }; struct BitwiseOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x | y; }; }; struct BitwiseXor { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x ^ y; }; }; struct LeftShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x << y; }; }; struct RightShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x >> y; }; }; struct ArcTan2 { template - T operator()(T y, T x) { + T operator()(T y, T x) thread { return metal::precise::atan2(y, x); } }; struct DivMod { template - metal::array operator()(T x, T y) { + metal::array operator()(T x, T y) thread { return {FloorDivide{}(x, y), Remainder{}(x, y)}; }; }; diff --git a/mlx/backend/metal/kernels/complex.h b/mlx/backend/metal/kernels/complex.h index 6e391483d3..5f654332b3 100644 --- a/mlx/backend/metal/kernels/complex.h +++ b/mlx/backend/metal/kernels/complex.h @@ -22,8 +22,9 @@ struct complex64_t { float imag; // Constructors - constexpr complex64_t(float real, float imag) : real(real), imag(imag) {}; - constexpr complex64_t() : real(0), imag(0) {}; + constexpr complex64_t(float real, float imag) thread : real(real), + imag(imag) {}; + constexpr complex64_t() thread : real(0), imag(0) {}; constexpr complex64_t() threadgroup : real(0), imag(0) {}; // Conversions to complex64_t diff --git a/mlx/backend/metal/kernels/fft/readwrite.h b/mlx/backend/metal/kernels/fft/readwrite.h index 0dc62992e6..3d1b23f4ed 100644 --- a/mlx/backend/metal/kernels/fft/readwrite.h +++ b/mlx/backend/metal/kernels/fft/readwrite.h @@ -57,16 +57,15 @@ struct ReadWriter { const short elems_per_thread_, const uint3 elem_, const uint3 grid_, - const bool inv_) - : in(in_), - buf(buf_), - out(out_), - n(n_), - batch_size(batch_size_), - elems_per_thread(elems_per_thread_), - elem(elem_), - grid(grid_), - inv(inv_) { + const bool inv_) thread : in(in_), + buf(buf_), + out(out_), + n(n_), + batch_size(batch_size_), + elems_per_thread(elems_per_thread_), + elem(elem_), + grid(grid_), + inv(inv_) { // Account for padding on last threadgroup threads_per_tg = elem.x == grid.x - 1 ? (batch_size - (grid.x - 1) * grid.y) * grid.z @@ -74,30 +73,30 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const { + METAL_FUNC float2 post_in(float2 elem) const thread { return inv ? float2(elem.x, -elem.y) : elem; } // Handle float case for generic RFFT alg - METAL_FUNC float2 post_in(float elem) const { + METAL_FUNC float2 post_in(float elem) const thread { return float2(elem, 0); } - METAL_FUNC float2 pre_out(float2 elem) const { + METAL_FUNC float2 pre_out(float2 elem) const thread { return inv ? float2(elem.x / n, -elem.y / n) : elem; } - METAL_FUNC float2 pre_out(float2 elem, int length) const { + METAL_FUNC float2 pre_out(float2 elem, int length) const thread { return inv ? float2(elem.x / length, -elem.y / length) : elem; } - METAL_FUNC bool out_of_bounds() const { + METAL_FUNC bool out_of_bounds() const thread { // Account for possible extra threadgroups int grid_index = elem.x * grid.y + elem.y; return grid_index >= batch_size; } - METAL_FUNC void load() const { + METAL_FUNC void load() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; short max_index = grid.y * n - 2; @@ -120,7 +119,7 @@ struct ReadWriter { } } - METAL_FUNC void write() const { + METAL_FUNC void write() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; short max_index = grid.y * n - 2; @@ -143,7 +142,8 @@ struct ReadWriter { } // Padded IO for Bluestein's algorithm - METAL_FUNC void load_padded(int length, const device float2* w_k) const { + METAL_FUNC void load_padded(int length, const device float2* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; @@ -160,7 +160,8 @@ struct ReadWriter { } } - METAL_FUNC void write_padded(int length, const device float2* w_k) const { + METAL_FUNC void write_padded(int length, const device float2* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; @@ -177,7 +178,7 @@ struct ReadWriter { } // Strided IO for four step FFT - METAL_FUNC void compute_strided_indices(int stride, int overall_n) { + METAL_FUNC void compute_strided_indices(int stride, int overall_n) thread { // Use the batch threadgroup dimension to coalesce memory accesses: // e.g. stride = 12 // device | shared mem @@ -199,7 +200,7 @@ struct ReadWriter { } // Four Step FFT First Step - METAL_FUNC void load_strided(int stride, int overall_n) { + METAL_FUNC void load_strided(int stride, int overall_n) thread { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { buf[strided_shared_idx + e] = @@ -207,7 +208,7 @@ struct ReadWriter { } } - METAL_FUNC void write_strided(int stride, int overall_n) { + METAL_FUNC void write_strided(int stride, int overall_n) thread { for (int e = 0; e < elems_per_thread; e++) { float2 output = buf[strided_shared_idx + e]; int combined_idx = (strided_device_idx + e * stride) % overall_n; @@ -223,7 +224,7 @@ struct ReadWriter { template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -237,7 +238,7 @@ METAL_FUNC void ReadWriter::load_strided( template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { float2 output = buf[strided_shared_idx + e]; @@ -253,14 +254,14 @@ METAL_FUNC void ReadWriter::write_strided( // // This roughly doubles the throughput over the regular FFT. template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for RFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -280,7 +281,7 @@ METAL_FUNC void ReadWriter::load() const { } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = @@ -317,7 +318,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -344,7 +345,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int length_over_2 = (length / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * length_over_2 * 2 + elem.y * length_over_2 * 2; @@ -389,14 +390,14 @@ METAL_FUNC void ReadWriter::write_padded( // x_k = Re(Z_k) // Y_k = Imag(Z_k) template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for IRFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * n_over_2 * 2 + elem.y * n_over_2 * 2; @@ -435,7 +436,7 @@ METAL_FUNC void ReadWriter::load() const { } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { int batch_idx = elem.x * grid.y * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -456,7 +457,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int n_over_2 = (n / 2) + 1; int length_over_2 = (length / 2) + 1; @@ -504,7 +505,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n + length - 1; @@ -531,7 +532,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -546,7 +547,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; int coalesce_width = grid.y; int tg_idx = elem.y * grid.z + elem.z; @@ -575,7 +576,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; auto conj = float2(1, -1); @@ -600,7 +601,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -614,7 +615,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { diff --git a/mlx/backend/metal/kernels/fp4.h b/mlx/backend/metal/kernels/fp4.h index 25642f2016..47bf4dda6f 100644 --- a/mlx/backend/metal/kernels/fp4.h +++ b/mlx/backend/metal/kernels/fp4.h @@ -1,7 +1,7 @@ #pragma once struct fp4_e2m1 { - fp4_e2m1(float x) { + fp4_e2m1(float x) thread { if (metal::isnan(x)) { bits = 0x7; return; @@ -30,17 +30,17 @@ struct fp4_e2m1 { bits |= sign_bit; } - operator float16_t() { + operator float16_t() thread { half converted = as_type(ushort((bits & 7) << 9)); converted *= 16384.0; return bits & 8 ? -converted : converted; } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 60d34be694..796dd21639 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -2,7 +2,7 @@ struct fp8_e4m3 { template - fp8_e4m3(T f) { + fp8_e4m3(T f) thread { // From PyTorch // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148 uint32_t fp8_max = 543 << 21; @@ -29,7 +29,7 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } - operator float16_t() { + operator float16_t() thread { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; @@ -37,11 +37,11 @@ struct fp8_e4m3 { return (sign ? -converted : converted); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } @@ -49,7 +49,7 @@ struct fp8_e4m3 { }; struct fp8_e8m0 { - fp8_e8m0(float x) { + fp8_e8m0(float x) thread { if (!metal::isfinite(x)) { bits = 0xFF; return; @@ -66,12 +66,12 @@ struct fp8_e8m0 { bits = static_cast(n + 127); } - operator bfloat16_t() { + operator bfloat16_t() thread { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } - operator float() { + operator float() thread { uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); return as_type(out); } diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 677183fa92..85d4284ac3 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -39,7 +39,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -50,7 +50,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -186,15 +186,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -203,7 +203,7 @@ struct QuantizedBlockLoader { scales_ + bi * src_ld / group_size + (bj * pack_factor) / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -215,7 +215,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -241,7 +241,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 0ea6af168e..7c452e6a64 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -39,7 +39,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -50,7 +50,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -112,14 +112,14 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -127,7 +127,7 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -143,7 +143,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -173,7 +173,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { scales += n_groups; diff --git a/mlx/backend/metal/kernels/gemv_masked.h b/mlx/backend/metal/kernels/gemv_masked.h index 407d14bb4b..58ab744050 100644 --- a/mlx/backend/metal/kernels/gemv_masked.h +++ b/mlx/backend/metal/kernels/gemv_masked.h @@ -10,7 +10,7 @@ using namespace metal; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -30,7 +30,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/mlx/backend/metal/kernels/logging.h b/mlx/backend/metal/kernels/logging.h index 7b3ee04674..d7446c7222 100644 --- a/mlx/backend/metal/kernels/logging.h +++ b/mlx/backend/metal/kernels/logging.h @@ -16,7 +16,7 @@ struct os_log { constexpr os_log(constant char*, constant char*) constant {} template - void log_debug(constant char*, Args...) const {} + void log_debug(constant char*, Args...) const thread {} template void log_debug(constant char*, Args...) const constant {} diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 345dfc711e..1820ce3423 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -610,15 +610,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -626,7 +626,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -639,7 +639,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -669,7 +669,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index d50af8679a..e67be9a06d 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -612,15 +612,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -628,7 +628,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -641,7 +641,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -671,7 +671,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -752,14 +752,14 @@ struct QuantizedBlockLoader< const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -768,7 +768,7 @@ struct QuantizedBlockLoader< scales(scales_ + bi * src_ld / group_size + group_id), biases(biases_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -781,7 +781,7 @@ struct QuantizedBlockLoader< } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -811,7 +811,7 @@ struct QuantizedBlockLoader< } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { // if (group_steps > 1) { diff --git a/mlx/backend/metal/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h index 11d8e83ac6..b7a9cacb39 100644 --- a/mlx/backend/metal/kernels/reduction/ops.h +++ b/mlx/backend/metal/kernels/reduction/ops.h @@ -7,12 +7,12 @@ #define DEFINE_SIMD_REDUCE() \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ return simd_reduce_impl(val); \ } \ \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ for (short i = simd_size / 2; i > 0; i /= 2) { \ val = operator()(val, simd_shuffle_down(val, i)); \ } \ @@ -28,7 +28,8 @@ union bool4_or_uint { struct None { template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_store_explicit(out, val, offset); } }; @@ -37,7 +38,7 @@ template struct And { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_all(val); } @@ -47,7 +48,7 @@ struct And { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (!val) { bool4_or_uint update; update.b = {true, true, true, true}; @@ -56,20 +57,20 @@ struct And { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (!val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out &= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a && b; } }; @@ -78,7 +79,7 @@ template struct Or { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_any(val); } @@ -88,7 +89,7 @@ struct Or { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (val) { bool4_or_uint update; update.b = {false, false, false, false}; @@ -97,20 +98,20 @@ struct Or { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out |= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a || b; } }; @@ -120,19 +121,20 @@ struct Sum { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_sum(val); } static constexpr constant U init = U(0); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_add_explicit(out, val, offset); } // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a + b; } }; @@ -142,19 +144,20 @@ struct Prod { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_product(val); } static constexpr constant U init = U(1); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_mul_explicit(out, val, offset); } // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a * b; } }; @@ -164,12 +167,14 @@ struct Min { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_min(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -179,18 +184,19 @@ struct Min { static constexpr constant U init = Limits::max; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_min_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a < b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -199,7 +205,7 @@ struct Min { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); @@ -221,12 +227,14 @@ struct Max { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_max(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -236,18 +244,19 @@ struct Max { static constexpr constant U init = Limits::min; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_max_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a > b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -256,7 +265,7 @@ struct Max { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); diff --git a/mlx/backend/metal/kernels/scan.h b/mlx/backend/metal/kernels/scan.h index 16682613b9..a6bfde0018 100644 --- a/mlx/backend/metal/kernels/scan.h +++ b/mlx/backend/metal/kernels/scan.h @@ -6,12 +6,12 @@ #define DEFINE_SIMD_SCAN() \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ return simd_scan_impl(val); \ } \ \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ for (int i = 1; i <= 16; i *= 2) { \ val = operator()(val, simd_shuffle_and_fill_up(val, init, i)); \ } \ @@ -20,12 +20,12 @@ #define DEFINE_SIMD_EXCLUSIVE_SCAN() \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ return simd_exclusive_scan_impl(val); \ } \ \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ val = simd_scan(val); \ return simd_shuffle_and_fill_up(val, init, 1); \ } @@ -38,15 +38,15 @@ struct CumSum { static constexpr constant U init = static_cast(0); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a + b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_sum(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_sum(x); } }; @@ -59,15 +59,15 @@ struct CumProd { static constexpr constant U init = static_cast(1.0f); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a * b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_product(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_product(x); } }; @@ -77,11 +77,11 @@ struct CumProd { static constexpr constant bool init = true; template - bool operator()(bool a, T b) { + bool operator()(bool a, T b) thread { return a & static_cast(b); } - bool simd_scan(bool x) { + bool simd_scan(bool x) thread { for (int i = 1; i <= 16; i *= 2) { bool other = simd_shuffle_and_fill_up(x, init, i); x &= other; @@ -89,7 +89,7 @@ struct CumProd { return x; } - bool simd_exclusive_scan(bool x) { + bool simd_exclusive_scan(bool x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -100,11 +100,11 @@ struct CumMax { static constexpr constant U init = Limits::min; template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return (a >= b) ? a : b; } - U simd_scan(U x) { + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); x = (x >= other) ? x : other; @@ -112,7 +112,7 @@ struct CumMax { return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -123,11 +123,11 @@ struct CumMin { static constexpr constant U init = Limits::max; template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return (a <= b) ? a : b; } - U simd_scan(U x) { + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); x = (x <= other) ? x : other; @@ -135,7 +135,7 @@ struct CumMin { return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -146,11 +146,11 @@ struct CumLogaddexp { static constexpr constant U init = Limits::min; template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return LogAddExp{}(a, static_cast(b)); } - U simd_scan(U x) { + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); x = LogAddExp{}(x, other); @@ -158,7 +158,7 @@ struct CumLogaddexp { return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } diff --git a/mlx/backend/metal/kernels/sort.h b/mlx/backend/metal/kernels/sort.h index 7d7fab0f9a..068d43d126 100644 --- a/mlx/backend/metal/kernels/sort.h +++ b/mlx/backend/metal/kernels/sort.h @@ -39,7 +39,7 @@ struct Init { template struct LessThan { static constexpr constant T init = Init::v; - METAL_FUNC bool operator()(T a, T b) const { + METAL_FUNC bool operator()(T a, T b) const thread { if constexpr (metal::is_floating_point_v) { bool an = metal::isnan(a); bool bn = metal::isnan(b); diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index 6c85da64f6..b48a9a942d 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -21,9 +21,9 @@ constant bool has_sinks [[function_constant(302)]]; template struct TransformScale { T scale; - METAL_FUNC TransformScale(T scale_) : scale(scale_) {} + METAL_FUNC TransformScale(T scale_) thread : scale(scale_) {} - METAL_FUNC T apply(T x) const { + METAL_FUNC T apply(T x) const thread { return scale * x; } }; diff --git a/mlx/backend/metal/kernels/steel/attn/loader.h b/mlx/backend/metal/kernels/steel/attn/loader.h index 7ec798146b..a2bdb0e46b 100644 --- a/mlx/backend/metal/kernels/steel/attn/loader.h +++ b/mlx/backend/metal/kernels/steel/attn/loader.h @@ -49,18 +49,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -71,7 +71,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -80,7 +80,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -128,7 +128,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; @@ -173,18 +173,18 @@ struct BlockLoaderT { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * kDstStrRow + bj * kDstStrCol), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -196,7 +196,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -207,7 +207,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -255,7 +255,7 @@ struct BlockLoaderT { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; diff --git a/mlx/backend/metal/kernels/steel/attn/mma.h b/mlx/backend/metal/kernels/steel/attn/mma.h index 737e930d8e..f5b775cfb1 100644 --- a/mlx/backend/metal/kernels/steel/attn/mma.h +++ b/mlx/backend/metal/kernels/steel/attn/mma.h @@ -24,7 +24,7 @@ struct Shape2D { RInt r; CInt c; - Shape2D(RInt r_, CInt c_) : r(r_), c(c_) {} + Shape2D(RInt r_, CInt c_) thread : r(r_), c(c_) {} }; template @@ -257,24 +257,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -283,16 +284,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const { + METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -304,7 +305,7 @@ struct MMATile { } template - METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) { + METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -316,7 +317,7 @@ struct MMATile { } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -333,7 +334,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -350,7 +351,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -365,7 +366,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -380,8 +381,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -400,8 +403,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -512,7 +517,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -530,7 +535,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -557,7 +562,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -571,7 +576,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -590,7 +595,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -604,7 +610,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -633,7 +639,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -678,7 +684,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -711,7 +717,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/mlx/backend/metal/kernels/steel/attn/nax.h b/mlx/backend/metal/kernels/steel/attn/nax.h index 072afe3ba3..6d978ebfb6 100644 --- a/mlx/backend/metal/kernels/steel/attn/nax.h +++ b/mlx/backend/metal/kernels/steel/attn/nax.h @@ -420,8 +420,8 @@ struct BaseNAXFrag { // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), + metal::remove_addrspace_t, + metal::remove_addrspace_t, CType>(); // Load A in to left operand registers @@ -492,8 +492,8 @@ struct BaseNAXFrag { // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), + metal::remove_addrspace_t, + metal::remove_addrspace_t, CType>(); // Load A in to left operand registers @@ -563,36 +563,39 @@ struct NAXTile { METAL_FUNC NAXTile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& frag_at() { + METAL_FUNC constexpr thread frag_type& frag_at() thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& - frag_at(const short i, const short j, metal::bool_constant) { + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { if constexpr (transpose) { return frag_at(j, i); } else { @@ -601,8 +604,10 @@ struct NAXTile { } template - METAL_FUNC constexpr const thread frag_type& - frag_at(const short i, const short j, metal::bool_constant) const { + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { if constexpr (transpose) { return frag_at(j, i); } else { @@ -611,7 +616,7 @@ struct NAXTile { } template - METAL_FUNC constexpr thread frag_type& frag_at() { + METAL_FUNC constexpr thread frag_type& frag_at() thread { if constexpr (transpose) { return frag_at(); } else { @@ -620,7 +625,7 @@ struct NAXTile { } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { if constexpr (transpose) { return frag_at(); } else { @@ -628,16 +633,17 @@ struct NAXTile { } } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { @@ -650,7 +656,8 @@ struct NAXTile { } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { @@ -663,7 +670,7 @@ struct NAXTile { } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( @@ -678,7 +685,7 @@ struct NAXTile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store( @@ -693,7 +700,7 @@ struct NAXTile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( @@ -708,7 +715,7 @@ struct NAXTile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store( @@ -724,7 +731,7 @@ struct NAXTile { template METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { + load_rows(const device U* src, const int ld, const short n_rows) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_rows( @@ -740,8 +747,10 @@ struct NAXTile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_safe( @@ -759,7 +768,7 @@ struct NAXTile { template METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { + const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_rows( @@ -775,8 +784,10 @@ struct NAXTile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_safe( @@ -797,7 +808,7 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_slice( diff --git a/mlx/backend/metal/kernels/steel/attn/transforms.h b/mlx/backend/metal/kernels/steel/attn/transforms.h index c0624d21b9..6878448bc0 100644 --- a/mlx/backend/metal/kernels/steel/attn/transforms.h +++ b/mlx/backend/metal/kernels/steel/attn/transforms.h @@ -24,7 +24,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -40,14 +40,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast(x * alpha + (beta * c)); } }; diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h index 9124e30459..15dc13cccf 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h @@ -64,10 +64,10 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -103,7 +103,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -131,7 +131,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -213,10 +213,10 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -287,7 +287,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -312,7 +312,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -394,11 +394,11 @@ struct Conv2DWeightBlockLoader { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -408,7 +408,7 @@ struct Conv2DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -435,7 +435,7 @@ struct Conv2DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_hw < (params->wS[1] * params->wS[0])) { src += weight_step; return; @@ -504,10 +504,10 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -559,7 +559,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -588,7 +588,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -683,10 +683,10 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -777,7 +777,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t d_mask = mask_t(1) << weight_d; mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -804,7 +804,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -897,11 +897,11 @@ struct Conv3DWeightBlockLoader { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -911,7 +911,7 @@ struct Conv3DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -938,7 +938,7 @@ struct Conv3DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_dhw < (params->wS[0] * params->wS[1] * params->wS[2])) { src += weight_step; return; diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h index 2312e1ca6e..6a6c16daf0 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h @@ -99,10 +99,10 @@ struct Conv2DInputBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -131,7 +131,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { @@ -187,7 +187,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; @@ -243,11 +243,11 @@ struct Conv2DWeightBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld), params(params_), @@ -256,7 +256,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (bi >= BROWS || bj >= BCOLS) return; @@ -310,7 +310,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h index 9b7ddc2eec..cf7794235d 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h @@ -67,10 +67,10 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), jump_params(jump_params_), @@ -101,7 +101,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -137,7 +137,7 @@ struct Conv2DInputBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -184,7 +184,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; @@ -263,11 +263,11 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -279,7 +279,7 @@ struct Conv2DWeightBlockLoaderGeneral { start_row(offsets.y + bi) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -308,7 +308,7 @@ struct Conv2DWeightBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -358,7 +358,7 @@ struct Conv2DWeightBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h index cc3ddd930b..bdf92cf485 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h @@ -11,7 +11,7 @@ using namespace mlx::steel; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -29,7 +29,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/mlx/backend/metal/kernels/steel/gemm/loader.h b/mlx/backend/metal/kernels/steel/gemm/loader.h index d421b2d1f9..c999fdc91a 100644 --- a/mlx/backend/metal/kernels/steel/gemm/loader.h +++ b/mlx/backend/metal/kernels/steel/gemm/loader.h @@ -49,18 +49,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -71,7 +71,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -80,7 +80,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -128,7 +128,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; diff --git a/mlx/backend/metal/kernels/steel/gemm/mma.h b/mlx/backend/metal/kernels/steel/gemm/mma.h index 809b38d0af..1d9e0a7820 100644 --- a/mlx/backend/metal/kernels/steel/gemm/mma.h +++ b/mlx/backend/metal/kernels/steel/gemm/mma.h @@ -234,24 +234,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -260,16 +261,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -286,7 +287,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -303,7 +304,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -318,7 +319,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -333,8 +334,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -353,8 +356,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -377,7 +382,7 @@ struct MMATile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -487,7 +492,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -505,7 +510,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -532,7 +537,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -545,8 +550,11 @@ struct BlockMMA { Ctile.template store(D, ldd); } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -566,7 +574,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -585,7 +593,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -599,7 +608,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -628,7 +637,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -673,7 +682,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -706,7 +715,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -819,7 +828,7 @@ struct BlockMMA< /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -839,7 +848,7 @@ struct BlockMMA< /* Karatsuba MMA: 3 real MMAs per K-chunk */ METAL_FUNC void mma( const threadgroup complex64_t* As, - const threadgroup complex64_t* Bs) { + const threadgroup complex64_t* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -897,7 +906,7 @@ struct BlockMMA< } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Adjust for simdgroup and thread location D += sm * ldd + sn; @@ -916,8 +925,11 @@ struct BlockMMA< } } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { D += sm * ldd + sn; start -= short2(sn, sm); stop -= short2(sn, sm); @@ -948,7 +960,7 @@ struct BlockMMA< } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { D += sm * ldd + sn; dst_tile_dims -= short2(sn, sm); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) @@ -974,7 +986,8 @@ struct BlockMMA< /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) { complex64_t out = epilogue_op.apply( @@ -990,7 +1003,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -1022,7 +1035,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -1071,7 +1084,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1106,7 +1119,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kernels/steel/gemm/nax.h index 072afe3ba3..6d978ebfb6 100644 --- a/mlx/backend/metal/kernels/steel/gemm/nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/nax.h @@ -420,8 +420,8 @@ struct BaseNAXFrag { // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), + metal::remove_addrspace_t, + metal::remove_addrspace_t, CType>(); // Load A in to left operand registers @@ -492,8 +492,8 @@ struct BaseNAXFrag { // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), + metal::remove_addrspace_t, + metal::remove_addrspace_t, CType>(); // Load A in to left operand registers @@ -563,36 +563,39 @@ struct NAXTile { METAL_FUNC NAXTile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& frag_at() { + METAL_FUNC constexpr thread frag_type& frag_at() thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& - frag_at(const short i, const short j, metal::bool_constant) { + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { if constexpr (transpose) { return frag_at(j, i); } else { @@ -601,8 +604,10 @@ struct NAXTile { } template - METAL_FUNC constexpr const thread frag_type& - frag_at(const short i, const short j, metal::bool_constant) const { + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { if constexpr (transpose) { return frag_at(j, i); } else { @@ -611,7 +616,7 @@ struct NAXTile { } template - METAL_FUNC constexpr thread frag_type& frag_at() { + METAL_FUNC constexpr thread frag_type& frag_at() thread { if constexpr (transpose) { return frag_at(); } else { @@ -620,7 +625,7 @@ struct NAXTile { } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { if constexpr (transpose) { return frag_at(); } else { @@ -628,16 +633,17 @@ struct NAXTile { } } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { @@ -650,7 +656,8 @@ struct NAXTile { } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { @@ -663,7 +670,7 @@ struct NAXTile { } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( @@ -678,7 +685,7 @@ struct NAXTile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store( @@ -693,7 +700,7 @@ struct NAXTile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( @@ -708,7 +715,7 @@ struct NAXTile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store( @@ -724,7 +731,7 @@ struct NAXTile { template METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { + load_rows(const device U* src, const int ld, const short n_rows) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_rows( @@ -740,8 +747,10 @@ struct NAXTile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_safe( @@ -759,7 +768,7 @@ struct NAXTile { template METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { + const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_rows( @@ -775,8 +784,10 @@ struct NAXTile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_safe( @@ -797,7 +808,7 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store_slice( diff --git a/mlx/backend/metal/kernels/steel/gemm/transforms.h b/mlx/backend/metal/kernels/steel/gemm/transforms.h index 0282a1223f..14e94e480a 100644 --- a/mlx/backend/metal/kernels/steel/gemm/transforms.h +++ b/mlx/backend/metal/kernels/steel/gemm/transforms.h @@ -24,7 +24,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -40,14 +40,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast( x * static_cast(alpha) + (static_cast(beta) * c)); } diff --git a/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/mlx/backend/metal/kernels/steel/utils/integral_constant.h index fa7b39986f..2f153f482f 100644 --- a/mlx/backend/metal/kernels/steel/utils/integral_constant.h +++ b/mlx/backend/metal/kernels/steel/utils/integral_constant.h @@ -20,7 +20,7 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } }; @@ -52,7 +52,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); diff --git a/mlx/backend/metal/kernels/ternary_ops.h b/mlx/backend/metal/kernels/ternary_ops.h index e0235d9dd3..d97e238c69 100644 --- a/mlx/backend/metal/kernels/ternary_ops.h +++ b/mlx/backend/metal/kernels/ternary_ops.h @@ -4,7 +4,7 @@ struct Select { template - T operator()(bool condition, T x, T y) { + T operator()(bool condition, T x, T y) thread { return condition ? x : y; } }; diff --git a/mlx/backend/metal/kernels/unary_ops.h b/mlx/backend/metal/kernels/unary_ops.h index 327bb5a940..1337f6b52e 100644 --- a/mlx/backend/metal/kernels/unary_ops.h +++ b/mlx/backend/metal/kernels/unary_ops.h @@ -16,125 +16,125 @@ constant float inf = metal::numeric_limits::infinity(); struct Abs { template - T operator()(T x) { + T operator()(T x) thread { return metal::abs(x); }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; }; }; struct ArcCos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acos(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcCosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acosh(x); }; }; struct ArcSin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asin(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcSinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asinh(x); }; }; struct ArcTan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atan(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcTanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atanh(x); }; }; struct BitwiseInvert { template - T operator()(T x) { + T operator()(T x) thread { return ~x; }; }; struct Ceil { template - T operator()(T x) { + T operator()(T x) thread { return metal::ceil(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Cos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cos(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cos(x.real) * metal::precise::cosh(x.imag), -metal::precise::sin(x.real) * metal::precise::sinh(x.imag)}; @@ -143,11 +143,11 @@ struct Cos { struct Cosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cosh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cosh(x.real) * metal::precise::cos(x.imag), metal::precise::sinh(x.real) * metal::precise::sin(x.imag)}; @@ -155,89 +155,89 @@ struct Cosh { }; struct Conjugate { - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return complex64_t{x.real, -x.imag}; } }; struct Erf { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erf(static_cast(x))); }; }; struct ErfInv { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erfinv(static_cast(x))); }; }; struct Exp { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::exp(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return cexpf(x); } }; struct Expm1 { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(expm1f(static_cast(x))); }; }; struct Floor { template - T operator()(T x) { + T operator()(T x) thread { return metal::floor(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Imag { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.imag; }; }; struct Log { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto r = metal::precise::log(Abs{}(x).real); auto i = metal::precise::atan2(x.imag, x.real); return {r, i}; @@ -246,11 +246,11 @@ struct Log { struct Log2 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log2(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN2_F, y.imag / M_LN2_F}; }; @@ -258,11 +258,11 @@ struct Log2 { struct Log10 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log10(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN10_F, y.imag / M_LN10_F}; }; @@ -270,44 +270,44 @@ struct Log10 { struct Log1p { template - T operator()(T x) { + T operator()(T x) thread { return log1p(x); }; }; struct LogicalNot { template - T operator()(T x) { + T operator()(T x) thread { return !x; }; }; struct Negative { template - T operator()(T x) { + T operator()(T x) thread { return -x; }; }; struct Real { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.real; }; }; struct Round { template - T operator()(T x) { + T operator()(T x) thread { return metal::rint(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::rint(x.real), metal::rint(x.imag)}; }; }; struct Sigmoid { template - T operator()(T x) { + T operator()(T x) thread { auto y = 1 / (1 + metal::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } @@ -315,13 +315,13 @@ struct Sigmoid { struct Sign { template - T operator()(T x) { + T operator()(T x) thread { return (x > T(0)) - (x < T(0)); }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x != 0; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x == complex64_t(0)) { return x; } @@ -332,11 +332,11 @@ struct Sign { struct Sin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sin(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sin(x.real) * metal::precise::cosh(x.imag), metal::precise::cos(x.real) * metal::precise::sinh(x.imag)}; @@ -345,11 +345,11 @@ struct Sin { struct Sinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sinh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sinh(x.real) * metal::precise::cos(x.imag), metal::precise::cosh(x.real) * metal::precise::sin(x.imag)}; @@ -358,18 +358,18 @@ struct Sinh { struct Square { template - T operator()(T x) { + T operator()(T x) thread { return x * x; }; }; struct Sqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x.real == 0.0 && x.imag == 0.0) { return {0.0, 0.0}; } @@ -383,22 +383,22 @@ struct Sqrt { struct Rsqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::rsqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return 1.0 / Sqrt{}(x); } }; struct Tan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tan(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tan_a = metal::precise::tan(x.real); float tanh_b = metal::precise::tanh(x.imag); float t1 = tan_a * tanh_b; @@ -409,11 +409,11 @@ struct Tan { struct Tanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tanh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tanh_a = metal::precise::tanh(x.real); float tan_b = metal::precise::tan(x.imag); float t1 = tanh_a * tan_b; @@ -422,19 +422,19 @@ struct Tanh { }; }; -complex64_t ArcCos::operator()(complex64_t x) { +complex64_t ArcCos::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(x + i * Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcSin::operator()(complex64_t x) { +complex64_t ArcSin::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(i * x + Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcTan::operator()(complex64_t x) { +complex64_t ArcTan::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto ix = i * x; return (1.0 / complex64_t{0.0, 2.0}) * Log{}((1.0 + ix) / (1.0 - ix)); @@ -442,13 +442,13 @@ complex64_t ArcTan::operator()(complex64_t x) { struct ToFP8 { template - uint8_t operator()(T f) { + uint8_t operator()(T f) thread { return fp8_e4m3(f).bits; } }; struct FromFP8 { - float operator()(uint8_t x) { + float operator()(uint8_t x) thread { return float(*(thread fp8_e4m3*)(&x)); } }; diff --git a/mlx/backend/metal/kernels/utils.h b/mlx/backend/metal/kernels/utils.h index d356450101..f2a8362ae2 100644 --- a/mlx/backend/metal/kernels/utils.h +++ b/mlx/backend/metal/kernels/utils.h @@ -205,9 +205,9 @@ struct LoopedElemToLoc { OffsetT offset{0}; int index{0}; - LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {} + LoopedElemToLoc(int dim) thread : dim(dim), inner_looper(dim - 1) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { if (dim == 0) { return; } @@ -220,7 +220,8 @@ struct LoopedElemToLoc { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { if (dim == 0) { return; } @@ -243,7 +244,7 @@ struct LoopedElemToLoc { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -254,9 +255,9 @@ struct LoopedElemToLoc<1, OffsetT, true> { OffsetT offset{0}; uint index{0}; - LoopedElemToLoc(int dim) : dim(dim) {} + LoopedElemToLoc(int dim) thread : dim(dim) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { index++; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -265,7 +266,8 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { index += n; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -274,7 +276,7 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -283,17 +285,18 @@ template struct LoopedElemToLoc<1, OffsetT, false> { OffsetT offset{0}; - LoopedElemToLoc(int) {} + LoopedElemToLoc(int) thread {} - void next(const constant int*, const constant int64_t* strides) { + void next(const constant int*, const constant int64_t* strides) thread { offset += OffsetT(strides[0]); } - void next(int n, const constant int*, const constant int64_t* strides) { + void next(int n, const constant int*, const constant int64_t* strides) + thread { offset += n * OffsetT(strides[0]); } - OffsetT location() { + OffsetT location() thread { return offset; } }; From a136dc8ca80bc3b65c637f5f04e9e2892567a840 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 01:07:16 -0700 Subject: [PATCH 039/222] Fix Transformer ignoring a custom encoder or decoder with no parameters (#3962) --- python/mlx/nn/layers/transformer.py | 46 ++++++++++++++++------------- python/tests/test_nn.py | 14 +++++++++ 2 files changed, 40 insertions(+), 20 deletions(-) diff --git a/python/mlx/nn/layers/transformer.py b/python/mlx/nn/layers/transformer.py index d7b65ed7cd..31ec916d19 100644 --- a/python/mlx/nn/layers/transformer.py +++ b/python/mlx/nn/layers/transformer.py @@ -327,27 +327,33 @@ def __init__( ): super().__init__() - self.encoder = custom_encoder or TransformerEncoder( - num_encoder_layers, - dims, - num_heads, - mlp_dims, - dropout, - activation, - norm_first, - checkpoint, - ) + if custom_encoder is not None: + self.encoder = custom_encoder + else: + self.encoder = TransformerEncoder( + num_encoder_layers, + dims, + num_heads, + mlp_dims, + dropout, + activation, + norm_first, + checkpoint, + ) - self.decoder = custom_decoder or TransformerDecoder( - num_decoder_layers, - dims, - num_heads, - mlp_dims, - dropout, - activation, - norm_first, - checkpoint, - ) + if custom_decoder is not None: + self.decoder = custom_decoder + else: + self.decoder = TransformerDecoder( + num_decoder_layers, + dims, + num_heads, + mlp_dims, + dropout, + activation, + norm_first, + checkpoint, + ) def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask): memory = self.encoder(src, src_mask) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 0df667a64d..951a0f22d2 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -2258,6 +2258,20 @@ def test_transformer(self): out = model(src, tgt, src_mask=None, tgt_mask=None, memory_mask=None) self.assertEqual(out.shape, tgt.shape) + def test_transformer_custom_modules(self): + # A custom encoder or decoder without any parameters is still a + # module and should not be replaced by the default one. + model = nn.Transformer( + dims=32, + num_heads=4, + num_encoder_layers=2, + num_decoder_layers=2, + custom_encoder=nn.Identity(), + custom_decoder=nn.Identity(), + ) + self.assertTrue(isinstance(model.encoder, nn.Identity)) + self.assertTrue(isinstance(model.decoder, nn.Identity)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From e50c0f2d39d024111075d4c3fd867e35879dae73 Mon Sep 17 00:00:00 2001 From: Recoordinate Date: Tue, 4 Aug 2026 21:36:40 +1200 Subject: [PATCH 040/222] docs: remove references to removed --no-verify-script launch flag (#3959) --- docs/src/usage/launching_distributed.rst | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/docs/src/usage/launching_distributed.rst b/docs/src/usage/launching_distributed.rst index e1e3128bb8..7c3a244920 100644 --- a/docs/src/usage/launching_distributed.rst +++ b/docs/src/usage/launching_distributed.rst @@ -160,11 +160,6 @@ host and on the same path. A good checklist to debug errors is the following: ``mlx.launch --print-python`` to see what that path is. * the script you want to run is available on all hosts at the same path -If you are launching from a node with a completely different setup than the -nodes that the program will run on, you can specify ``--no-verify-script`` so -that ``mlx.launch`` does not attempt to verify that the executable and script -exist locally before launching the distributed job. - .. _ring_specifics: Ring Specifics @@ -207,7 +202,7 @@ multi-gpu jobs. For instance .. code-block:: - mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 --no-verify-script -- ./my-job.sh + mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 -- ./my-job.sh will attempt to launch 16 processes, 8 on each node that will all run ``my-job.sh``. From 255f953f99c3403df19fa4d92462143139c3dfff Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Tue, 4 Aug 2026 15:16:05 +0530 Subject: [PATCH 041/222] Add eye(0) support (#3952) --- mlx/ops.cpp | 7 ++++++- python/tests/test_ops.py | 4 ++++ tests/ops_tests.cpp | 5 +++++ 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index da1901678d..a5dc03ec66 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -360,10 +360,15 @@ array ones_like(const array& a, StreamOrDevice s /* = {} */) { } array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s /* = {} */) { - if (n <= 0 || m <= 0) { + if (n < 0 || m < 0) { throw std::invalid_argument("[eye] N and M must be positive integers."); } array result = zeros({n, m}, dtype, s); + + if (n == 0 || m == 0) { + return result; + } + if (k >= m || -k >= n) { return result; } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 5d16c6e96b..86d92039f0 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2513,8 +2513,12 @@ def test_large_binary(self): def test_eye(self): self.assertCmpNumpy([3], mx.eye, np.eye) + # Test for zero rows and columns + self.assertCmpNumpy([0], mx.eye, np.eye) # Test for non-square matrix self.assertCmpNumpy([3, 4], mx.eye, np.eye) + # Test for zero rows + self.assertCmpNumpy([0, 4], mx.eye, np.eye) # Test with positive k parameter self.assertCmpNumpy([3, 4], mx.eye, np.eye, k=1) # Test with negative k parameter diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 987ee4df9b..741530aaf7 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3140,6 +3140,11 @@ TEST_CASE("test eye") { CHECK_EQ(eye_3x2.shape(), Shape{3, 2}); auto expected_eye_3x2 = array({1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f}, {3, 2}); CHECK(array_equal(eye_3x2, expected_eye_3x2).item()); + + auto eye_0x0 = eye(0, 0); + CHECK_EQ(eye_0x0.shape(), Shape{0, 0}); + CHECK_EQ(eye_0x0.size(), 0); + CHECK_EQ(eye_0x0.dtype(), float32); } TEST_CASE("test tri") { From 022ea8cb498401a48b809af5996447dd11387dad Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 4 Aug 2026 13:19:26 -0700 Subject: [PATCH 042/222] Refactor the JACCL ring and add threads for the multiple rings (#3900) --- .../jaccl/lib/jaccl/reduction_ops.h | 100 +++ mlx/distributed/jaccl/lib/jaccl/ring.cpp | 14 +- mlx/distributed/jaccl/lib/jaccl/ring.h | 2 + mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 791 +++++++++--------- mlx/distributed/jaccl/lib/jaccl/threadpool.h | 137 +++ 5 files changed, 633 insertions(+), 411 deletions(-) create mode 100644 mlx/distributed/jaccl/lib/jaccl/threadpool.h diff --git a/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h b/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h index a9b3e5e796..4dd395539f 100644 --- a/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h +++ b/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h @@ -8,6 +8,10 @@ namespace jaccl { +// Each reduction op has an in place form out[i] OP= in[i] and an out of place +// form out[i] = a[i] OP b[i]. The out of place pointers are __restrict, so +// callers must only use it when a, b and output are distinct buffers. + template struct SumOp { void operator()(const T* input, T* output, size_t N) const { @@ -15,6 +19,15 @@ struct SumOp { output[i] = output[i] + input[i]; } } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = a[i] + b[i]; + } + } }; template @@ -24,6 +37,15 @@ struct MaxOp { output[i] = (output[i] > input[i]) ? output[i] : input[i]; } } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] > b[i]) ? a[i] : b[i]; + } + } }; template @@ -33,6 +55,15 @@ struct MinOp { output[i] = (output[i] < input[i]) ? output[i] : input[i]; } } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] < b[i]) ? a[i] : b[i]; + } + } }; // @@ -73,6 +104,36 @@ native_bf16_min(const void* input, void* output, size_t N) { } } +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_sum(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = pa[i] + pb[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_max(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (pa[i] > pb[i]) ? pa[i] : pb[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_min(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (pa[i] < pb[i]) ? pa[i] : pb[i]; + } +} + template <> struct SumOp { void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const { @@ -84,6 +145,19 @@ struct SumOp { } } } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_sum(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = a[i] + b[i]; + } + } + } }; template <> @@ -97,6 +171,19 @@ struct MaxOp { } } } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_max(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] > b[i]) ? a[i] : b[i]; + } + } + } }; template <> @@ -110,6 +197,19 @@ struct MinOp { } } } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_min(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] < b[i]) ? a[i] : b[i]; + } + } + } }; #endif // defined(__aarch64__) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index c2099cb907..7fbfac04ca 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -17,7 +17,8 @@ RingGroup::RingGroup( n_conns_(left_devices.size()), side_channel_(std::move(sc)), left_(create_connections(left_devices)), - right_(create_connections(right_devices)) { + right_(create_connections(right_devices)), + pool_(n_conns_ > 0 ? n_conns_ - 1 : 0) { if (left_.size() > RING_MAX_CONNS || right_.size() > RING_MAX_CONNS) { std::ostringstream msg; msg << "[jaccl] Up to " << RING_MAX_CONNS << " per direction supported but " @@ -32,7 +33,8 @@ RingGroup::RingGroup( side_channel_.barrier(); // Create the ring implementation object - ring_ = RingImpl(rank_, size_, left_, right_, send_buffers_, recv_buffers_); + ring_ = RingImpl( + rank_, size_, left_, right_, send_buffers_, recv_buffers_, &pool_); } void RingGroup::initialize() { @@ -204,13 +206,9 @@ void RingGroup::all_reduce( auto in_ptr = static_cast(input); auto out_ptr = static_cast(output); int64_t count = n_bytes / sizeof(T); - if (count < size_ * 2 * n_conns_) { - ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); - return; - } - if (n_bytes <= 65536) { - ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); + if (n_bytes <= 32768 || count < size_ * 2 * n_conns_) { + ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); return; } diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h index a8b427ba80..94799dd267 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring.h @@ -78,6 +78,8 @@ class RingGroup : public Group { std::vector right_; std::vector send_buffers_; std::vector recv_buffers_; + // Declared before ring_ so it outlives the RingImpl that points at it. + ThreadPool pool_; RingImpl ring_; }; diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 766f55386d..412b178dc3 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -2,9 +2,11 @@ #pragma once +#include #include #include "jaccl/rdma.h" +#include "jaccl/threadpool.h" constexpr int RING_MAX_CONNS = 4; @@ -18,14 +20,16 @@ class RingImpl { std::vector& left, std::vector& right, std::vector& send_buffers, - std::vector& recv_buffers) + std::vector& recv_buffers, + ThreadPool* pool = nullptr) : rank_(rank), size_(size), n_conns_(left.size()), left_(left), right_(right), send_buffers_(send_buffers), - recv_buffers_(recv_buffers) {} + recv_buffers_(recv_buffers), + pool_(pool) {} RingImpl( int rank, @@ -34,16 +38,44 @@ class RingImpl { Connection* right_begin, size_t n_conns, std::vector& send_buffers, - std::vector& recv_buffers) + std::vector& recv_buffers, + ThreadPool* pool = nullptr) : rank_(rank), size_(size), n_conns_(n_conns), left_(left_begin, n_conns), right_(right_begin, n_conns), send_buffers_(send_buffers), - recv_buffers_(recv_buffers) {} + recv_buffers_(recv_buffers), + pool_(pool) {} - RingImpl() : rank_(0), size_(1), n_conns_(0) {} + RingImpl() : rank_(0), size_(1), n_conns_(0), pool_(nullptr) {} + + // Copy the received chunk into the output. Used by the all gather passes. + struct CopyOp { + template + inline void operator()(const T* recv, const T*, T* out, int64_t n) const { + std::copy(recv, recv + std::max(0, n), out); + } + }; + + // Reduce the received chunk with this rank's own input for the chunk. When + // in place (input aliases output) we use the fast two argument kernel; + // otherwise the fused out of place kernel seeds the output from the input. + template + struct ReduceRecvOp { + ReduceOp reduce_op; + template + inline void operator()(const T* recv, const T* base, T* out, int64_t n) + const { + n = std::max(0, n); + if constexpr (INPLACE) { + reduce_op(recv, out, n); + } else { + reduce_op(base, recv, out, n); + } + } + }; template void all_reduce( @@ -52,248 +84,108 @@ class RingImpl { int64_t size, int n_wires, ReduceOp reduce_op) { - // If not inplace all reduce then copy the input to the output first - if (in_ptr != out_ptr) { - std::memcpy(out_ptr, in_ptr, size * sizeof(T)); - } - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * MAX_DIR; int64_t chunk_size = (size + size_ - 1) / size_; int64_t size_per_wire = (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires); - auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T)); - N /= sizeof(T); - int64_t n_steps = (size_per_wire + N - 1) / N; - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t chunk_multiple_size = size_ * chunk_size; + // Split the reduce scatter + all gather across the available wires. Each + // wire handles a contiguous slice of each chunk in every direction. + dispatch_wires(n_wires, [&](int lw) { + all_reduce_wire( + in_ptr, + out_ptr, + size, + chunk_size, + size_per_wire, + n_wires, + lw, + reduce_op); + }); + } + + // Perform the ring all reduce (reduce scatter followed by all gather) for a + // single wire lw. + // + // Every chunk of chunk_size elements is divided into MAX_DIR directional + // regions and each region is further split into n_wires contiguous slices of + // size_per_wire elements. This function is responsible for slice lw in every + // direction and only touches left_[lw] / right_[lw] and the buffers that + // belong to wire lw, so several wires can run concurrently. + template + void all_reduce_wire( + const T* in_ptr, + T* out_ptr, + int64_t size, + int64_t chunk_size, + int64_t size_per_wire, + int n_wires, + int lw, + ReduceOp reduce_op) { + // The element offset (within a chunk) of this wire's slice in each + // direction and the end of each direction's region. Wire slices are + // contiguous rather than interleaved. Direction lr owns the chunk region + // [lr * n_wires * size_per_wire, (lr + 1) * n_wires * size_per_wire) (the + // last region is clamped to chunk_size). + int64_t wire_offset[MAX_DIR]; + int64_t region_end[MAX_DIR]; int64_t send_offset[MAX_DIR]; int64_t recv_offset[MAX_DIR]; - int64_t send_limits[MAX_DIR]; - int64_t recv_limits[MAX_DIR]; - int send_count[MAX_DIR * RING_MAX_CONNS] = {0}; - int recv_count[MAX_DIR * RING_MAX_CONNS] = {0}; - send_offset[0] = rank_ * chunk_size; + for (int lr = 0; lr < MAX_DIR; lr++) { + wire_offset[lr] = lr * n_wires * size_per_wire + + static_cast(lw) * size_per_wire; + region_end[lr] = std::min(chunk_size, (lr + 1) * n_wires * size_per_wire); + send_offset[lr] = rank_ * chunk_size; + } recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; if constexpr (MAX_DIR == 2) { - send_offset[1] = rank_ * chunk_size; recv_offset[1] = ((rank_ + 1) % size_) * chunk_size; - send_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - - // First reduce scatter - // - // Possible perf improvement by not syncing at every step but running ahead - // as needed. - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - reduce_op( - recv_buffer(sz, buff, lr, lw).begin(), - out_ptr + recv_offset[lr] + offset, - std::max(0, std::min(N, recv_limits[lr] - offset))); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } } - // Secondly all gather - // - // The offsets are correct from the scatter reduce - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, recv_limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } + // First reduce scatter, then all gather. The offsets carry over from the + // reduce scatter into the all gather. Pick the in place vs out of place + // reduction kernel once so the hot loop stays branch free. + if (in_ptr == out_ptr) { + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + region_end, + send_offset, + recv_offset, + in_ptr, + out_ptr, + ReduceRecvOp{reduce_op}); + } else { + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + region_end, + send_offset, + recv_offset, + in_ptr, + out_ptr, + ReduceRecvOp{reduce_op}); } + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + region_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); } void @@ -301,50 +193,128 @@ class RingImpl { // Copy our data to the appropriate place std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * 2; + // Split the all gather across the available wires. Each wire handles a + // contiguous slice of every rank's data in both directions. size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); - size_t out_bytes = n_bytes * size_; - auto [sz, N] = buffer_size_from_message(n_bytes_per_wire); - int n_steps = (n_bytes_per_wire + N - 1) / N; + dispatch_wires(n_wires, [&](int lw) { + all_gather_wire( + out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + }); + } + + // Perform the ring all gather for a single wire lw. + // + // The wire is responsible for the contiguous slice + // [lw * n_bytes_per_wire, (lw + 1) * n_bytes_per_wire) + // of every rank's n_bytes region, sent in the left direction (lr == 0), and + // the mirrored slice sent in the right direction (lr == 1). + // + // This function only ever touches left_[lw] / right_[lw] and the buffers that + // belong to wire lw so several wires can run concurrently. + void all_gather_wire( + char* out_ptr, + int64_t n_bytes, + int64_t n_bytes_per_wire, + int lw) { + // Both directions send the same contiguous slice of each rank's region. + int64_t slice = static_cast(lw) * n_bytes_per_wire; + int64_t wire_offset[2] = {slice, slice}; + int64_t region_end[2] = {n_bytes, n_bytes}; + int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; + int64_t recv_offset[2] = { + ((rank_ + size_ - 1) % size_) * n_bytes, + ((rank_ + 1) % size_) * n_bytes}; + + ring_pass<2, char>( + lw, + n_bytes, + n_bytes, + n_bytes * size_, + n_bytes_per_wire, + wire_offset, + region_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); + } + + // Run size_ - 1 pipelined ring steps for a single wire in every direction. + // + // At every step each direction sends its current slice to a neighbor and + // receives the neighbor's slice, which recv_op either reduces or copies into + // out_ptr. After each step the send and recv windows rotate around the ring + // (direction 0 backward, direction 1 forward). This is the shared engine + // behind both all reduce passes and all gather. + // + // The first step sends from in_ptr, later steps from out_ptr; this lets the + // all reduce seed its output without an up front copy. Callers that stage + // into out_ptr (both all gather paths) pass out_ptr for in_ptr. wire_offset + // is this wire's element offset within a chunk, region_end the end of each + // direction's region. send_offset and recv_offset are the starting windows, + // updated in place so callers can chain passes. stride is the elements a + // window moves per step and total is the wrap around modulus. + template + inline void ring_pass( + int lw, + int64_t size, + int64_t stride, + int64_t total, + int64_t size_per_wire, + const int64_t (&wire_offset)[MAX_DIR], + const int64_t (®ion_end)[MAX_DIR], + int64_t (&send_offset)[MAX_DIR], + int64_t (&recv_offset)[MAX_DIR], + const T* in_ptr, + T* out_ptr, + RecvOp recv_op) { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * 2 * MAX_DIR; + auto [sz, buffer_bytes] = + buffer_size_from_message(size_per_wire * sizeof(T)); + int64_t N = buffer_bytes / sizeof(T); + int64_t n_steps = (size_per_wire + N - 1) / N; - // Counters to maintain the state of transfers int in_flight = 0; - int64_t send_offset[2]; - int64_t recv_offset[2]; - int64_t limits[2]; - int send_count[2 * RING_MAX_CONNS] = {0}; - int recv_count[2 * RING_MAX_CONNS] = {0}; - send_offset[0] = send_offset[1] = rank_ * n_bytes; - recv_offset[0] = ((rank_ + size_ - 1) % size_) * n_bytes; - recv_offset[1] = ((rank_ + 1) % size_) * n_bytes; - limits[0] = n_wires * n_bytes_per_wire; - limits[1] = n_bytes; - - // Possible perf improvement by not syncing at every step but running ahead - // as needed. + int64_t send_limits[MAX_DIR]; + int64_t recv_limits[MAX_DIR]; + int send_count[MAX_DIR] = {0}; + int recv_count[MAX_DIR] = {0}; + auto set_limits = [&]() { + for (int lr = 0; lr < MAX_DIR; lr++) { + send_limits[lr] = std::min( + region_end[lr], std::max(0, size - send_offset[lr])); + recv_limits[lr] = std::min( + region_end[lr], std::max(0, size - recv_offset[lr])); + } + }; + set_limits(); + for (int k = 0; k < size_ - 1; k++) { + // Step 0 sends this rank's own input; later steps send the accumulated + // partial from out_ptr. + const T* send_base = (k == 0) ? in_ptr : out_ptr; + // Prefill the pipeline int buff = 0; while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff); - for (int lr = 0; lr < 2; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } + for (int lr = 0; lr < MAX_DIR; lr++) { + recv_from(sz, buff, lr, lw); + } + for (int lr = 0; lr < MAX_DIR; lr++) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_offset[lr] + offset, + send_base + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); } - post_send_all(sz, buff); buff++; - in_flight += 2 * 2 * n_wires; + in_flight += 2 * MAX_DIR; } // Main loop @@ -352,39 +322,37 @@ class RingImpl { // Keep going until we have no longer data in flight. while (in_flight > 0) { ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); + int n = poll_wire(lw, WC_NUM, wc); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; + int lr = wc[i].wr_id & 0xff; in_flight--; - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; + if (work_type == SEND_WR && send_count[lr] < n_steps) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); + send_base + send_offset[lr] + offset, + send_base + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).template begin()); send_to(sz, buff, lr, lw); in_flight++; - send_count[wire]++; + send_count[lr]++; } else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { + int64_t offset = wire_offset[lr] + recv_count[lr] * N; + // The base operand is this rank's own input for the chunk; the all + // gather passes copy and ignore it. + recv_op( + recv_buffer(sz, buff, lr, lw).template begin(), + in_ptr + recv_offset[lr] + offset, + out_ptr + recv_offset[lr] + offset, + std::min(N, recv_limits[lr] - offset)); + recv_count[lr]++; + if (recv_count[lr] + (PIPELINE - 1) < n_steps) { recv_from(sz, buff, lr, lw); in_flight++; } @@ -392,12 +360,16 @@ class RingImpl { } } - send_offset[0] = (send_offset[0] + out_bytes - n_bytes) % out_bytes; - recv_offset[0] = (recv_offset[0] + out_bytes - n_bytes) % out_bytes; - send_offset[1] = (send_offset[1] + n_bytes) % out_bytes; - recv_offset[1] = (recv_offset[1] + n_bytes) % out_bytes; - for (int i = 0; i < 2 * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; + // Rotate the windows around the ring for the next step. + send_offset[0] = (send_offset[0] + total - stride) % total; + recv_offset[0] = (recv_offset[0] + total - stride) % total; + if constexpr (MAX_DIR == 2) { + send_offset[1] = (send_offset[1] + stride) % total; + recv_offset[1] = (recv_offset[1] + stride) % total; + } + set_limits(); + for (int lr = 0; lr < MAX_DIR; lr++) { + send_count[lr] = recv_count[lr] = 0; } } } @@ -408,37 +380,50 @@ class RingImpl { // In the case that size_ == 2 then left == right so we bias send towards // left and recv towards right so that the selections will be correct for // the 2 node case. - auto& conns = (dst == left) ? left_ : right_; int dir = dst == left; + int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; + + // Split the send across the available wires. Each wire handles the + // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). + dispatch_wires(n_wires, [&](int lw) { + send_wire(in_ptr, n_bytes, dir, bytes_per_wire, lw); + }); + } + + // Perform a point-to-point send for a single wire lw. + // + // Only touches the connection and buffers of wire lw so several wires can run + // concurrently. + void send_wire( + const char* in_ptr, + int64_t n_bytes, + int dir, + int64_t bytes_per_wire, + int lw) { + auto& conns = dir ? left_ : right_; + constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS; + constexpr int WC_NUM = PIPELINE; - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; auto [sz, N] = buffer_size_from_message(bytes_per_wire); int in_flight = 0; - int64_t read_offset[RING_MAX_CONNS]; - int64_t limits[RING_MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - read_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } + int64_t read_offset = std::min(lw * bytes_per_wire, n_bytes); + int64_t limit = std::min((lw + 1) * bytes_per_wire, n_bytes); // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (read_offset[lw] < limits[lw] && buff < PIPELINE) { - std::copy( - in_ptr + read_offset[lw], - in_ptr + std::min(read_offset[lw] + N, limits[lw]), - send_buffer(sz, buff, dir, lw).begin()); - send_to(sz, buff, dir, lw); - - buff++; - read_offset[lw] += N; - in_flight++; - } + int buff = 0; + while (read_offset < limit && buff < PIPELINE) { + std::copy( + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, limit), + send_buffer(sz, buff, dir, lw).begin()); + send_to(sz, buff, dir, lw); + + buff++; + read_offset += N; + in_flight++; } // Main loop @@ -448,22 +433,20 @@ class RingImpl { // If a send was completed and we have more data to send then go ahead // and send them. ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); + int n = conns[lw].poll(WC_NUM, wc); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % RING_MAX_CONNS; in_flight--; - if (read_offset[lw] < limits[lw]) { + if (read_offset < limit) { std::copy( - in_ptr + read_offset[lw], - in_ptr + std::min(read_offset[lw] + N, limits[lw]), + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, limit), send_buffer(sz, buff, dir, lw).begin()); send_to(sz, buff, dir, lw); - read_offset[lw] += N; + read_offset += N; in_flight++; } } @@ -476,32 +459,45 @@ class RingImpl { // In the case that size_ == 2 then left == right so we bias send towards // left and recv towards right so that the selections will be correct for // the 2 node case. - auto& conns = (src == right) ? right_ : left_; int dir = src == right; + int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; + + // Split the recv across the available wires. Each wire handles the + // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). + dispatch_wires(n_wires, [&](int lw) { + recv_wire(out_ptr, n_bytes, dir, bytes_per_wire, lw); + }); + } + + // Perform a point-to-point recv for a single wire lw. + // + // Only touches the connection and buffers of wire lw so several wires can run + // concurrently. + void recv_wire( + char* out_ptr, + int64_t n_bytes, + int dir, + int64_t bytes_per_wire, + int lw) { + auto& conns = dir ? right_ : left_; + constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS; + constexpr int WC_NUM = PIPELINE; - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; auto [sz, N] = buffer_size_from_message(bytes_per_wire); int in_flight = 0; - int64_t write_offset[RING_MAX_CONNS]; - int64_t limits[RING_MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - write_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } + int64_t write_offset = std::min(lw * bytes_per_wire, n_bytes); + int64_t limit = std::min((lw + 1) * bytes_per_wire, n_bytes); // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (N * buff < limits[lw] - write_offset[lw] && buff < PIPELINE) { - recv_from(sz, buff, dir, lw); + int buff = 0; + while (write_offset + N * buff < limit && buff < PIPELINE) { + recv_from(sz, buff, dir, lw); - buff++; - in_flight++; - } + buff++; + in_flight++; } // Main loop @@ -511,11 +507,9 @@ class RingImpl { // If a recv was completed copy it to the output and if we have more // data to fetch post another recv. ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); + int n = conns[lw].poll(WC_NUM, wc); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % RING_MAX_CONNS; in_flight--; @@ -523,11 +517,11 @@ class RingImpl { recv_buffer(sz, buff, dir, lw).begin(), recv_buffer(sz, buff, dir, lw).begin() + std::max( - 0, std::min(limits[lw] - write_offset[lw], N)), - out_ptr + write_offset[lw]); - write_offset[lw] += N; + 0, std::min(limit - write_offset, N)), + out_ptr + write_offset); + write_offset += N; - if (write_offset[lw] + (PIPELINE - 1) * N < limits[lw]) { + if (write_offset + (PIPELINE - 1) * N < limit) { recv_from(sz, buff, dir, lw); in_flight++; @@ -538,36 +532,17 @@ class RingImpl { private: void send_to(int sz, int buff, int left_right, int wire) { - if (left_right) { - left_[wire].post_send( - send_buffer_left(sz, buff, wire), - SEND_WR << 16 | buff << 8 | (RING_MAX_CONNS + wire)); - } else { - right_[wire].post_send( - send_buffer_right(sz, buff, wire), SEND_WR << 16 | buff << 8 | wire); - } + auto& conns = left_right ? left_ : right_; + conns[wire].post_send( + send_buffer(sz, buff, left_right, wire), + SEND_WR << 16 | buff << 8 | left_right); } void recv_from(int sz, int buff, int left_right, int wire) { - if (left_right) { - right_[wire].post_recv( - recv_buffer_right(sz, buff, wire), - RECV_WR << 16 | buff << 8 | (RING_MAX_CONNS + wire)); - } else { - left_[wire].post_recv( - recv_buffer_left(sz, buff, wire), RECV_WR << 16 | buff << 8 | wire); - } - } - - SharedBuffer& send_buffer_right(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + wire]; - } - - SharedBuffer& send_buffer_left(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + n_conns_ + - wire]; + auto& conns = left_right ? right_ : left_; + conns[wire].post_recv( + recv_buffer(sz, buff, left_right, wire), + RECV_WR << 16 | buff << 8 | left_right); } SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) { @@ -576,47 +551,56 @@ class RingImpl { left_right * n_conns_ + wire]; } - SharedBuffer& recv_buffer_left(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + wire]; - } - - SharedBuffer& recv_buffer_right(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + n_conns_ + - wire]; - } - SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) { return recv_buffers_ [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + left_right * n_conns_ + wire]; } - template - void post_recv_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { + // Poll the completion queues that belong to a single wire. + // + // A wire always uses both its left and right connections (direction 0 sends + // right / receives left, direction 1 sends left / receives right) so both are + // polled here. Restricting polling to a single wire's connections is what + // allows several wires to run concurrently. + int poll_wire(int wire, int num_completions, ibv_wc* work_completions) { + return poll( + std::span(&left_[wire], 1), + std::span(&right_[wire], 1), + num_completions, + work_completions); + } + + // Run fn(lw) for each wire, the first n_wires - 1 on the pool and the last + // inline, then wait for the pool calls before returning. + template + void dispatch_wires(int n_wires, Fn&& fn) { + if (n_wires <= 1 || pool_ == nullptr) { for (int lw = 0; lw < n_wires; lw++) { - recv_from(sz, buff, lr, lw); + fn(lw); } + return; } - } - void post_recv_all(int sz, int buff) { - post_recv_all<2>(sz, buff, n_conns_); - } + std::vector> futures; + futures.reserve(n_wires - 1); + for (int lw = 0; lw < n_wires - 1; lw++) { + futures.emplace_back(pool_->enqueue(fn, lw)); + } - template - void post_send_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - send_to(sz, buff, lr, lw); + // Wait for the pool calls even if the inline one throws, so they never + // outlive this frame. + try { + fn(n_wires - 1); + } catch (...) { + for (auto& f : futures) { + f.wait(); } + throw; + } + for (auto& f : futures) { + f.wait(); } - } - - void post_send_all(int sz, int buff) { - post_send_all<2>(sz, buff, n_conns_); } int rank_; @@ -626,6 +610,7 @@ class RingImpl { std::span right_; std::span send_buffers_; std::span recv_buffers_; + ThreadPool* pool_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/threadpool.h b/mlx/distributed/jaccl/lib/jaccl/threadpool.h new file mode 100644 index 0000000000..b5d05153e2 --- /dev/null +++ b/mlx/distributed/jaccl/lib/jaccl/threadpool.h @@ -0,0 +1,137 @@ +// This code was modified from https://github.com/progschj/ThreadPool +// The original License is copied below: +// +// Copyright (c) 2012 Jakob Progsch, Václav Zeman +// This software is provided 'as-is', without any express or implied +// warranty. In no event will the authors be held liable for any damages +// arising from the use of this software. +// +// Permission is granted to anyone to use this software for any purpose, +// including commercial applications, and to alter it and redistribute it +// freely, subject to the following restrictions: +// +// 1. The origin of this software must not be misrepresented; you must not +// claim that you wrote the original software. If you use this software +// in a product, an acknowledgment in the product documentation would be +// appreciated but is not required. +// +// 2. Altered source versions must be plainly marked as such, and must not be +// misrepresented as being the original software. +// +// 3. This notice may not be removed or altered from any source +// distribution. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jaccl { + +class ThreadPool { + public: + ThreadPool(size_t); + template + auto enqueue(F&& f, Args&&... args) + -> std::future>; + void resize(size_t); + ~ThreadPool(); + + private: + void stop_and_wait(); + void start_threads(size_t); + + std::vector workers; + std::queue> tasks; + std::mutex queue_mutex; + std::condition_variable condition; + bool stop; +}; + +inline ThreadPool::ThreadPool(size_t threads) : stop(false) { + start_threads(threads); +} + +template +auto ThreadPool::enqueue(F&& f, Args&&... args) + -> std::future> { + using return_type = typename std::invoke_result_t; + + auto task = std::make_shared>( + std::bind(std::forward(f), std::forward(args)...)); + + std::future res = task->get_future(); + { + std::unique_lock lock(queue_mutex); + + if (stop) { + throw std::runtime_error( + "[ThreadPool::enqueue] Not allowed on stopped ThreadPool"); + } + + tasks.emplace([task]() { (*task)(); }); + } + condition.notify_one(); + return res; +} + +inline void ThreadPool::resize(size_t threads) { + if (workers.size() == threads) { + return; + } + + if (workers.size() > threads) { + stop_and_wait(); + } + start_threads(threads - workers.size()); +} + +inline ThreadPool::~ThreadPool() { + stop_and_wait(); +} + +inline void ThreadPool::stop_and_wait() { + // Stop the current threads and wait until they finish + { + std::unique_lock lock(queue_mutex); + stop = true; + } + condition.notify_all(); + for (std::thread& worker : workers) { + worker.join(); + } + + // Reset the member variables so that the threadpool is reusable + stop = false; + workers.clear(); +} + +inline void ThreadPool::start_threads(size_t threads) { + for (size_t i = 0; i < threads; ++i) { + workers.emplace_back([this] { + for (;;) { + std::function task; + + { + std::unique_lock lock(this->queue_mutex); + this->condition.wait( + lock, [this] { return this->stop || !this->tasks.empty(); }); + if (this->stop && this->tasks.empty()) + return; + task = std::move(this->tasks.front()); + this->tasks.pop(); + } + + task(); + } + }); + } +} + +} // namespace jaccl From 6dfd81406dbac598a45f0610f59f8e99a33d43f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Devansh=20Varshney=20=28=E0=A4=A6=E0=A5=87=E0=A4=B5?= =?UTF-8?q?=E0=A4=BE=E0=A4=82=E0=A4=B6=20=E0=A4=B5=E0=A4=BE=E0=A4=B0?= =?UTF-8?q?=E0=A5=8D=E0=A4=B7=E0=A5=8D=E0=A4=A3=E0=A5=87=E0=A4=AF=29?= Date: Wed, 5 Aug 2026 03:33:38 +0530 Subject: [PATCH 043/222] Fix shapeless matmul with dynamic batch dimensions (#3813) --- mlx/ops.cpp | 12 ++++++-- python/tests/test_export_import.py | 44 ++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index a5dc03ec66..0356ad0cae 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3357,9 +3357,14 @@ array matmul( } // We can batch the multiplication by reshaping a - if (in_a.ndim() > 2 && in_b.ndim() <= 2) { + bool flattened_a = false; + + // Avoid flatten/unflatten during dynamic tracing because unflatten stores the + // trace-time shape, which breaks shapeless replay with dynamic batch sizes. + if (in_a.ndim() > 2 && in_b.ndim() <= 2 && !detail::in_dynamic_tracing()) { a = flatten(a, 0, -2, s); - } else if (in_b.ndim() > 2) { + flattened_a = true; + } else if (in_a.ndim() > 2 || in_b.ndim() > 2) { std::tie(a, b) = broadcast_arrays(a, b, {-2, -1}, s); } @@ -3371,7 +3376,8 @@ array matmul( out_type, std::make_shared(to_stream(s)), {a, b}); - if (in_a.ndim() > 2 && in_b.ndim() <= 2) { + + if (flattened_a) { auto orig_shape = in_a.shape(); orig_shape.pop_back(); out = unflatten(out, 0, std::move(orig_shape), s); diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index c45dd4a606..6b5f0ab146 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -704,6 +704,50 @@ def fun(x, y, z): imported = mx.import_function(path) self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z))) + def test_export_matmul_shapeless_mid_dim(self): + path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn") + + E, H = 64, 17 + arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H)) + + def fn(x): + return mx.matmul(x, arr) + + sample = mx.zeros((1, 40, E), dtype=mx.float32) + mx.export_function(path, fn, sample, shapeless=True) + imported = mx.import_function(path) + + for seq_len in (40, 248, 623): + with self.subTest(seq_len=seq_len): + x = mx.arange(seq_len * E, dtype=mx.float32).reshape((1, seq_len, E)) + expected = fn(x) + (y,) = imported(x) + self.assertEqual(y.shape, (1, seq_len, H)) + self.assertTrue(mx.allclose(y, expected)) + + def test_export_matmul_shapeless_batch_and_mid_dim(self): + path = os.path.join(self.test_dir, "matmul_shapeless_batch.mlxfn") + + B, E, H = 2, 32, 8 + arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H)) + + def fn(x): + return mx.matmul(x, arr) + + sample = mx.zeros((B, 10, E), dtype=mx.float32) + mx.export_function(path, fn, sample, shapeless=True) + imported = mx.import_function(path) + + for seq_len in (10, 50, 100): + with self.subTest(seq_len=seq_len): + x = mx.arange(B * seq_len * E, dtype=mx.float32).reshape( + (B, seq_len, E) + ) + expected = fn(x) + (y,) = imported(x) + self.assertEqual(y.shape, (B, seq_len, H)) + self.assertTrue(mx.allclose(y, expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From b34e3323fd7b71320580c1f210e3abac7a34989d Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:03:49 -0700 Subject: [PATCH 044/222] Fix JIT build with old macOS SDK (#3853) --- mlx/backend/metal/CMakeLists.txt | 35 +++++++++++++++++++++---------- mlx/backend/metal/jit_kernels.cpp | 34 ++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 11 deletions(-) diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index 5eb8f543c7..e7a4d9d2af 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -84,19 +84,32 @@ if(MLX_METAL_JIT) make_jit_source(steel/attn/kernels/steel_attention) - make_jit_source( - steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h - kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) - make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) - make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) - make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) - make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) + if(MLX_METAL_VERSION GREATER_EQUAL 400 + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 + AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) - make_jit_source(quantized_nax kernels/quantized_utils.h) - make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h - kernels/fp4.h) + make_jit_source( + steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h + kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) + make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) + + make_jit_source(quantized_nax kernels/quantized_utils.h) + make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h + kernels/fp4.h) + + make_jit_source(steel/attn/kernels/steel_attention_nax) - make_jit_source(steel/attn/kernels/steel_attention_nax) + else() + message( + WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and " + "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, " + "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). " + "Building without NAX kernels.") + target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX) + endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp) diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index d639719d81..1384e06c50 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -8,6 +8,40 @@ using namespace fmt::literals; namespace mlx::core { +#ifdef MLX_METAL_NO_NAX +// NAX JIT preambles are only generated (via make_jit_source) when the SDK +// requirement is met. On older SDKs they are skipped and MLX_METAL_NO_NAX is +// defined, so is_nax_available() returns false and the get_*_nax_kernel entry +// points below are never reached. These empty definitions only exist to satisfy +// the linker for this translation unit. +namespace metal { +const char* gemm_nax() { + return ""; +} +const char* steel_gemm_fused_nax() { + return ""; +} +const char* steel_gemm_gather_nax() { + return ""; +} +const char* steel_gemm_splitk_nax() { + return ""; +} +const char* steel_gemm_segmented_nax() { + return ""; +} +const char* quantized_nax() { + return ""; +} +const char* fp_quantized_nax() { + return ""; +} +const char* steel_attention_nax() { + return ""; +} +} // namespace metal +#endif // MLX_METAL_NO_NAX + MTL::ComputePipelineState* get_arange_kernel( metal::Device& d, const std::string& kernel_name, From 1c9845031dfe2571fa77144f0f5b3a2e20911428 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 15:48:32 -0700 Subject: [PATCH 045/222] chore: Fix cholesky_inv arg name (#3950) --- python/src/linalg.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/src/linalg.cpp b/python/src/linalg.cpp index bfebe3bb47..0bf7b6f12c 100644 --- a/python/src/linalg.cpp +++ b/python/src/linalg.cpp @@ -326,7 +326,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cholesky_inv(L: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cholesky_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition. @@ -347,7 +347,8 @@ void init_linalg(nb::module_& parent_module) { If the input matrix is not a triangular matrix behaviour is undefined. Args: - L (array): Input array. + a (array): Input array. This is the Cholesky factor + :math:`\mathbf{L}`, not :math:`\mathbf{A}` itself. upper (bool, optional): If ``True``, return the upper triangular Cholesky factor. If ``False``, return the lower triangular Cholesky factor. Default: ``False``. stream (Stream, optional): Stream or device. Defaults to ``None`` From 9a1e04d1a1c0375bbce9d1093e3994d1bc164cd4 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 16:21:01 -0700 Subject: [PATCH 046/222] docs: Add new_thread_unsafe_stream (#3968) --- docs/src/python/devices_and_streams.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/python/devices_and_streams.rst b/docs/src/python/devices_and_streams.rst index 9d0d15f2a5..843f122c98 100644 --- a/docs/src/python/devices_and_streams.rst +++ b/docs/src/python/devices_and_streams.rst @@ -15,6 +15,7 @@ Devices and Streams default_stream new_stream new_thread_local_stream + new_thread_unsafe_stream set_default_stream stream synchronize From a79332de6982cc32725b8d84ba80c4fdcbd481f1 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 5 Aug 2026 08:26:39 +0900 Subject: [PATCH 047/222] Add gather_qqmm (#3757) --- mlx/backend/cpu/quantized.cpp | 4 + mlx/backend/cuda/device/qmm_naive.cuh | 14 + mlx/backend/cuda/quantized/qmm/qmm.h | 1 + mlx/backend/cuda/quantized/qmm/qmm_naive.cu | 5 + mlx/backend/cuda/quantized/qqmm.cpp | 202 ++++-- mlx/backend/cuda/quantized/quantized.cpp | 2 + mlx/backend/metal/kernels/fp4.h | 3 + mlx/backend/metal/kernels/fp_quantized.h | 243 ++++++-- mlx/backend/metal/kernels/fp_quantized.metal | 79 ++- mlx/backend/metal/kernels/quantized.h | 55 +- mlx/backend/metal/quantized.cpp | 617 +++++++++++++------ mlx/backend/no_cpu/primitives.cpp | 1 + mlx/backend/no_gpu/primitives.cpp | 1 + mlx/ops.cpp | 168 ++--- mlx/ops.h | 34 +- mlx/primitives.cpp | 16 + mlx/primitives.h | 35 ++ python/src/ops.cpp | 52 ++ python/tests/test_quantized.py | 177 +++++- 19 files changed, 1220 insertions(+), 489 deletions(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index c0f1a3c315..3469d99788 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -1359,4 +1359,8 @@ void QQMatmul::eval_cpu(const std::vector& inputs, array& out) { } } +void GatherQQMM::eval_cpu(const std::vector& inputs, array& out) { + throw std::runtime_error("[GatherQQMM] NYI"); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/device/qmm_naive.cuh b/mlx/backend/cuda/device/qmm_naive.cuh index 01e5f444d5..e2ad855b9f 100644 --- a/mlx/backend/cuda/device/qmm_naive.cuh +++ b/mlx/backend/cuda/device/qmm_naive.cuh @@ -2,6 +2,7 @@ #include "mlx/backend/cuda/device/cute_dequant.cuh" #include "mlx/backend/cuda/device/gemm_sm70.cuh" +#include "mlx/backend/cuda/device/utils.cuh" #include @@ -25,6 +26,7 @@ CUTE_DEVICE void qmm_naive_mainloop( TensorS gS, TensorZ gZ, TensorC gC, + const float* global_scale, int m_max_coord, int n_max_coord, int k_residue, @@ -32,6 +34,7 @@ CUTE_DEVICE void qmm_naive_mainloop( // Get the types of operands. using Element = typename decltype(gA)::value_type; using Quant = typename decltype(gB)::value_type; + using Scale = typename decltype(gS)::value_type; // Shift tensor so we handle residue of K in the 0th tile. gA = domain_offset(make_coord(0, k_residue, 0), gA); @@ -196,6 +199,15 @@ CUTE_DEVICE void qmm_naive_mainloop( CUTE_UNROLL for (int i = 0; i < size(tCrC); ++i) { if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) { + if constexpr ( + cuda::std::is_same_v && + cuda::std::is_same_v) { + // Only nvfp4 supports global scale. + if (global_scale) { + tCgC(i) = Element(tCrC(i) * (*global_scale / (F8E4M3_MAX * F4E2M1_MAX))); + continue; + } + } tCgC(i) = Element(tCrC(i)); } } @@ -224,6 +236,7 @@ void qmm_naive_kernel( const Quant* B, const Scale* S, const Element* Z, + const float* global_scale, const uint32_t* lhs_indices, const uint32_t* rhs_indices, Element* C, @@ -295,6 +308,7 @@ void qmm_naive_kernel( gS, gZ, gC, + global_scale, m_max_coord, n_max_coord, k_residue, thread_idx); } diff --git a/mlx/backend/cuda/quantized/qmm/qmm.h b/mlx/backend/cuda/quantized/qmm/qmm.h index 8d998cda40..64cfecfd5a 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/mlx/backend/cuda/quantized/qmm/qmm.h @@ -74,6 +74,7 @@ void qmm_naive( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const std::optional& lhs_indices, const std::optional& rhs_indices, array& out, diff --git a/mlx/backend/cuda/quantized/qmm/qmm_naive.cu b/mlx/backend/cuda/quantized/qmm/qmm_naive.cu index cb47d7f1aa..c7c5c6049a 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_naive.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm_naive.cu @@ -29,6 +29,7 @@ void qmm_naive( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const std::optional& lhs_indices, const std::optional& rhs_indices, array& out, @@ -75,6 +76,9 @@ void qmm_naive( if (biases) { encoder.set_input_array(*biases); } + if (global_scale) { + encoder.set_input_array(*global_scale); + } if (lhs_indices) { encoder.set_input_array(*lhs_indices); } @@ -103,6 +107,7 @@ void qmm_naive( gpu_ptr(w), gpu_ptr(scales), biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, lhs_indices ? gpu_ptr(*lhs_indices) : nullptr, rhs_indices ? gpu_ptr(*rhs_indices) : nullptr, gpu_ptr(out), diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index eaec2ac8f4..196bd9c05e 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -21,7 +21,7 @@ std::tuple quantize_input( QuantizationMode mode, int bits, int group_size, - std::optional global_scale = std::nullopt) { + std::optional global_scale) { const array x = ensure_contiguous(input, encoder, s); // Compute output shapes @@ -52,6 +52,27 @@ std::tuple quantize_input( return {std::move(x_q), std::move(scales_x)}; } +array quantize_dequantize_input( + const array& x_pre, + const std::optional& global_scale, + int bits, + int group_size, + cu::CommandEncoder& encoder, + Stream s) { + bool donate_x = x_pre.is_donatable(); + array x = ensure_row_contiguous(x_pre, encoder, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = donate_x + ? x + : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype()); + if (!donate_x) { + encoder.add_temporary(xhat); + } + fp_quantize_dequantize(x, xhat, group_size, bits, global_scale, encoder, s); + return xhat; +} + GemmScalars create_nvfp4_scalars( const array& global_scale_x, const array& global_scale_w, @@ -75,77 +96,81 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); auto& device = encoder.device(); - bool w_quantized = (inputs[1].dtype() == uint32); + + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // - 2 inputs: x, w (non-quantized w) // - 3 inputs: x, w, scales_w (quantized w) + bool w_quantized = (w_pre.dtype() == uint32); int base_size = w_quantized ? 3 : 2; - assert( - inputs.size() == base_size || - (mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2)); - // For nvfp4, global scales are optional but must be both present or both // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) bool has_global_scales = - mode_ == QuantizationMode::Nvfp4 && inputs.size() > base_size; - std::optional global_scale_x = std::nullopt; - std::optional global_scale_w = std::nullopt; + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; if (has_global_scales) { global_scale_x = inputs[inputs.size() - 2]; global_scale_w = inputs[inputs.size() - 1]; } - if (w_quantized && inputs[0].shape(-2) == 1) { - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - - bool donate_x = inputs[0].is_donatable(); - array x = ensure_row_contiguous(inputs[0], encoder, s); - // If x is a copy it should be donatable - donate_x |= x.is_donatable(); - auto xhat = donate_x - ? x - : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype()); - if (!donate_x) { - encoder.add_temporary(xhat); + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input( + w_pre, encoder, s, mode_, bits_, group_size_, global_scale_w) + : std::make_tuple( + ensure_contiguous(w_pre, encoder, s), + ensure_contiguous(inputs[base_size - 1], encoder, s)); + + // Reroute to qmm when: no support in cuBLAS, or doing GEMV. + bool can_use_cublas = + (mode_ == QuantizationMode::Nvfp4 || mode_ == QuantizationMode::Mxfp8) && + (device.compute_capability_major() >= 10); + int M = x_pre.shape(-2); + bool use_qmm = (!can_use_cublas) || (M == 1); + + if (use_qmm) { + array x = quantize_dequantize_input( + x_pre, global_scale_x, bits_, group_size_, encoder, s); + if (M < 8) { + qmv(x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + bits_, + group_size_, + mode_, + encoder); + } else { + qmm_naive( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + std::nullopt, + std::nullopt, + out, + true, // transpose + bits_, + group_size_, + mode_, + encoder); } - fp_quantize_dequantize( - x, xhat, group_size_, bits_, global_scale_x, encoder, s); - - const array& w = inputs[1]; - const array& scales = inputs[2]; - qmv(xhat, - w, - scales, - std::nullopt, - global_scale_w, - out, - bits_, - group_size_, - mode_, - encoder); return; } - auto cc = device.compute_capability_major() * 100 + - device.compute_capability_minor() * 10; - if (cc < 1000) { - throw std::runtime_error( - "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher."); - } - - // Quantize inputs (or use pre-quantized) - auto [x_q, scale_x_pre] = quantize_input( - inputs[0], encoder, s, mode_, bits_, group_size_, global_scale_x); - auto [w_q, scale_w_pre] = !w_quantized - ? quantize_input( - inputs[1], encoder, s, mode_, bits_, group_size_, global_scale_w) - : std::make_tuple( - ensure_contiguous(inputs[1], encoder, s), - ensure_contiguous(inputs[2], encoder, s)); - - out.set_data(cu::malloc_async(out.nbytes(), encoder)); + // Quantize activation. + auto [x_q, scales_x] = quantize_input( + x_pre, encoder, s, mode_, bits_, group_size_, global_scale_x); - int M = x_q.shape(-2); int N = w_q.shape(-2); // transposed int K = x_q.shape(-1) * (32 / bits_); @@ -155,8 +180,8 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int64_t ldb = K; // Repack scales to tiled layout for tensor cores - array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s); - array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s); + scales_x = pad_and_swizzle_scales(scales_x, encoder, s); + scales_w = pad_and_swizzle_scales(scales_w, encoder, s); GemmScalars scalars; if (has_global_scales) { @@ -175,10 +200,69 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { out, x_q, w_q, - scale_x, - scale_w, + scales_x, + scales_w, mode_, scalars); } +void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("QQMatmul::eval_gpu"); + + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + const array& lhs_indices = ensure_row_contiguous(inputs[2], encoder, s); + const array& rhs_indices = ensure_row_contiguous(inputs[3], encoder, s); + + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + // - 4 inputs: x, w, lhs_indices, rhs_indices (non-quantized w) + // - 5 inputs: x, w, lhs_indices, rhs_indices, scales_w (quantized w) + bool w_quantized = (w_pre.dtype() == uint32); + int base_size = w_quantized ? 5 : 4; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; + } + + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input( + w_pre, encoder, s, mode_, bits_, group_size_, global_scale_w) + : std::make_tuple( + ensure_contiguous(w_pre, encoder, s), + ensure_contiguous(inputs[base_size - 1], encoder, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, bits_, group_size_, encoder, s); + + // Reroute to qmm. + qmm_naive( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + true, // transpose + bits_, + group_size_, + mode_, + encoder); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 2a1a268c91..4d25f3c3e0 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -72,6 +72,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { biases, std::nullopt, std::nullopt, + std::nullopt, out, transpose_, bits_, @@ -211,6 +212,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, diff --git a/mlx/backend/metal/kernels/fp4.h b/mlx/backend/metal/kernels/fp4.h index 47bf4dda6f..af4cbb02cf 100644 --- a/mlx/backend/metal/kernels/fp4.h +++ b/mlx/backend/metal/kernels/fp4.h @@ -1,5 +1,8 @@ #pragma once +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { fp4_e2m1(float x) thread { if (metal::isnan(x)) { diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 85d4284ac3..f3aa24bf7a 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -321,10 +321,11 @@ METAL_FUNC void fp_qmv_quad_impl( } } -template +template METAL_FUNC void fp_qmv_fast_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -374,18 +375,28 @@ METAL_FUNC void fp_qmv_fast_impl( x += block_size; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } -template +template METAL_FUNC void fp_qmv_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -421,6 +432,11 @@ METAL_FUNC void fp_qmv_impl( return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // In this case we need to properly guard all our reads because there isn't // even 1 tile in the matrix if (out_vec_size < (num_simdgroups * results_per_simdgroup)) { @@ -471,7 +487,11 @@ METAL_FUNC void fp_qmv_impl( row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } @@ -519,7 +539,11 @@ METAL_FUNC void fp_qmv_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } @@ -644,10 +668,11 @@ METAL_FUNC void fp_qmv_wide_impl( } } -template +template METAL_FUNC void fp_qvm_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const int in_vec_size, @@ -690,6 +715,11 @@ METAL_FUNC void fp_qvm_impl( return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // Loop over in_vec in blocks of block_size int remaining = in_vec_size % block_size; if (remaining == 0) { @@ -739,7 +769,11 @@ METAL_FUNC void fp_qvm_impl( if (simd_lid == 0) { #pragma clang loop unroll(full) for (int k = 0; k < tn * pack_factor; k++) { - y[k] = static_cast(result[k]); + if constexpr (has_global_scale) { + y[k] = static_cast(result[k] * inv_scale_enc); + } else { + y[k] = static_cast(result[k]); + } } } } @@ -872,11 +906,11 @@ METAL_FUNC void fp_qmm_t_impl( template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + int BM = 32, + int BK = 32, + int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, @@ -1128,10 +1162,16 @@ template w, scales, x, y, in_vec_size, out_vec_size, tid, quad_gid, quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void fp_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1163,14 +1203,29 @@ template s_strides, tid); } - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void fp_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1202,8 +1257,17 @@ template s_strides, tid); } - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } template < @@ -1251,10 +1315,16 @@ template < w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void fp_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1286,9 +1356,10 @@ template s_strides, tid); } - fp_qvm_impl( + fp_qvm_impl( w, scales, + global_scale, x, y, in_vec_size, @@ -1341,9 +1412,10 @@ template // The in_vec_stride is the full K dimension, not the partition size int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; - fp_qvm_impl( + fp_qvm_impl( w, scales, + nullptr, x, y, in_vec_size_adj, @@ -1411,12 +1483,13 @@ template < template < typename T, - const int group_size, - const int bits, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1465,10 +1538,11 @@ template < w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void fp_gather_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1510,14 +1584,24 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1559,14 +1643,24 @@ template w_strides, s_strides, tid); - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1608,9 +1702,10 @@ template w_strides, s_strides, tid); - fp_qvm_impl( + fp_qvm_impl( w, scales, + global_scale, x, y, in_vec_size, @@ -1741,11 +1836,12 @@ template < template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1996,39 +2092,48 @@ template < } } -template +template [[kernel]] void fp_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], device uint8_t* scales [[buffer(2)]], + const device float* global_scale [[buffer(3)]], uint2 tidx [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]], uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + scale_dec_b = simd_max(abs(w_thread)); } else { float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); - scale = simd_lid < 16 ? w_max_l : w_max_r; + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); + auto s = ScaleType(scale_dec_b); uint8_t q_scale = s.bits; - scale = float(s); - size_t gindex = index / group_size; if (index % group_size == 0) { scales[gindex] = q_scale; } - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = float(s); + scale_enc_b = (scale_enc_b == 0) ? 0.0f : (scale_enc / scale_enc_b); + uint8_t output = Quantize{}(w_thread * scale_enc_b); if (bits == 4) { uint8_t sval = simd_shuffle_down(output, 1); output |= sval << bits; @@ -2039,10 +2144,11 @@ template } } -template +template [[kernel]] void fp_dequantize( const device uint8_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], + const device float* global_scale [[buffer(2)]], device T* out [[buffer(3)]], uint2 index [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]]) { @@ -2054,9 +2160,17 @@ template out += oindex; + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + using ScaleType = metal::conditional_t; auto q_scale = ((device ScaleType*)(scales))[gindex]; auto scale = float(q_scale); + if constexpr (has_global_scale) { + scale *= inv_scale_enc; + } uint val = w[offset]; #pragma clang loop unroll(full) @@ -2071,32 +2185,41 @@ template } } -template +template [[kernel]] void fp_quantize_dequantize( const device T* w [[buffer(0)]], - device T* out [[buffer(1)]], + const device float* global_scale [[buffer(1)]], + device T* out [[buffer(2)]], uint2 tidx [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]], uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + scale_dec_b = simd_max(abs(w_thread)); } else { float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); - scale = simd_lid < 16 ? w_max_l : w_max_r; + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); - scale = float(s); + auto scale = float(ScaleType(scale_dec_b)); - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = (scale == 0) ? 0.0f : (scale_enc / scale); + uint8_t output = Quantize{}(w_thread * scale_enc_b); - out[index] = static_cast(scale * Dequantize{}(output)); + out[index] = static_cast((scale / scale_enc) * Dequantize{}(output)); } diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 76980164f4..7404f2023f 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -9,19 +9,34 @@ #define instantiate_quantized(mode, name, type, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits, \ - fp_ ## name, \ - type, \ - group_size, \ - bits) + fp_ ## name, \ + type, \ + group_size, \ + bits) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + true) #define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \ fp_ ## name, \ - type, \ - group_size, \ - bits, \ - batched) + type, \ + group_size, \ + bits, \ + batched) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + batched, \ + true) #define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \ instantiate_kernel( \ @@ -137,25 +152,28 @@ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits) -#define instantiate_quantize_dequantize(type, mode, group_size, bits) \ - instantiate_kernel( \ - #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits, \ +#define instantiate_quantize_dequantize(type, mode, group_size, bits, has_global_scale) \ + instantiate_kernel( \ + #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ fp_quantize_dequantize, \ - type, \ - group_size, \ - bits) \ - instantiate_kernel( \ - #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits, \ - fp_quantize, \ - type, \ - group_size, \ - bits) \ - instantiate_kernel( \ - #mode "_dequantize_" #type "_gs_" #group_size "_b_" #bits, \ - fp_dequantize, \ - type, \ - group_size, \ - bits) + type, \ + group_size, \ + bits, \ + has_global_scale) \ + instantiate_kernel( \ + #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ + fp_quantize, \ + type, \ + group_size, \ + bits, \ + has_global_scale) \ + instantiate_kernel( \ + #mode "_dequantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ + fp_dequantize, \ + type, \ + group_size, \ + bits, \ + has_global_scale) #define instantiate_quantized_modes(type, mode, group_size, bits) \ instantiate_quantized_all_batched(type, mode, group_size, bits) \ @@ -164,13 +182,16 @@ instantiate_quantized_all_wide(type, mode, group_size, bits) \ instantiate_quantized_all_splitk(type, mode, group_size, bits) \ instantiate_quantized_all_aligned(type, mode, group_size, bits) \ - instantiate_quantized_all_rhs(type, mode, group_size, bits) \ - instantiate_quantize_dequantize(type, mode, group_size, bits) + instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ instantiate_quantized_modes(type, nvfp4, 16, 4) \ instantiate_quantized_modes(type, mxfp8, 32, 8) \ - instantiate_quantized_modes(type, mxfp4, 32, 4) + instantiate_quantized_modes(type, mxfp4, 32, 4) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ + instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ + instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 1820ce3423..6d87dc770f 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -1591,7 +1591,12 @@ template quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void affine_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1643,7 +1648,12 @@ template simd_lid); } -template +template < + typename T, + int group_size, + const int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void affine_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1754,7 +1764,12 @@ template < simd_lid); } -template +template < + typename T, + const int group_size, + const int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void affine_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2001,12 +2016,13 @@ template < template < typename T, - const int group_size, - const int bits, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2059,7 +2075,7 @@ template < w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void affine_gather_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2121,7 +2137,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2183,7 +2199,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2330,11 +2346,12 @@ template < template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_gather_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2592,7 +2609,7 @@ template < } } -template +template [[kernel]] void affine_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], @@ -2697,7 +2714,7 @@ template } } -template +template [[kernel]] void affine_dequantize( const device uint8_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 94c563070a..57692dfcd5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -134,7 +134,7 @@ inline int add_strides_and_shapes( const std::optional& biases, int offset) { if (skip) { - return 0; + return offset; } // TODO: Collapse batch dimensions @@ -172,6 +172,209 @@ inline int add_gather_strides_and_shapes( return offset; } +auto get_quantize_kernel_dims( + MTL::ComputePipelineState* kernel, + const array& w, + const array& out, + int group_size, + int bits, + bool dequantize = false) { + // Treat uint32 as uint8 in kernel + constexpr int uint8_per_uint32 = 4; + constexpr int simd_size = 32; + int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; + int per_thread = + dequantize ? packs_per_int : std::max(group_size / simd_size, 1); + size_t nthreads = + dequantize ? out.size() / packs_per_int : w.size() / per_thread; + + NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + if (thread_group_size > nthreads) { + thread_group_size = nthreads; + } + auto group_dims = MTL::Size(thread_group_size, 1, 1); + bool use_2d = nthreads > UINT_MAX; + auto grid_shape = w.shape(); + if (dequantize) { + grid_shape.back() *= uint8_per_uint32; + } else { + grid_shape.back() /= per_thread; + } + MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) + : MTL::Size(nthreads, 1, 1); + return std::make_tuple(grid_dims, group_dims); +} + +void quantize_impl( + const std::vector& inputs, + std::vector& outputs, + QuantizationMode mode, + int group_size, + int bits, + bool dequantize, + Stream s) { + auto& w_pre = inputs[0]; + auto& out = outputs[0]; + out.set_data(allocator::malloc(out.nbytes())); + + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + bool has_biases = (mode == QuantizationMode::Affine); + bool has_global_scale = !has_biases && (inputs.size() > (1 + dequantize)); + + auto w = ensure_row_contiguous(w_pre, d, s); + if (dequantize) { + auto scales = ensure_row_contiguous(inputs[1], d, s); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (has_biases) { + auto biases = ensure_row_contiguous(inputs[2], d, s); + compute_encoder.set_input_array(biases, 2); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[2], 2); + } + compute_encoder.set_output_array(out, 3); + } else { + auto& scales = outputs[1]; + scales.set_data(allocator::malloc(scales.nbytes())); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_output_array(scales, 2); + if (has_biases) { + auto& biases = outputs[2]; + biases.set_data(allocator::malloc(biases.nbytes())); + compute_encoder.set_output_array(biases, 3); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[1], 3); + } + } + + auto type_string = dequantize ? get_type_string(out.dtype()) + : get_type_string(w_pre.dtype()); + auto mode_string = quantization_mode_to_string(mode); + std::string kname; + concatenate( + kname, + mode_string + (dequantize ? "_dequantize" : "_quantize"), + "_", + type_string, + "_gs_", + group_size, + "_b_", + bits); + if (!has_biases) { + concatenate(kname, "_hgs_", has_global_scale ? "true" : "false"); + } + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + dequantize ? "dequantize" : "quantize", + mode_string, + type_string, + group_size, + bits, + has_global_scale); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits, dequantize); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +auto quantize_input( + const array& w, + const std::optional& global_scale, + QuantizationMode mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + auto wq_shape = w.shape(); + wq_shape.back() = w.shape(-1) * bits / 32; + auto scales_shape = w.shape(); + scales_shape.back() = w.shape(-1) / group_size; + + std::vector inputs{w}; + if (global_scale) { + inputs.push_back(*global_scale); + } + std::vector outputs{ + array(wq_shape, uint32, nullptr, {}), + array(scales_shape, uint8, nullptr, {})}; + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.add_temporary(outputs[0]); + compute_encoder.add_temporary(outputs[1]); + quantize_impl(inputs, outputs, mode, group_size, bits, false, s); + return std::make_tuple(outputs[0], outputs[1]); +} + +void fp_quantize_dequantize( + const array& in, + const std::optional& global_scale, + array& out, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + const Stream& s) { + auto& compute_encoder = metal::get_command_encoder(s); + + auto w = ensure_row_contiguous(in, d, s); + compute_encoder.set_input_array(w, 0); + if (global_scale) { + compute_encoder.set_input_array(*global_scale, 1); + } + compute_encoder.set_output_array(out, 2); + auto type_string = get_type_string(in.dtype()); + std::string kname; + concatenate( + kname, + mode + "_quantize_dequantize_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_hgs_", + global_scale ? "true" : "false"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "quantize_dequantize", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +array quantize_dequantize_input( + const array& x_pre, + const std::optional& global_scale, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + bool donate_x = x_pre.is_donatable(); + array x = ensure_row_contiguous(x_pre, d, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = + donate_x ? x : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); + if (!donate_x) { + metal::get_command_encoder(s).add_temporary(xhat); + } + fp_quantize_dequantize(x, global_scale, xhat, mode, group_size, bits, d, s); + return xhat; +} + } // namespace void qmv_quad( @@ -237,6 +440,7 @@ void qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -266,7 +470,8 @@ void qmv( group_size, "_b_", bits, - B > 1 ? "_batch_1" : "_batch_0"); + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( d, kname, @@ -275,17 +480,20 @@ void qmv( type_string, group_size, bits, - B > 1); + B > 1, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(K, c++); @@ -511,6 +719,7 @@ void qvm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -539,18 +748,29 @@ void qvm( group_size, "_b_", bits, - B > 1 ? "_batch_1" : "_batch_0"); + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( - d, kname, "qvm", mode, type_string, group_size, bits, B > 1); + d, + kname, + "qvm", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(K, c++); @@ -1058,6 +1278,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -1087,7 +1308,8 @@ void gather_qmv( "_gs_", group_size, "_b_", - bits); + bits, + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( d, @@ -1096,17 +1318,20 @@ void gather_qmv( mode, type_string, group_size, - bits); + bits, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -1124,6 +1349,7 @@ void gather_qvm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -1153,18 +1379,28 @@ void gather_qvm( "_gs_", group_size, "_b_", - bits); + bits, + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( - d, kname, "gather_qvm", mode, type_string, group_size, bits); + d, + kname, + "gather_qvm", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -1463,6 +1699,7 @@ void dispatch_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -1473,18 +1710,31 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - if ((K == 128 || K == 64) && is_power_of_2(bits)) { + if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } // Small batch so route to qmv_wide, which reuses each weight group across the // M vectors. - if (M >= 2 && use_qmv_wide(mode, d)) { + if (M >= 2 && use_qmv_wide(mode, d) && !global_scale) { qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } - qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + qmv(x, + w, + scales, + biases, + global_scale, + out, + group_size, + bits, + M, + N, + K, + d, + s, + mode); } void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { @@ -1540,13 +1790,39 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { // Run of the mill qmv if (transpose_) { dispatch_qmv( - x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); return; } // Run of the mill qvm if (K < 1024) { - qvm(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + qvm(x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); return; } @@ -1632,6 +1908,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1651,6 +1928,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1664,194 +1942,133 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { mode); } -void quantize_dequantize( - const array& in, - array& out, - std::string mode, - int group_size, - int bits, - metal::Device& d, - const Stream& s) { - auto& compute_encoder = metal::get_command_encoder(s); - - auto w = ensure_row_contiguous(in, d, s); - compute_encoder.set_input_array(w, 0); - compute_encoder.set_output_array(out, 1); - auto type_string = get_type_string(in.dtype()); - std::string kname; - concatenate( - kname, - mode + "_quantize_dequantize_", - type_string, - "_gs_", - group_size, - "_b_", - bits); - auto kernel = get_quantized_kernel_wrapped( - d, kname, "quantize_dequantize", mode, type_string, group_size, bits); - - compute_encoder.set_compute_pipeline_state(kernel); - - constexpr int uint8_per_uint32 = 4; - constexpr int simd_size = 32; - int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; - int per_thread = std::max(group_size / simd_size, 1); - size_t nthreads = w.size() / per_thread; - - NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - if (thread_group_size > nthreads) { - thread_group_size = nthreads; - } - auto group_dims = MTL::Size(thread_group_size, 1, 1); - bool use_2d = nthreads > UINT_MAX; - auto grid_shape = w.shape(); - grid_shape.back() /= per_thread; - MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) - : MTL::Size(nthreads, 1, 1); - compute_encoder.dispatch_threads(grid_dims, group_dims); -} - void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; auto mode = quantization_mode_to_string(mode_); + + out.set_data(allocator::malloc(out.nbytes())); + + // - 2 inputs: x, w (non-quantized w) + // - 3 inputs: x, w, scales_w (quantized w) bool w_quantized = (inputs[1].dtype() == uint32); - // Tensor-scale nvfp4 (global_scale_x / global_scale_w) is packed into - // inputs by ops.cpp but no Metal qqmm kernel currently consumes the - // global scales. Reject the request rather than silently dropping them - // in the gemv path below. int base_size = w_quantized ? 3 : 2; - if (mode_ == QuantizationMode::Nvfp4 && - static_cast(inputs.size()) > base_size) { - throw std::runtime_error( - "[QQMatmul] Global scale (tensor-scale nvfp4) is not supported " - "on the Metal backend."); - } - if (w_quantized && inputs[0].shape(-2) == 1) { - out.set_data(allocator::malloc(out.nbytes())); - - bool donate_x = inputs[0].is_donatable(); - array x = ensure_row_contiguous(inputs[0], d, s); - // If x is a copy it should be donatable - donate_x |= x.is_donatable(); - auto xhat = donate_x - ? x - : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); - quantize_dequantize(x, xhat, mode, group_size_, bits_, d, s); - - // Make sure the last two dims of w and s are contiguous - array w = ensure_row_contiguous_matrix(inputs[1], d, s); - array scales = ensure_row_contiguous_matrix(inputs[2], d, s); - - bool non_batched = w.ndim() == 2; - int K = x.shape(-1); - int M = non_batched ? x.size() / K : x.shape(-2); - int N = out.shape(-1); - dispatch_qmv( - xhat, - w, - scales, - std::nullopt, - out, - group_size_, - bits_, - M, - N, - K, - d, - s, - mode); - return; - } else { - throw std::runtime_error("[QQMatmul] NYI for the general case"); + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; } -} -void fast::Quantize::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - auto& w_pre = inputs[0]; - auto& out = outputs[0]; - out.set_data(allocator::malloc(out.nbytes())); + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + dispatch_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} + +void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); - auto& compute_encoder = metal::get_command_encoder(s); - auto w = ensure_row_contiguous(w_pre, d, s); - if (dequantize_) { - auto scales = ensure_row_contiguous(inputs[1], d, s); - if (mode_ == QuantizationMode::Affine) { - auto biases = ensure_row_contiguous(inputs[2], d, s); - compute_encoder.set_input_array(biases, 2); - } - compute_encoder.set_input_array(w, 0); - compute_encoder.set_input_array(scales, 1); - compute_encoder.set_output_array(out, 3); - } else { - auto& scales = outputs[1]; - scales.set_data(allocator::malloc(scales.nbytes())); - if (mode_ == QuantizationMode::Affine) { - auto& biases = outputs[2]; - biases.set_data(allocator::malloc(biases.nbytes())); - compute_encoder.set_output_array(biases, 3); - } - compute_encoder.set_input_array(w, 0); - compute_encoder.set_output_array(out, 1); - compute_encoder.set_output_array(scales, 2); + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + const array& lhs_indices = ensure_row_contiguous(inputs[2], d, s); + const array& rhs_indices = ensure_row_contiguous(inputs[3], d, s); + auto mode = quantization_mode_to_string(mode_); + + out.set_data(allocator::malloc(out.nbytes())); + + // - 4 inputs: x, w (non-quantized w) + // - 5 inputs: x, w, scales_w (quantized w) + bool w_quantized = (inputs[1].dtype() == uint32); + int base_size = w_quantized ? 5 : 4; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; } - auto type_string = dequantize_ ? get_type_string(out.dtype()) - : get_type_string(w_pre.dtype()); - auto mode = quantization_mode_to_string(mode_); - std::string kname; - concatenate( - kname, - mode + (dequantize_ ? "_dequantize" : "_quantize"), - "_", - type_string, - "_gs_", - group_size_, - "_b_", - bits_); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - dequantize_ ? "dequantize" : "quantize", - mode, - type_string, - group_size_, - bits_); + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); - compute_encoder.set_compute_pipeline_state(kernel); + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); - // Treat uint32 as uint8 in kernel - constexpr int uint8_per_uint32 = 4; - constexpr int simd_size = 32; - int packs_per_int = (bits_ == 3 || bits_ == 5) ? 8 - : bits_ == 6 ? 4 - : 8 / bits_; - int per_thread = - dequantize_ ? packs_per_int : std::max(group_size_ / simd_size, 1); - size_t nthreads = - dequantize_ ? out.size() / packs_per_int : w.size() / per_thread; + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + gather_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} - NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - if (thread_group_size > nthreads) { - thread_group_size = nthreads; - } - auto group_dims = MTL::Size(thread_group_size, 1, 1); - bool use_2d = nthreads > UINT_MAX; - auto grid_shape = w.shape(); - if (dequantize_) { - grid_shape.back() *= uint8_per_uint32; - } else { - grid_shape.back() /= per_thread; - } - MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) - : MTL::Size(nthreads, 1, 1); - compute_encoder.dispatch_threads(grid_dims, group_dims); +void fast::Quantize::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + quantize_impl( + inputs, outputs, mode_, group_size_, bits_, dequantize_, stream()); } void fast::ConvertFP8::eval_gpu( diff --git a/mlx/backend/no_cpu/primitives.cpp b/mlx/backend/no_cpu/primitives.cpp index ae51dd9b2f..faaeb0c7c4 100644 --- a/mlx/backend/no_cpu/primitives.cpp +++ b/mlx/backend/no_cpu/primitives.cpp @@ -71,6 +71,7 @@ NO_CPU(Gather) NO_CPU(GatherAxis) NO_CPU(GatherMM) NO_CPU(GatherQMM) +NO_CPU(GatherQQMM) NO_CPU(Greater) NO_CPU(GreaterEqual) NO_CPU(Hadamard) diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 4819ed2724..0e05e9d19f 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -98,6 +98,7 @@ NO_GPU(Gather) NO_GPU(GatherAxis) NO_GPU(GatherMM) NO_GPU(GatherQMM) +NO_GPU(GatherQQMM) NO_GPU(Greater) NO_GPU(GreaterEqual) NO_GPU(Hadamard) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 0356ad0cae..323c6a9a6e 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -65,11 +65,26 @@ Dtype at_least_float(const Dtype& d) { } array indices_or_default( - std::optional indices, + std::string_view tag, + const std::optional& indices, const array& x, StreamOrDevice s) { + if (x.ndim() < 2) { + std::ostringstream msg; + msg << tag + << " Input must have at least two dimensions but got input with shape " + << x.shape() << "."; + throw std::invalid_argument(msg.str()); + } + if (indices.has_value()) { - return indices.value(); + if (!issubdtype(indices->dtype(), integer)) { + std::ostringstream msg; + msg << tag + << " Got indices with invalid dtype. Indices must be integral."; + throw std::invalid_argument(msg.str()); + } + return astype(indices.value(), uint32); } Shape shape(x.shape().begin(), x.shape().end() - 2); @@ -4637,10 +4652,10 @@ void validate_global_scale( } array quantized_matmul( - array x, - array w, - array scales, - std::optional biases /* = std::nullopt */, + const array& x, + const array& w, + const array& scales, + const std::optional& biases /* = std::nullopt */, bool transpose /* = true */, std::optional group_size_ /* = std::nullopt */, std::optional bits_ /* = std::nullopt */, @@ -4689,13 +4704,13 @@ array quantized_matmul( } void validate_qqmm_inputs( - array x, - array w, - std::optional scales_w, + const array& x, + const array& w, + const std::optional& scales_w, int group_size, int bits, - std::optional global_scale_x, - std::optional global_scale_w, + const std::optional& global_scale_x, + const std::optional& global_scale_w, QuantizationMode qmode) { // check 2D (for now) if (x.ndim() > 2 || w.ndim() > 2) { @@ -4749,9 +4764,9 @@ void validate_qqmm_inputs( } std::pair extract_qqmm_dims( - array x, - array w, - std::optional scales_w, + const array& x, + const array& w, + const std::optional& scales_w, int group_size, int bits) { if (w.dtype() != uint32) { @@ -4779,24 +4794,17 @@ std::pair extract_qqmm_dims( } array qqmm( - array in_x, - array w, - std::optional scales_w, + const array& in_x, + const array& w, + const std::optional& scales_w, std::optional group_size_ /* = std::nullopt */, std::optional bits_ /* = std::nullopt */, const std::string& mode /* = "nvfp4" */, - const std::optional global_scale_x /* = std::nullopt */, - const std::optional global_scale_w /* = std::nullopt */, + const std::optional& global_scale_x /* = std::nullopt */, + const std::optional& global_scale_w /* = std::nullopt */, StreamOrDevice s /* = {} */) { auto stream = to_stream(s); auto qmode = string_to_quantization_mode(mode, "qqmm"); - // cuBLAS block scaled matmul only supports nvfp4 and mxfp8 - if (qmode != QuantizationMode::Nvfp4 && qmode != QuantizationMode::Mxfp8) { - std::ostringstream msg; - msg << "[qqmm] Only 'nvfp4' and 'mxfp8' quantization modes are supported but '" - << mode << "' was provided."; - throw std::invalid_argument(msg.str()); - } // we need to check 2 cases: // 1. w is quantized, scales is provided // 2. w is not quantized, scales is not provided @@ -5099,12 +5107,6 @@ std::vector quantize( << " matrix has shape " << w.shape(); throw std::invalid_argument(msg.str()); } - if (to_stream(s).device == Device::gpu && metal::is_available() && - global_scale.has_value()) { - std::ostringstream msg; - msg << "[quantize] Global scale is not supported on the Metal backend."; - throw std::invalid_argument(msg.str()); - } validate_global_scale("quantize", qmode, global_scale); if (qmode == QuantizationMode::Affine) { return affine_quantize(w, group_size, bits, s); @@ -5364,13 +5366,6 @@ array dequantize( << "but it has only " << w.ndim() << "."; throw std::invalid_argument(msg.str()); } - if (global_scale.has_value()) { - if (to_stream(s).device == Device::gpu && metal::is_available()) { - std::ostringstream msg; - msg << "[dequantize] Global scale is not supported on the Metal backend."; - throw std::invalid_argument(msg.str()); - } - } validate_global_scale("dequantize", qmode, global_scale); if (qmode == QuantizationMode::Affine) { @@ -5463,30 +5458,11 @@ array gather_qmm( } // Extract indices and broadcast them - array lhs_indices = indices_or_default(lhs_indices_, x, s); - array rhs_indices = indices_or_default(rhs_indices_, w, s); + array lhs_indices = indices_or_default("[gather_qmm]", lhs_indices_, x, s); + array rhs_indices = indices_or_default("[gather_qmm]", rhs_indices_, w, s); std::tie(lhs_indices, rhs_indices) = broadcast_arrays(lhs_indices, rhs_indices, s); - if (!issubdtype(lhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_qmm] Got lhs_indices with invalid dtype. Indices must be integral."); - } - - if (!issubdtype(rhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_qmm] Got rhs_indices with invalid dtype. Indices must be integral."); - } - if (x.ndim() < 2) { - std::ostringstream msg; - msg << "[gather_qmm] Non-quantized input must have at least two" - << " dimensions but got input with shape " << x.shape() << "."; - throw std::invalid_argument(msg.str()); - } - - lhs_indices = astype(lhs_indices, uint32, s); - rhs_indices = astype(rhs_indices, uint32, s); - // Compute the full output shape auto out_shape = lhs_indices.shape(); out_shape.push_back(x.shape(-2)); @@ -5522,6 +5498,56 @@ array gather_qmm( std::move(inputs)); } +array gather_qqmm( + const array& x, + const array& w, + const std::optional& scales_w, + const std::optional& lhs_indices_, + const std::optional& rhs_indices_, + std::optional group_size_, + std::optional bits_, + const std::string& mode, + const std::optional& global_scale_x, + const std::optional& global_scale_w, + bool sorted_indices, + StreamOrDevice s) { + auto stream = to_stream(s); + auto qmode = string_to_quantization_mode(mode, "gather_qqmm"); + auto [group_size, bits] = + quantization_params_from_mode(qmode, group_size_, bits_); + + // Extract indices and broadcast them + array lhs_indices = indices_or_default("[gather_qqmm]", lhs_indices_, x, s); + array rhs_indices = indices_or_default("[gather_qqmm]", rhs_indices_, w, s); + std::tie(lhs_indices, rhs_indices) = + broadcast_arrays(lhs_indices, rhs_indices, s); + + std::vector inputs = { + x, + w, + lhs_indices, + rhs_indices, + }; + if (scales_w.has_value()) { + inputs.push_back(*scales_w); + } + if (global_scale_x.has_value() && global_scale_w.has_value()) { + inputs.push_back(*global_scale_x); + inputs.push_back(*global_scale_w); + } + + auto [w_inner_dims, w_outer_dims] = + extract_qqmm_dims(x, w, scales_w, group_size, bits); + auto out_shape = lhs_indices.shape(); + out_shape.push_back(x.shape(-2)); + out_shape.push_back(w_outer_dims); + return array( + std::move(out_shape), + x.dtype(), + std::make_shared(stream, group_size, bits, qmode), + std::move(inputs)); +} + array tensordot( const array& a, const array& b, @@ -6020,28 +6046,14 @@ array gather_mm( b = astype(b, out_type, s); // Handle broadcasting - array lhs_indices = indices_or_default(lhs_indices_, a, s); - array rhs_indices = indices_or_default(rhs_indices_, b, s); - - if (!issubdtype(lhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_mm] Got lhs_indices with invalid dtype. Indices must be integral."); - } - - if (!issubdtype(rhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_mm] Got rhs_indices with invalid dtype. Indices must be integral."); - } - - lhs_indices = astype(lhs_indices, uint32, s); - rhs_indices = astype(rhs_indices, uint32, s); + array lhs_indices = indices_or_default("[gather_mm]", lhs_indices_, a, s); + array rhs_indices = indices_or_default("[gather_mm]", rhs_indices_, b, s); + std::tie(lhs_indices, rhs_indices) = + broadcast_arrays(lhs_indices, rhs_indices, s); int M = a.shape(-2); int N = b.shape(-1); - std::tie(lhs_indices, rhs_indices) = - broadcast_arrays(lhs_indices, rhs_indices, s); - auto out_shape = lhs_indices.shape(); out_shape.push_back(M); out_shape.push_back(N); diff --git a/mlx/ops.h b/mlx/ops.h index ed617c3441..20568024f6 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -1546,10 +1546,10 @@ MLX_API array conv_transpose3d( /** Quantized matmul multiplies x with a quantized matrix w*/ MLX_API array quantized_matmul( - array x, - array w, - array scales, - std::optional biases = std::nullopt, + const array& x, + const array& w, + const array& scales, + const std::optional& biases = std::nullopt, bool transpose = true, std::optional group_size = std::nullopt, std::optional bits = std::nullopt, @@ -1578,15 +1578,15 @@ MLX_API array dequantize( StreamOrDevice s = {}); MLX_API array qqmm( - array x, // input activations - array w, // maybe quantized weights - const std::optional w_scales = std::nullopt, // optional scales if w - // is quantized + const array& x, // input activations + const array& w, // maybe quantized weights + const std::optional& w_scales = std::nullopt, // optional scales if w + // is quantized std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "nvfp4", - const std::optional global_scale_x = std::nullopt, - const std::optional global_scale_w = std::nullopt, + const std::optional& global_scale_x = std::nullopt, + const std::optional& global_scale_w = std::nullopt, StreamOrDevice s = {}); /** Convert an E4M3 float8 to the given floating point dtype. */ @@ -1610,6 +1610,20 @@ MLX_API array gather_qmm( bool sorted_indices = false, StreamOrDevice s = {}); +MLX_API array gather_qqmm( + const array& x, + const array& w, + const std::optional& scales_w = std::nullopt, + const std::optional& lhs_indices = std::nullopt, + const std::optional& rhs_indices = std::nullopt, + std::optional group_size = std::nullopt, + std::optional bits = std::nullopt, + const std::string& mode = "nvfp4", + const std::optional& global_scale_x = std::nullopt, + const std::optional& global_scale_w = std::nullopt, + bool sorted_indices = false, + StreamOrDevice s = {}); + /** Returns a contraction of a and b over multiple dimensions. */ MLX_API array tensordot( const array& a, diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index ae434af94d..67afad54b0 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3827,6 +3827,22 @@ std::vector GatherQMM::output_shapes(const std::vector& inputs) { return {out_shape}; } +bool GatherQQMM::is_equivalent(const Primitive& other) const { + const GatherQQMM& qm_other = static_cast(other); + return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ && + mode_ == qm_other.mode_; +} + +std::vector GatherQQMM::output_shapes(const std::vector& inputs) { + const auto& x = inputs[0]; + const auto& w = inputs[1]; + const auto& lhs_indices = inputs[2]; + auto out_shape = lhs_indices.shape(); + out_shape.push_back(x.shape(-2)); + out_shape.push_back(w.shape(-2)); + return {out_shape}; +} + std::pair, std::vector> RandomBits::vmap( const std::vector& inputs, const std::vector& axes) { diff --git a/mlx/primitives.h b/mlx/primitives.h index 5b8517c56d..403490dbe8 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -1716,6 +1716,41 @@ class GatherQMM : public UnaryPrimitive { bool right_sorted_; }; +class GatherQQMM : public UnaryPrimitive { + public: + explicit GatherQQMM( + Stream stream, + int group_size, + int bits, + QuantizationMode mode, + bool left_sorted = false, + bool right_sorted = false) + : UnaryPrimitive(stream), + group_size_(group_size), + bits_(bits), + mode_(mode), + left_sorted_(left_sorted), + right_sorted_(right_sorted) {} + + void eval_cpu(const std::vector& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_NAME(GatherQQMM) + bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; + auto state() const { + return std::make_tuple( + group_size_, bits_, mode_, left_sorted_, right_sorted_); + } + + private: + int group_size_; + int bits_; + QuantizationMode mode_; + bool left_sorted_; + bool right_sorted_; +}; + class RandomBits : public UnaryPrimitive { public: explicit RandomBits(Stream stream, const Shape& shape, int width) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 5658a78648..976eaae264 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4801,6 +4801,58 @@ void init_ops(nb::module_& m) { array: The result of the multiplication of ``x`` with ``w`` after gathering using ``lhs_indices`` and ``rhs_indices``. )pbdoc"); + m.def( + "gather_qqmm", + &mx::gather_qqmm, + nb::arg(), + nb::arg(), + "scales"_a = nb::none(), + "lhs_indices"_a = nb::none(), + "rhs_indices"_a = nb::none(), + "group_size"_a = nb::none(), + "bits"_a = nb::none(), + "mode"_a = "nvfp4", + "global_scale_x"_a = nb::none(), + "global_scale_w"_a = nb::none(), + nb::kw_only(), + "sorted_indices"_a = false, + "stream"_a = nb::none(), + nb::sig( + "def gather_qqmm(x: array, w: array, /, scales: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + R"pbdoc( + Fused :func:`qqmm` with matrix-level gather. + + Similar to :func:`gather_mm`, the indices ``lhs_indices`` and + ``rhs_indices`` contain flat indices along the batch dimensions (i.e. + all but the last two dimensions) of ``x`` and ``w`` respectively. + + Args: + x (array): Input array. + w (array): Weight matrix. If quantized, it is packed in unsigned integers. + scales (array, optional): The scales to use per ``group_size`` elements of + ``w`` if ``w`` is quantized. Default: ``None``. + lhs_indices (array, optional): Integer indices for ``x``. Default: ``None``. + rhs_indices (array, optional): Integer indices for ``w``. Default: ``None``. + group_size (int, optional): Number of elements in ``x`` and ``w`` that + share a scale. See supported values and defaults in the + :ref:`table of quantization modes `. Default: ``None``. + bits (int, optional): Number of bits used to represent each element of + ``x`` and ``w``. See supported values and defaults in the + :ref:`table of quantization modes `. Default: ``None``. + mode (str, optional): The quantization mode. Default: ``"nvfp4"``. + Supported modes are ``nvfp4`` and ``mxfp8``. See the + :ref:`table of quantization modes ` for details. + global_scale_x (array, optional): The per-input float32 scale used for x + with ``"nvfp4"`` quantization. Default: ``None``. + global_scale_w (array, optional): The per-input float32 scale used for w + with ``"nvfp4"`` quantization. Default: ``None``. + sorted_indices (bool, optional): May allow a faster implementation + if the passed indices are sorted. Default: ``False``. + + Returns: + array: The result of the multiplication of quantized ``x`` with quantized ``w``. + needed). + )pbdoc"); m.def( "segmented_mm", &mx::segmented_mm, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 7fbbe6938e..63254ee9c7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -180,11 +180,7 @@ def test_nvfp4_quantize_dequantize(self): self.assertTrue(mx.all(w_hat == 0)) # Test nvfp4 quantize/dequantize with tensor-scale global_scale - # currently supported only on cpu and cuda - if not mx.metal.is_available(): - global_scale = w.abs().max().astype(mx.float32) - else: - global_scale = None + global_scale = w.abs().max().astype(mx.float32) w_q, scales = mx.quantize(w, mode="nvfp4", global_scale=global_scale) w_hat = mx.dequantize( @@ -197,7 +193,7 @@ def test_qqmv(self): k1, k2 = mx.random.split(key) tests = product( [256, 512, 67], # M - [64, 256], # N + [64, 256, 512], # N ["nvfp4", "mxfp8"], # mode ) for M, N, mode in tests: @@ -205,12 +201,8 @@ def test_qqmv(self): x_shape = (1, N) w_shape = (M, N) - # TODO: Fix qmv with global scale in Metal/CPU backends. - has_global_scale = ( - mode == "nvfp4" - and mx.cuda.is_available() - and mx.default_device() == mx.gpu - ) + # TODO: Fix qmv with global scale in CPU backend. + has_global_scale = mode == "nvfp4" and mx.default_device() == mx.gpu x = mx.random.normal(shape=x_shape, key=k1) global_scale_x = mx.max(mx.abs(x)) if has_global_scale else None @@ -243,31 +235,54 @@ def test_qqmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) - def test_qqmm_metal_global_scale_rejected(self): - # Tensor-scale nvfp4 (global_scale_x / global_scale_w) is not - # implemented in the Metal qqmm kernels. mx.qqmm must reject the - # request on Metal rather than silently dropping the global scales - # in the gemv path and producing incorrect results. - if not mx.metal.is_available(): + def test_qqmm(self): + if mx.default_device() == mx.cpu: + self.skipTest("Not implemented for CPU") return - w = mx.random.normal(shape=(64, 64)) - w_q, scales = mx.quantize(w, mode="nvfp4") - x = mx.random.normal(shape=(1, 64)) - gx = mx.array(1.0, dtype=mx.float32) - gw = mx.array(1.0, dtype=mx.float32) + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + tests = product( + [8, 32, 33, 64], # M + [128, 256], # N + [128, 256], # K + ["nvfp4", "mxfp8"], # mode + ) + for M, N, K, mode in tests: + with self.subTest(shape=(M, N, K), mode=mode): + x_shape = (M, K) + w_shape = (N, K) - with self.assertRaises(RuntimeError): - y = mx.qqmm( - x, - w_q, - scales, - mode="nvfp4", - global_scale_x=gx, - global_scale_w=gw, - stream=mx.gpu, - ) - mx.eval(y) + x = mx.random.normal(shape=x_shape, key=k1) + global_scale_x = mx.max(mx.abs(x)) if mode == "nvfp4" else None + x_hat = mx.dequantize( + *mx.quantize(x, mode=mode, global_scale=global_scale_x), + mode=mode, + dtype=mx.float32, + global_scale=global_scale_x, + ) + + w = mx.random.normal(shape=w_shape, key=k2) + global_scale_w = mx.max(mx.abs(w)) if mode == "nvfp4" else None + w_q, scales = mx.quantize(w, mode=mode, global_scale=global_scale_w) + w_hat = mx.dequantize( + w_q, + scales, + mode=mode, + global_scale=global_scale_w, + dtype=mx.float32, + ) + y_q = mx.qqmm( + x, + w_q, + scales, + mode=mode, + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + y_hat = x_hat @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) def test_qmm(self): key = mx.random.key(0) @@ -1091,6 +1106,100 @@ def test_shape( test_shape(32, 512, 32, transpose=False, **kwargs) test_shape(1, 512, 32, transpose=False, **kwargs) + def test_gather_qqmm(self): + if mx.default_device() == mx.cpu: + self.skipTest("Not implemented for CPU") + return + + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + batches = ( + { + "batch_A": (1,), + "lhs_indices": (0,), + "batch_B": (3,), + "rhs_indices": (2, 1), + }, + { + "batch_A": (1,), + "lhs_indices": None, + "batch_B": (3,), + "rhs_indices": (2, 1), + }, + { + "batch_A": (2,), + "lhs_indices": None, + "batch_B": (3,), + "rhs_indices": (2, 1), + }, + { + "batch_A": (3,), + "lhs_indices": (0, 2), + "batch_B": (1,), + "rhs_indices": (0,), + }, + { + "batch_A": (5,), + "lhs_indices": (0, 2), + "batch_B": (3,), + "rhs_indices": (2, 1), + }, + ) + tests = product( + batches, + [1, 32], # M + [32, 256], # N + [32, 256], # K + ["nvfp4", "mxfp8"], # mode + ) + + for batch, M, N, K, mode in tests: + with self.subTest(shape=(M, N, K), mode=mode, **batch): + batch_A, lhs_indices, batch_B, rhs_indices = batch.values() + x_shape = (*batch_A, M, K) + w_shape = (*batch_B, N, K) + + x = mx.random.normal(shape=x_shape, key=k1) + global_scale_x = mx.max(mx.abs(x)) if mode == "nvfp4" else None + x_hat = mx.dequantize( + *mx.quantize(x, mode=mode, global_scale=global_scale_x), + mode=mode, + dtype=mx.float32, + global_scale=global_scale_x, + ) + + w = mx.random.normal(shape=w_shape, key=k2) + global_scale_w = mx.max(mx.abs(w)) if mode == "nvfp4" else None + w_q, scales = mx.quantize(w, mode=mode, global_scale=global_scale_w) + w_hat = mx.dequantize( + w_q, + scales, + mode=mode, + global_scale=global_scale_w, + dtype=mx.float32, + ) + + if lhs_indices is not None: + lhs_indices = mx.array(lhs_indices) + if rhs_indices is not None: + rhs_indices = mx.array(rhs_indices) + + y_q = mx.gather_qqmm( + x, + w_q, + scales, + lhs_indices, + rhs_indices, + mode=mode, + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + y_hat = mx.gather_mm( + x_hat, mx.swapaxes(w_hat, -1, -2), lhs_indices, rhs_indices + ) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmm_fp_type(self): indices = mx.array([[2], [0], [1]], dtype=mx.uint32) From bd813fbda132b30daa05bb4d8ed0ac56cbb67f07 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 17:09:17 -0700 Subject: [PATCH 048/222] docs: Fix missing printoptions page (#3985) --- docs/src/conf.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/src/conf.py b/docs/src/conf.py index 446d95cd20..b95a3fe74c 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -28,7 +28,10 @@ python_use_unqualified_type_names = True autosummary_generate = True -autosummary_filename_map = {"mlx.core.Stream": "stream_class"} +autosummary_filename_map = { + "mlx.core.Stream": "stream_class", + "mlx.core.PrintOptions": "printoptions_class", +} intersphinx_mapping = { "python": ("https://docs.python.org/3", None), From a681f8bf2b8660b4fc20c7e9ad6ece810ca570ed Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 17:29:19 -0700 Subject: [PATCH 049/222] docs: ThreadLocalStream and iinfo (#3986) --- docs/src/python/data_types.rst | 1 + docs/src/python/devices_and_streams.rst | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/src/python/data_types.rst b/docs/src/python/data_types.rst index 18abd235f7..9793648d88 100644 --- a/docs/src/python/data_types.rst +++ b/docs/src/python/data_types.rst @@ -76,3 +76,4 @@ documentation for more information. Use :func:`issubdtype` to determine if one DtypeCategory issubdtype finfo + iinfo diff --git a/docs/src/python/devices_and_streams.rst b/docs/src/python/devices_and_streams.rst index 843f122c98..7743cd5d0f 100644 --- a/docs/src/python/devices_and_streams.rst +++ b/docs/src/python/devices_and_streams.rst @@ -10,6 +10,7 @@ Devices and Streams Device Stream + ThreadLocalStream default_device set_default_device default_stream From 0268c81b62b135dc99773565016edf2b35a67898 Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Wed, 5 Aug 2026 05:59:42 +0530 Subject: [PATCH 050/222] Make "stop" optional in arange (#3982) --- python/src/ops.cpp | 16 ++++++++++------ python/tests/test_ops.py | 4 ++++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 976eaae264..7057075c9b 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1473,21 +1473,25 @@ void init_ops(nb::module_& m) { m.def( "arange", [](Scalar start, - Scalar stop, + std::optional stop, const std::optional& step, const std::optional& dtype_, mx::StreamOrDevice s) { + if (!stop) { + stop = start; + start = 0; + } // Determine the final dtype based on input types mx::Dtype dtype = dtype_ ? *dtype_ : mx::promote_types( scalar_to_dtype(start), step ? mx::promote_types( - scalar_to_dtype(stop), scalar_to_dtype(*step)) - : scalar_to_dtype(stop)); + scalar_to_dtype(*stop), scalar_to_dtype(*step)) + : scalar_to_dtype(*stop)); return mx::arange( scalar_to_double(start), - scalar_to_double(stop), + scalar_to_double(*stop), step ? scalar_to_double(*step) : 1.0, dtype, s); @@ -1499,7 +1503,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arange(start : Union[int, float], stop : Union[int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arange(start : Union[int, float], stop : Union[None, int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( Generates ranges of numbers. @@ -1508,7 +1512,7 @@ void init_ops(nb::module_& m) { Args: start (float or int, optional): Starting value which defaults to ``0``. - stop (float or int): Stopping value. + stop (float or int, optional): Stopping value. step (float or int, optional): Increment which defaults to ``1``. dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 86d92039f0..a0dcc0689e 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1501,6 +1501,10 @@ def test_arange_overload_dispatch(self): expected = [0, -1, -2] self.assertListEqual(a.tolist(), expected) + a = mx.arange(-3, None, -1) + expected = [0, -1, -2] + self.assertListEqual(a.tolist(), expected) + a = mx.arange(stop=2, step=0.5) expected = [0, 0.5, 1.0, 1.5] self.assertListEqual(a.tolist(), expected) From 5391a8e0522e41175f53efa58aa6765e22048bad Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 19:06:47 -0700 Subject: [PATCH 051/222] docs: Add softsign to the nn functions (#3989) --- docs/src/python/nn/functions.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/python/nn/functions.rst b/docs/src/python/nn/functions.rst index d3c533d0cb..68849e2206 100644 --- a/docs/src/python/nn/functions.rst +++ b/docs/src/python/nn/functions.rst @@ -36,5 +36,6 @@ simple functions. softmin softplus softshrink + softsign step tanh From 7c293d182976141e44926c741fd353841018f9ec Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 19:22:35 -0700 Subject: [PATCH 052/222] docs: Fix broken all_sum and Group references (#3996) --- docs/src/examples/data_parallelism.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/src/examples/data_parallelism.rst b/docs/src/examples/data_parallelism.rst index b1592611d6..9fb4ceba55 100644 --- a/docs/src/examples/data_parallelism.rst +++ b/docs/src/examples/data_parallelism.rst @@ -34,7 +34,8 @@ dataset, and optimizer initialization. mx.eval(loss, model.parameters()) All we have to do to average the gradients across machines is perform an -:func:`all_sum` and divide by the size of the :class:`Group`. Namely we +:func:`mlx.core.distributed.all_sum` and divide by the size of the +:class:`mlx.core.distributed.Group`. Namely we have to :func:`mlx.utils.tree_map` the gradients with following function. .. code:: python From f99e1dc03302d5450e0c6d65d3e33561c6b8e7f0 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 19:23:14 -0700 Subject: [PATCH 053/222] docs: Fix ast.metal_kernel typo (#3997) --- docs/src/dev/custom_metal_kernels.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/src/dev/custom_metal_kernels.rst b/docs/src/dev/custom_metal_kernels.rst index e907190cbf..6a9471fd51 100644 --- a/docs/src/dev/custom_metal_kernels.rst +++ b/docs/src/dev/custom_metal_kernels.rst @@ -92,8 +92,8 @@ function. This means we will launch ``mx.prod(grid)`` threads, subdivided into ``threadgroup`` size threadgroups. For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension. -Passing ``verbose=True`` to :func:`ast.metal_kernel.__call__` will print the -generated code for debugging purposes. +Passing ``verbose=True`` when calling the kernel returned by +:func:`fast.metal_kernel` will print the generated code for debugging purposes. Math Mode --------- From 391e140e626a53f3381621c5eec64524368746ff Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 4 Aug 2026 19:24:14 -0700 Subject: [PATCH 054/222] chore: Fix unresolved mx.array docstring (#3990) --- python/mlx/nn/layers/base.py | 2 +- python/mlx/nn/losses.py | 6 +++--- python/mlx/optimizers/optimizers.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index c2bcf7ff36..03b8be7a51 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -129,7 +129,7 @@ def load_weights( Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list. Args: - file_or_weights (str or list(tuple(str, mx.array))): The path to + file_or_weights (str or list(tuple(str, array))): The path to the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list of pairs of parameter names and arrays. strict (bool, optional): If ``True`` then checks that the provided diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 9218691e14..9226c10eee 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -540,8 +540,8 @@ def cosine_similarity_loss( \frac{x_1 \cdot x_2}{\max(\|x_1\| \cdot \|x_2\|, \epsilon)} Args: - x1 (mx.array): The first set of inputs. - x2 (mx.array): The second set of inputs. + x1 (array): The first set of inputs. + x2 (array): The second set of inputs. axis (int, optional): The embedding axis. Default: ``1``. eps (float, optional): The minimum value of the denominator used for numerical stability. Default: ``1e-8``. @@ -549,7 +549,7 @@ def cosine_similarity_loss( ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: - mx.array: The computed cosine similarity loss. + array: The computed cosine similarity loss. """ x1_norm = mx.linalg.norm(x1, axis=axis) x2_norm = mx.linalg.norm(x2, axis=axis) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index ca3fbc395a..3f167fd36b 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/python/mlx/optimizers/optimizers.py @@ -77,7 +77,7 @@ def init_single(self, parameter: mx.array, state: dict): state initialization. Args: - parameter (mx.array): A single parameter that will be optimized. + parameter (array): A single parameter that will be optimized. state (dict): The optimizer's state. """ raise NotImplementedError() @@ -112,8 +112,8 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): """To be extended by derived classes to implement the optimizer's update. Args: - gradient (mx.array): The ``parameter`` gradient. - parameter (mx.array): The ``parameter`` to update. + gradient (array): The ``parameter`` gradient. + parameter (array): The ``parameter`` to update. state (dict): The optimizer's state. """ raise NotImplementedError() From 36e891e9e9c2fe41c61ecc8cc6d351ca32befc39 Mon Sep 17 00:00:00 2001 From: Tanish Jain Date: Wed, 5 Aug 2026 09:25:01 +0530 Subject: [PATCH 055/222] python: fix bfloat16 buffer format itemsize mismatch (#3975) --- python/src/array.cpp | 7 +++++++ python/src/buffer.h | 2 +- python/tests/test_array.py | 7 +++---- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/python/src/array.cpp b/python/src/array.cpp index 1fcbbd25f0..a5a47dcec1 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -539,6 +539,13 @@ void init_array(nb::module_& m) { return nb::make_tuple(1, 0); } }) + .def( + "__array__", + [](const mx::array& self, nb::object dtype, nb::object copy) { + return mlx_to_np_array(self); + }, + "dtype"_a = nb::none(), + "copy"_a = nb::none()) .def("__copy__", [](const mx::array& self) { return mx::array(self); }) .def( "__deepcopy__", diff --git a/python/src/buffer.h b/python/src/buffer.h index 272a918883..4b194b3d25 100644 --- a/python/src/buffer.h +++ b/python/src/buffer.h @@ -43,7 +43,7 @@ std::string buffer_format(const mx::array& a) { case mx::float32: return "f"; case mx::bfloat16: - return "B"; + return "bfloat16"; case mx::float64: return "d"; case mx::complex64: diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 7279fa6276..cfd96e9236 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1868,11 +1868,10 @@ def test_buffer_protocol(self): mv_mx = memoryview(a_mx) self.assertEqual(mv_mx.strides, (8, 2)) self.assertEqual(mv_mx.shape, (3, 4)) - self.assertEqual(mv_mx.format, "B") - with self.assertRaises(RuntimeError) as cm: + self.assertIn(mv_mx.format, "bfloat16") + with self.assertRaises(ValueError) as cm: np.array(a_mx) - e = cm.exception - self.assertTrue("Item size 2 for PEP 3118 buffer format string" in str(e)) + self.assertIn("bfloat16", str(cm.exception)) # Test buffer protocol with non-arrays ie bytes a = ord("a") * 257 + mx.arange(10).astype(mx.int16) From c4111f8c3e98a188c15b4deb5da8a6c215c3e888 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:52:00 -0700 Subject: [PATCH 056/222] Template Metal complex scalar lanes (#3970) --- mlx/backend/metal/kernels/complex.h | 246 ++++++++++++++++++++-------- mlx/backend/metal/kernels/utils.h | 16 +- 2 files changed, 184 insertions(+), 78 deletions(-) diff --git a/mlx/backend/metal/kernels/complex.h b/mlx/backend/metal/kernels/complex.h index 5f654332b3..06a3050e8a 100644 --- a/mlx/backend/metal/kernels/complex.h +++ b/mlx/backend/metal/kernels/complex.h @@ -4,164 +4,263 @@ #include +#include "mlx/backend/metal/kernels/bf16.h" + using namespace metal; -struct complex64_t; +template +struct complex_t; template -static constexpr constant bool can_convert_to_complex64 = - !is_same_v && is_convertible_v; +static constexpr constant bool is_complex_v = false; + +template +static constexpr constant bool is_complex_v> = true; + +// Metal accepts explicit bfloat casts that is_convertible_v reports as false. +template +static constexpr constant bool is_lane_convertible_v = + is_convertible_v || + (is_same_v && is_convertible_v) || + (is_same_v && is_convertible_v); template -static constexpr constant bool can_convert_from_complex64 = - !is_same_v && - (is_convertible_v || is_convertible_v); +struct complex_t { + using value_type = T; -struct complex64_t { - float real; - float imag; + T real; + T imag; // Constructors - constexpr complex64_t(float real, float imag) thread : real(real), - imag(imag) {}; - constexpr complex64_t() thread : real(0), imag(0) {}; - constexpr complex64_t() threadgroup : real(0), imag(0) {}; + constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {}; + constexpr complex_t() thread : real(0), imag(0) {}; + constexpr complex_t() threadgroup : real(0), imag(0) {}; + + // Conversions from scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) thread : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) threadgroup : real(static_cast(x)), + imag(static_cast(0)) {} - // Conversions to complex64_t template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) thread : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) device : real(static_cast(x)), + imag(static_cast(0)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) threadgroup : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) constant : real(static_cast(x)), + imag(static_cast(0)) {} + // Conversions between complex types template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) device : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) thread : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) constant : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) threadgroup + : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} - // Conversions from complex64_t template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const thread { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) device : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + // Conversions to scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const thread { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const threadgroup { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const threadgroup { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const device { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const device { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const constant { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const constant { + return static_cast(real); } }; -constexpr complex64_t operator-(complex64_t x) { +using complex64_t = complex_t; + +static_assert(sizeof(complex64_t) == 2 * sizeof(float)); +static_assert(sizeof(complex_t) == 2 * sizeof(half)); +static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); + +template +constexpr complex_t operator-(complex_t x) { return {-x.real, -x.imag}; } -constexpr bool operator>=(complex64_t a, complex64_t b) { +template +constexpr bool operator>=(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); } -constexpr bool operator>(complex64_t a, complex64_t b) { +template +constexpr bool operator>(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); } -constexpr bool operator<=(complex64_t a, complex64_t b) { +template +constexpr bool operator<=(complex_t a, complex_t b) { return operator>=(b, a); } -constexpr bool operator<(complex64_t a, complex64_t b) { +template +constexpr bool operator<(complex_t a, complex_t b) { return operator>(b, a); } -constexpr bool operator==(complex64_t a, complex64_t b) { +template +constexpr bool operator==(complex_t a, complex_t b) { return a.real == b.real && a.imag == b.imag; } -constexpr complex64_t operator+(complex64_t a, complex64_t b) { +template +constexpr complex_t operator+(complex_t a, complex_t b) { return {a.real + b.real, a.imag + b.imag}; } -constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) { +template +constexpr thread complex_t& operator+=( + thread complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr threadgroup complex64_t& operator+=( - threadgroup complex64_t& a, - complex64_t b) { +template +constexpr threadgroup complex_t& operator+=( + threadgroup complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) { +template +constexpr device complex_t& operator+=( + device complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr complex64_t operator+(float a, complex64_t b) { - return {a + b.real, b.imag}; +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(U a, complex_t b) { + return {static_cast(a) + b.real, b.imag}; } -constexpr complex64_t operator+(complex64_t a, float b) { - return {a.real + b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(complex_t a, U b) { + return {a.real + static_cast(b), a.imag}; } -constexpr complex64_t operator-(complex64_t a, complex64_t b) { +template +constexpr complex_t operator-(complex_t a, complex_t b) { return {a.real - b.real, a.imag - b.imag}; } -constexpr complex64_t operator-(float a, complex64_t b) { - return {a - b.real, -b.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(U a, complex_t b) { + return {static_cast(a) - b.real, -b.imag}; } -constexpr complex64_t operator-(complex64_t a, float b) { - return {a.real - b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(complex_t a, U b) { + return {a.real - static_cast(b), a.imag}; } -constexpr complex64_t operator*(complex64_t a, complex64_t b) { +template +constexpr complex_t operator*(complex_t a, complex_t b) { return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; } -constexpr complex64_t operator/(complex64_t a, complex64_t b) { +template +constexpr complex_t operator/(complex_t a, complex_t b) { auto denom = b.real * b.real + b.imag * b.imag; auto x = a.real * b.real + a.imag * b.imag; auto y = a.imag * b.real - a.real * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator/(float a, complex64_t b) { +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator/(U a, complex_t b) { + auto scalar = static_cast(a); auto denom = b.real * b.real + b.imag * b.imag; - auto x = a * b.real; - auto y = -a * b.imag; + auto x = scalar * b.real; + auto y = -scalar * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator%(complex64_t a, complex64_t b) { +template +constexpr complex_t operator%(complex_t a, complex_t b) { auto real = a.real - (b.real * static_cast(a.real / b.real)); auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); if (real != 0 && (real < 0 != b.real < 0)) { @@ -172,3 +271,10 @@ constexpr complex64_t operator%(complex64_t a, complex64_t b) { } return {real, imag}; } + +static_assert( + (complex_t{1.0h, 2.0h} * complex_t{3.0h, 4.0h}).real == -5.0h); +static_assert( + (complex_t{bfloat16_t(1.0f), bfloat16_t(2.0f)} * + complex_t{bfloat16_t(3.0f), bfloat16_t(4.0f)}) + .real == bfloat16_t(-5.0f)); diff --git a/mlx/backend/metal/kernels/utils.h b/mlx/backend/metal/kernels/utils.h index f2a8362ae2..266f27e91c 100644 --- a/mlx/backend/metal/kernels/utils.h +++ b/mlx/backend/metal/kernels/utils.h @@ -75,14 +75,14 @@ struct Limits { static constexpr constant bool min = false; }; -template <> -struct Limits { - static constexpr constant complex64_t max = complex64_t( - metal::numeric_limits::infinity(), - metal::numeric_limits::infinity()); - static constexpr constant complex64_t min = complex64_t( - -metal::numeric_limits::infinity(), - -metal::numeric_limits::infinity()); +template +struct Limits> { + inline static constexpr constant complex_t max = complex_t( + metal::numeric_limits::infinity(), + metal::numeric_limits::infinity()); + inline static constexpr constant complex_t min = complex_t( + -metal::numeric_limits::infinity(), + -metal::numeric_limits::infinity()); }; /////////////////////////////////////////////////////////////////////////////// From 8584610b3fd4ca0586d5259b93827af850994cb0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Wed, 5 Aug 2026 09:37:08 +0300 Subject: [PATCH 057/222] Pad 2D conv input channels to reach the specialized Metal kernel (#3904) --- mlx/backend/metal/conv.cpp | 48 ++++++++++++++++++++++++++++++++++++-- python/tests/test_conv.py | 12 ++++++++++ 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 49f87d9a5b..e3c8462c5b 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -977,6 +977,36 @@ void depthwise_conv_2D_gpu( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +void pad_in_channels_conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const MLXConvParams<2>& conv_params) { + // Only the input channels are padded, so the kernel writes straight into out + // and the output stays contiguous. Assumes conv_params.groups == 1. + int extra_c = ((conv_params.C + 15) / 16) * 16 - conv_params.C; + + // Pad function + auto pad_array = [&](const array& x) { + auto xshape = x.shape(); + xshape.back() += extra_c; + array x_copy(xshape, x.dtype(), nullptr, {}); + array zero(0, x.dtype()); + pad_gpu(x, zero, x_copy, {-1}, {0}, s); + metal::get_command_encoder(s).add_temporary(x_copy); + + return x_copy; + }; + + array in = pad_array(in_pre); + array wt = pad_array(wt_pre); + auto new_params = + MLXConvParams<2>::with_padded_channels(conv_params, 0, extra_c); + implicit_gemm_conv_2D_gpu(s, d, in, wt, out, new_params); +} + void dispatch_conv_2D_gpu( const Stream& s, metal::Device& d, @@ -1026,9 +1056,23 @@ void dispatch_conv_2D_gpu( return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies); } + // Whether the specialized implicit gemm kernel can take the channels as-is. + bool specialized_channels = (conv_params.C <= 4 || conv_params.C % 16 == 0) && + (conv_params.O <= 16 || conv_params.O % 16 == 0); + + // Pad the input channels up to a multiple of 16 to use the faster specialized + // kernel instead of the general one. Only worth the padded work for stride-1 + // convs with a large enough output and kernel, and only when the output + // channels are already aligned so the kernel writes contiguously into out + // (padding those too needs a copy that can cost more than the kernel saves). + bool out_channels_aligned = conv_params.O <= 16 || conv_params.O % 16 == 0; + if (is_idil_one && is_stride_one && out_large && !specialized_channels && + out_channels_aligned && (conv_params.wS[0] * conv_params.wS[1]) >= 9) { + return pad_in_channels_conv_2D_gpu(s, d, in, wt, out, conv_params); + } + // Direct to implicit gemm conv - if (is_idil_one && (conv_params.C <= 4 || conv_params.C % 16 == 0) && - (conv_params.O <= 16 || conv_params.O % 16 == 0)) { + if (is_idil_one && specialized_channels) { return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params); } diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index 0243b87b3f..c5f9a2c1b2 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -1202,6 +1202,18 @@ def test_conv2d_unaligned_channels(self): y_hat = mx.conv2d(x, w) self.assertTrue(mx.allclose(y, y_hat)) + x = mx.random.uniform(shape=(2, 16, 16, 24)) + w = mx.random.uniform(shape=(32, 5, 5, 24)) + y = mx.conv2d(x, w, padding=2, stream=mx.cpu) + y_hat = mx.conv2d(x, w, padding=2) + self.assertTrue(mx.allclose(y, y_hat)) + + x = mx.random.uniform(shape=(2, 16, 16, 24)) + w = mx.random.uniform(shape=(32, 3, 3, 24)) + y = mx.conv_transpose2d(x, w, stream=mx.cpu) + y_hat = mx.conv_transpose2d(x, w) + self.assertTrue(mx.allclose(y, y_hat)) + def test_conv2d_large_filter_small_channels(self): x = mx.random.normal(shape=(1, 181, 181, 1)) w = mx.random.normal(shape=(1, 182, 182, 1)) From ab946f90e4f1a1659f6f5a2ea796bea3eebea8e9 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Wed, 5 Aug 2026 00:39:47 -0700 Subject: [PATCH 058/222] Support head dimension 96 in Metal full attention (#3943) --- benchmarks/python/sdpa_bench.py | 11 +++++++- .../steel/attn/kernels/steel_attention.metal | 1 + .../attn/kernels/steel_attention_nax.metal | 1 + .../metal/scaled_dot_product_attention.cpp | 3 ++- python/tests/test_fast_sdpa.py | 27 +++++++++++++++++++ 5 files changed, 41 insertions(+), 2 deletions(-) diff --git a/benchmarks/python/sdpa_bench.py b/benchmarks/python/sdpa_bench.py index 4e1c0234e6..bd279f0ead 100644 --- a/benchmarks/python/sdpa_bench.py +++ b/benchmarks/python/sdpa_bench.py @@ -189,6 +189,15 @@ def get_gflop_count(B, M, N, K): ( 1, 2048, 32121, 80, 32, 8), ) + shapes_96 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 96, 32, 8), + ( 1, 2048, 2048, 96, 32, 8), + ( 1, 4096, 4096, 96, 32, 8), + ( 1, 4096, 5000, 96, 32, 8), + ( 1, 2048, 32121, 96, 32, 8), + ) + shapes_128 = ( # ( B, qsl, ksl, head_dim, n_qh, n_kvh) ( 1, 1024, 1024, 128, 32, 8), @@ -199,7 +208,7 @@ def get_gflop_count(B, M, N, K): ) # fmt: on - shapes = shapes_64 + shapes_80 + shapes_128 + shapes = shapes_64 + shapes_80 + shapes_96 + shapes_128 masks = [None, "bool", "causal"] diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index 0ff9d91b00..7bddfcb054 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal @@ -13,6 +13,7 @@ #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index d40df50a45..c2b60b9cf0 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal @@ -13,6 +13,7 @@ #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 64, 4, 1, mname, mtype) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index ee5e09c870..aa8213edd6 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -626,7 +626,8 @@ bool ScaledDotProductAttention::use_fallback( query_head_dim == 256)) || (query_head_dim == 192 && value_head_dim == 128); const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128); + (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 96 || + query_head_dim == 128); const bool sdpa_full_supported_mask = !has_mask || has_arr_mask || (query_sequence_length <= key_sequence_length && do_causal); diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index f7440a0f8a..2418997bc8 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -117,6 +117,33 @@ def mlx_primitives_sdpa(q, k, v, scale, mask=None): class TestFastSDPA(mlx_tests.MLXTestCase): + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_96(self): + B, D, qH, kH = (1, 96, 8, 2) + for qL, kL, dtype, mask_str in product( + (64, 65), + (128, 127), + (mx.float16, mx.bfloat16, mx.float32), + (None, "additive", "bool", "causal"), + ): + with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, False, dtype + ) + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + + if dtype == mx.float32: + atol = 1e-5 + elif dtype == mx.bfloat16: + atol = 5e-3 + else: + atol = 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 From 2c46b953db88965c4270cc7306eda6887a3247f2 Mon Sep 17 00:00:00 2001 From: sashko-zakharchuk Date: Wed, 5 Aug 2026 11:15:04 +0300 Subject: [PATCH 059/222] Fix mx.remainder floored-mod for float16/bfloat16 on CPU (#3976) Co-authored-by: Cheng --- mlx/backend/cpu/simd/accelerate_simd.h | 2 +- mlx/backend/cpu/simd/base_simd.h | 10 +++++++++- python/tests/test_ops.py | 16 +++++++++++++++- 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/mlx/backend/cpu/simd/accelerate_simd.h b/mlx/backend/cpu/simd/accelerate_simd.h index f62c67d38b..e11cc7f5f9 100644 --- a/mlx/backend/cpu/simd/accelerate_simd.h +++ b/mlx/backend/cpu/simd/accelerate_simd.h @@ -241,7 +241,7 @@ Simd remainder(Simd a, Simd b) { } else { r = a - b * (a / b); } - if constexpr (std::is_signed_v) { + if constexpr (is_signed_v) { auto mask = r != 0 && (r < 0 != b < 0); r = select(mask, r + b, r); } diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index 775f5dfd10..be10a89ee4 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -13,6 +13,8 @@ #include // For _BitScanReverse #endif +#include "mlx/types/half_types.h" + namespace mlx::core::simd { template struct Simd; @@ -61,6 +63,12 @@ template constexpr bool is_complex().real())>> = true; +// std::is_signed_v is false for the custom float16_t/bfloat16_t types, so it +// skips the floored-mod sign correction for them. +template +inline constexpr bool is_signed_v = std::is_signed_v || + std::is_same_v || std::is_same_v; + template Simd rint(Simd in) { if constexpr (is_complex) { @@ -210,7 +218,7 @@ Simd remainder(Simd a_, Simd b_) { } else { r = std::remainder(a, b); } - if constexpr (std::is_signed_v) { + if constexpr (is_signed_v) { if (r != 0 && (r < 0 != b < 0)) { r += b; } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a0dcc0689e..214c59f542 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -344,7 +344,7 @@ def test_divide(self): self.assertEqual(z.item(), 2) def test_remainder(self): - for dt in [mx.int32, mx.float32]: + for dt in [mx.int32, mx.float32, mx.float16, mx.bfloat16]: x = mx.array(2, dtype=dt) y = mx.array(4, dtype=dt) @@ -370,6 +370,20 @@ def test_remainder(self): self.assertEqual(z.dtype, dt) self.assertEqual(z.item(), -1) + # floored remainder takes the sign of the divisor; check both + # correction directions (positive and negative divisor). + av = [-1.0, -1.0, -1.0, -3.0, 1.0, 1.0, 1.0, 3.0] + bv = [3.0, 5.0, 7.0, 8.0, -3.0, -5.0, -7.0, -8.0] + got = np.array( + mx.remainder(mx.array(av, dtype=dt), mx.array(bv, dtype=dt)).astype( + mx.float32 + ) + ) + expected = np.remainder( + np.array(av, dtype=np.float32), np.array(bv, dtype=np.float32) + ) + self.assertTrue(np.allclose(got, expected, atol=1e-2)) + x = mx.arange(10).astype(dt) - 5 y = x % 5 z = x % -5 From 268ecd62e3ac830a37af6fef513c8a3c8b31ff70 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:46:54 -0700 Subject: [PATCH 060/222] Fix sorted gather_mm activation row stride (#3960) --- mlx/backend/metal/matmul.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index ec3cb10c74..b61a58adb3 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -2114,7 +2114,7 @@ void gather_mm_rhs( int K = a.shape(-1); int M = a.size() / K; int N = b.shape(-1); - int lda = a.strides()[a.ndim() - 2]; // should be K + int lda = K; // Define the dispatch blocks int bm = 16, bn = 64, bk = 16; @@ -2247,7 +2247,7 @@ void gather_mm_rhs_nax( int K = a.shape(-1); int M = a.size() / K; int N = b.shape(-1); - int lda = a.strides()[a.ndim() - 2]; // should be K + int lda = K; int E = b.shape(0); // Define the dispatch blocks From fe92a0565d91a250512e5d4bda948066c383713f Mon Sep 17 00:00:00 2001 From: Ranran Date: Wed, 5 Aug 2026 17:49:35 -0500 Subject: [PATCH 061/222] Fix state corruption when a primitive throws during eval (#3675) Signed-off-by: ran Co-authored-by: Cheng --- mlx/transforms.cpp | 160 ++++++++++++++++++++++---------------- python/tests/test_eval.py | 32 ++++++++ 2 files changed, 126 insertions(+), 66 deletions(-) diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp index 9a8207339e..694bca53d8 100644 --- a/mlx/transforms.cpp +++ b/mlx/transforms.cpp @@ -225,85 +225,113 @@ array eval_impl(std::vector outputs, bool async) { } std::set open_streams; - while (!tape.empty()) { - auto arr = std::move(tape.back()); - tape.pop_back(); - - auto stream = arr.primitive().stream(); - open_streams.insert(stream); - - if (async) { - // Lookup corresponding event - auto e = events.find(stream.index); - if (e == events.end()) { - e = events.emplace(stream.index, Event{stream}).first; - } - e->second.set_value(1); - arr.attach_event(e->second); - for (auto& s : arr.siblings()) { - s.attach_event(e->second); + try { + while (!tape.empty()) { + auto arr = std::move(tape.back()); + tape.pop_back(); + + auto stream = arr.primitive().stream(); + open_streams.insert(stream); + + if (async) { + // Lookup corresponding event + auto e = events.find(stream.index); + if (e == events.end()) { + e = events.emplace(stream.index, Event{stream}).first; + } + e->second.set_value(1); + arr.attach_event(e->second); + for (auto& s : arr.siblings()) { + s.attach_event(e->second); + } } - } - for (auto& in : arr.inputs()) { - if (auto it = needs_fence.find(in.id()); it != needs_fence.end()) { - // Use fence to wait within a single eval - // Get the input array's stream fence and wait on the - // output arrays stream - fences[it->second.first].wait(stream, in); - } else if (in.event().valid()) { - if (in.event().is_signaled()) { - in.detach_event(); - } else if (in.event().stream() != stream) { - // Use event to wait across async eval - in.event().wait(stream); + for (auto& in : arr.inputs()) { + if (auto it = needs_fence.find(in.id()); it != needs_fence.end()) { + // Use fence to wait within a single eval + // Get the input array's stream fence and wait on the + // output arrays stream + fences[it->second.first].wait(stream, in); + } else if (in.event().valid()) { + if (in.event().is_signaled()) { + in.detach_event(); + } else if (in.event().stream() != stream) { + // Use event to wait across async eval + in.event().wait(stream); + } } } - } - if (arr.primitive().device() == Device::gpu) { - gpu::eval(arr); - } else { - cpu::eval(arr); - } + if (arr.primitive().device() == Device::gpu) { + gpu::eval(arr); + } else { + cpu::eval(arr); + } - if (scheduler::n_active_tasks() > MAX_ACTIVE_TASKS || - (get_active_memory() > get_memory_limit() && - scheduler::n_active_tasks() > 0)) { - // Commit any open streams - for (auto& s : open_streams) { - if (s.device == Device::gpu) { - gpu::finalize(s); + if (scheduler::n_active_tasks() > MAX_ACTIVE_TASKS || + (get_active_memory() > get_memory_limit() && + scheduler::n_active_tasks() > 0)) { + // Commit any open streams + for (auto& s : open_streams) { + if (s.device == Device::gpu) { + gpu::finalize(s); + } } - } - scheduler::wait_for_one(); - while (get_active_memory() > get_memory_limit() && - scheduler::n_active_tasks() > 0) { scheduler::wait_for_one(); - } - } - - auto maybe_update_fence = [&fences, &needs_fence, stream](const array& a) { - if (auto nf = needs_fence.find(a.id()); nf != needs_fence.end()) { - auto it = fences.find(stream.index); - if (it == fences.end()) { - it = fences.emplace(stream.index, Fence{stream}).first; + while (get_active_memory() > get_memory_limit() && + scheduler::n_active_tasks() > 0) { + scheduler::wait_for_one(); } - it->second.update(stream, a, nf->second.second); } - }; - arr.set_status(array::Status::evaluated); - // TODO Maybe always want the fence coherent kernel in the same cbuf - // as the other kernels? - maybe_update_fence(arr); - for (auto& sib : arr.siblings()) { - sib.set_status(array::Status::evaluated); - maybe_update_fence(sib); + auto maybe_update_fence = + [&fences, &needs_fence, stream](const array& a) { + if (auto nf = needs_fence.find(a.id()); nf != needs_fence.end()) { + auto it = fences.find(stream.index); + if (it == fences.end()) { + it = fences.emplace(stream.index, Fence{stream}).first; + } + it->second.update(stream, a, nf->second.second); + } + }; + + arr.set_status(array::Status::evaluated); + // TODO Maybe always want the fence coherent kernel in the same cbuf + // as the other kernels? + maybe_update_fence(arr); + for (auto& sib : arr.siblings()) { + sib.set_status(array::Status::evaluated); + maybe_update_fence(sib); + } + if (!arr.is_tracer()) { + arr.detach(); + } } - if (!arr.is_tracer()) { - arr.detach(); + } catch (...) { + // A primitive threw from inside its eval (e.g. argument validation in + // eval_gpu, or a JIT compile failure). Arrays evaluated earlier in this + // tape are already marked evaluated, but their kernels sit in pending + // command buffers that only the epilogue below would commit, and events + // attached during this eval would never be signaled. Left that way, a + // later read of an affected array returns an unwritten buffer or blocks + // forever. Signal the events and flush the touched streams, then let the + // exception propagate. + for (auto& [idx, e] : events) { + try { + auto es = e.stream(); + e.signal(es); + open_streams.insert(es); + } catch (...) { + } + } + for (auto& s : open_streams) { + try { + synchronize(s); + } catch (...) { + // Preserve the original exception. + } } + throw; } // Signal the event in its stream diff --git a/python/tests/test_eval.py b/python/tests/test_eval.py index 5d6daaec21..da7f0ea8df 100644 --- a/python/tests/test_eval.py +++ b/python/tests/test_eval.py @@ -195,6 +195,38 @@ def test_multistream_deadlock(self): mx.eval(z) mx.set_memory_limit(old_limit) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_eval_exception_does_not_corrupt_state(self): + # An exception thrown from inside a primitive's eval (here a Metal + # compile error raised lazily at eval time) must not corrupt arrays + # evaluated earlier in the same batch: they are already marked + # evaluated, so their pending command buffers must still be + # committed before the exception propagates. + a = mx.full((1024,), 3.0) + b = a * 2.0 # encoded in the same eval batch as the failing kernel + + kernel = mx.fast.metal_kernel( + name="test_eval_exception_bad_kernel", + input_names=["inp"], + output_names=["out"], + source="this is not metal code {", + ) + with self.assertRaises(Exception): + (y,) = kernel( + inputs=[b], + output_shapes=[b.shape], + output_dtypes=[b.dtype], + grid=(1, 1, 1), + threadgroup=(1, 1, 1), + ) + mx.eval(y) + + self.assertTrue(mx.all(b == 6.0).item()) + + # Fresh computations after the failure stay correct. + x = mx.full((512,), 2.0) + self.assertEqual((x + 1.0).sum().item(), 512.0 * 3.0) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 49ee7265f7ea19bf7889c08607e34f3d73d07b28 Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Thu, 6 Aug 2026 04:54:27 +0530 Subject: [PATCH 062/222] python: added __complex__ support (#3984) Co-authored-by: Cheng --- python/src/array.cpp | 5 +++++ python/tests/test_array.py | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/python/src/array.cpp b/python/src/array.cpp index a5a47dcec1..653b4a4c5b 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -1032,6 +1032,11 @@ void init_array(nb::module_& m) { nb::rv_policy::none) .def("__int__", [](mx::array& a) { return nb::int_(to_scalar(a)); }) .def("__float__", [](mx::array& a) { return nb::float_(to_scalar(a)); }) + .def( + "__complex__", + [](mx::array& a) { + return nb::cast>(to_scalar(a)); + }) .def( "__format__", [](mx::array& a, nb::object format_spec) { diff --git a/python/tests/test_array.py b/python/tests/test_array.py index cfd96e9236..75bab7da1f 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -2673,16 +2673,23 @@ def test_to_scalar(self): a = mx.array(1) self.assertEqual(int(a), 1) self.assertEqual(float(a), 1) + self.assertEqual(complex(a), 1 + 0j) a = mx.array(1.5) self.assertEqual(float(a), 1.5) self.assertEqual(int(a), 1) + self.assertEqual(complex(a), 1.5 + 0j) + + a = mx.array(1 + 2j, dtype=mx.complex64) # type: ignore + self.assertEqual(complex(a), 1 + 2j) a = mx.zeros((2, 1)) with self.assertRaises(ValueError): float(a) with self.assertRaises(ValueError): int(a) + with self.assertRaises(ValueError): + complex(a) def test_format(self): a = mx.arange(3) From 25e449d61bb063fb6a874d8fee642790556341a7 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 5 Aug 2026 17:44:15 -0700 Subject: [PATCH 063/222] python: Validate freeze and unfreeze keys against the whole model when recursing (#3966) --- python/mlx/nn/layers/base.py | 16 ++++++++++++++++ python/tests/test_nn.py | 17 +++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 03b8be7a51..e980f8343b 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -461,6 +461,16 @@ def _validate_keys(self, keys, strict): raise KeyError(f"Module doesn't contain member {k}.") return keys + def _validate_keys_recursive(self, keys): + # A key such as "bias" is expected to be present somewhere in the model + # but not necessarily in every single submodule, so validate against the + # whole model rather than each module in isolation. + keys = keys if isinstance(keys, list) else [keys] + modules = self.modules() + for k in keys: + if not any(k in m for m in modules): + raise KeyError(f"Module doesn't contain member {k}.") + def freeze( self, *, @@ -511,6 +521,9 @@ def _freeze_impl(_, m): m._no_grad.update(local_keys) if recurse: + if strict and keys is not None: + self._validate_keys_recursive(keys) + strict = False self.apply_to_modules(_freeze_impl) else: _freeze_impl("", self) @@ -561,6 +574,9 @@ def _unfreeze_impl(_, m): m._no_grad.difference_update(local_keys) if recurse: + if strict and keys is not None: + self._validate_keys_recursive(keys) + strict = False self.apply_to_modules(_unfreeze_impl) else: _unfreeze_impl("", self) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 951a0f22d2..5c85133a9f 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -199,6 +199,23 @@ def test_module_state(self): m.state["hello"] = "world" self.assertEqual(m.state["hello"], "world") + def test_freeze_strict_keys(self): + # "bias" is present in the model but not in every submodule, so a + # recursive freeze should accept it rather than raising on the first + # submodule that happens not to have one. + m = nn.Sequential(nn.Linear(2, 2, bias=False), nn.Linear(2, 2)) + m.freeze(keys="bias", strict=True) + trainable = dict(tree_flatten(m.trainable_parameters())) + self.assertFalse(any(k.endswith("bias") for k in trainable)) + + m.unfreeze(keys="bias", strict=True) + trainable = dict(tree_flatten(m.trainable_parameters())) + self.assertTrue(any(k.endswith("bias") for k in trainable)) + + # A key that is nowhere in the model is still an error. + with self.assertRaises(KeyError): + m.freeze(keys="not_a_member", strict=True) + def test_chaining(self): m = nn.Sequential(nn.Linear(2, 2), nn.ReLU(), nn.Linear(2, 1)) pre_freeze_num_params = len(m.parameters()) From eac436b26b51503aa1d2b81b2b3e2aacbe0be9ce Mon Sep 17 00:00:00 2001 From: jonathan308 Date: Wed, 5 Aug 2026 18:44:44 -0700 Subject: [PATCH 064/222] Fix mlx.launch --python: flag is parsed but never forwarded to the launch script (#4002) Co-authored-by: jonathan308 Co-authored-by: Cheng --- python/mlx/_distributed_utils/launch.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 6838897da4..3b6af453e2 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -47,7 +47,9 @@ def terminate(self): class RemoteProcess(CommandProcess): def __init__(self, rank, host, python, cwd, files, env, command): is_local = host == "127.0.0.1" - cmd = RemoteProcess.make_launch_script(rank, cwd, files, env, command, is_local) + cmd = RemoteProcess.make_launch_script( + rank, python, cwd, files, env, command, is_local + ) if not is_local: cmd = f"ssh -tt -o LogLevel=QUIET {shlex.quote(host)} {shlex.quote(cmd)}" @@ -104,7 +106,7 @@ def terminate(self): self._killed = c.stdout.strip() == "1" @staticmethod - def make_launch_script(rank, cwd, files, env, command, is_local): + def make_launch_script(rank, python, cwd, files, env, command, is_local): script = "" # Disable echo @@ -146,6 +148,10 @@ def make_launch_script(rank, cwd, files, env, command, is_local): # Finally add the rank script += f"export MLX_RANK={rank}; " + # Run the command with the explicitly requested python interpreter + if python is not None: + command = [python, *command] + # Replace the process with the script script += f"cmd=({' '.join(map(shlex.quote, command))}); " script += 'exec "${cmd[@]}"' @@ -514,13 +520,15 @@ def main(): help="The port to use for the NCCL communication (only for nccl backend)", ) parser.add_argument( - "--python", default=sys.executable, help="Use this python on the remote hosts" + "--python", + default=None, + help="Launch the command with this python interpreter on the remote hosts", ) args, rest = parser.parse_known_args() if args.print_python: - print(args.python) + print(args.python or sys.executable) return if len(rest) == 0: From f0eade3341265b960794bb4320840d622c9d88de Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 5 Aug 2026 19:07:55 -0700 Subject: [PATCH 065/222] chore: Fix broken fully_shard reference docstring (#4007) --- python/mlx/nn/layers/distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index 69804756c9..16979b5097 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -669,7 +669,7 @@ class FullyShardedModule(Module): Every parameter is sharded along axis 0, so each parameter's size along that axis must be divisible by the size of ``group``. - Use :func:`fully_shard` to wrap a module. + Use :func:`~mlx.nn.layers.distributed.fully_shard` to wrap a module. Args: module (mlx.nn.Module): The module whose parameters will be sharded. From 679f3efaae4193180bf63c86a080e76fb02247a4 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 5 Aug 2026 19:08:11 -0700 Subject: [PATCH 066/222] chore: Use the current interpreter in the comparative benchmark runner (#4016) --- benchmarks/python/comparative/compare.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/benchmarks/python/comparative/compare.py b/benchmarks/python/comparative/compare.py index 68b4a5bd32..aa81369fe2 100644 --- a/benchmarks/python/comparative/compare.py +++ b/benchmarks/python/comparative/compare.py @@ -4,6 +4,7 @@ import argparse import re +import sys from pathlib import Path from subprocess import run @@ -22,15 +23,15 @@ def run_or_raise(*args, **kwargs): def compare(args): - t_mlx = run_or_raise(["python", BENCH_MLX] + args) - t_torch = run_or_raise(["python", BENCH_TORCH] + args) + t_mlx = run_or_raise([sys.executable, BENCH_MLX] + args) + t_torch = run_or_raise([sys.executable, BENCH_TORCH] + args) print((t_torch - t_mlx) / t_torch, " ".join(args), sep="\t") def compare_mlx_dtypes(args, dt1, dt2): - t_mlx_dt1 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt1]) - t_mlx_dt2 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt2]) + t_mlx_dt1 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt1]) + t_mlx_dt2 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt2]) print((t_mlx_dt2 - t_mlx_dt1) / t_mlx_dt2, " ".join(args), sep="\t") From 752738deb83af3e3546e2e881de3f9640adea131 Mon Sep 17 00:00:00 2001 From: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:08:26 +0800 Subject: [PATCH 067/222] docs: Document MLX environment variables (#4000) --- docs/src/index.rst | 1 + docs/src/usage/compile.rst | 2 +- docs/src/usage/distributed.rst | 27 ++-- docs/src/usage/environment_variables.rst | 184 +++++++++++++++++++++++ docs/src/usage/precision.rst | 4 +- 5 files changed, 202 insertions(+), 16 deletions(-) create mode 100644 docs/src/usage/environment_variables.rst diff --git a/docs/src/index.rst b/docs/src/index.rst index d9beb16753..e29d396a89 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -44,6 +44,7 @@ are the CPU and GPU. usage/compile usage/numpy usage/precision + usage/environment_variables usage/distributed usage/using_streams usage/export diff --git a/docs/src/usage/compile.rst b/docs/src/usage/compile.rst index 662f69410c..f26ac6673c 100644 --- a/docs/src/usage/compile.rst +++ b/docs/src/usage/compile.rst @@ -155,7 +155,7 @@ contents) inside compiled functions. For debugging, inspecting arrays can be helpful. One way to do that is to globally disable compilation using the :func:`disable_compile` function or -``MLX_DISABLE_COMPILE`` flag. For example the following is okay even though +:envvar:`MLX_DISABLE_COMPILE` flag. For example the following is okay even though ``fun`` is compiled: .. code-block:: python diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 4586d7e8ae..866c627d13 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -341,11 +341,12 @@ of a gigantic model using MLX LM. .. note:: - Defining the environment variable ``MLX_METAL_FAST_SYNCH=1`` enables a - different, faster way of synchronizing between the GPU and the CPU. It is - not specific to the JACCL backend and can be used in all cases where the CPU - and GPU need to collaborate for some computation and is pretty critical for - low-latency communication since the communication is done by the CPU. + Defining the environment variable :envvar:`MLX_METAL_FAST_SYNCH` to ``1`` + enables a different, faster way of synchronizing between the GPU and the + CPU. It is not specific to the JACCL backend and can be used in all cases + where the CPU and GPU need to collaborate for some computation and is pretty + critical for low-latency communication since the communication is done by + the CPU. Custom side channel ^^^^^^^^^^^^^^^^^^^ @@ -506,10 +507,10 @@ Below we list the environment variables required to use each backend. Ring ^^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_HOSTFILE** should contain the path to a json file that contains IPs and +:envvar:`MLX_HOSTFILE` should contain the path to a json file that contains IPs and ports for each rank to listen to, something like the following: .. code-block:: json @@ -521,19 +522,19 @@ ports for each rank to listen to, something like the following: ["123.123.4.1:5000", "123.123.4.2:5000"] ] -**MLX_RING_VERBOSE** is optional and if set to 1 it enables some more logging +:envvar:`MLX_RING_VERBOSE` is optional and if set to 1 it enables some more logging from the distributed backend. JACCL ^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen +:envvar:`MLX_JACCL_COORDINATOR` should contain the IP and port that rank 0 can listen to all the other ranks connect to in order to establish the RDMA connections. -**MLX_IBV_DEVICES** should contain the path to a json file that contains the +:envvar:`MLX_IBV_DEVICES` should contain the path to a json file that contains the ibverbs device names that connect each node to each other node, something like the following: @@ -550,10 +551,10 @@ the following: NCCL ^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_WORLD_SIZE** should contain the total number of processes that will be +:envvar:`MLX_WORLD_SIZE` should contain the total number of processes that will be launched. **NCCL_HOST_IP** and **NCCL_PORT** should contain the IP and port that all diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst new file mode 100644 index 0000000000..3584f4c22e --- /dev/null +++ b/docs/src/usage/environment_variables.rst @@ -0,0 +1,184 @@ +.. _environment_variables: + +Environment Variables +===================== + +MLX uses environment variables to configure compilation, numerical precision, +backend behavior, and distributed execution. Set them before starting the +process. Many variables are read when the corresponding subsystem is first +initialized and changing them later may have no effect. + +Boolean variables use ``0`` to disable and a nonzero integer to enable unless +otherwise noted. + +General +------- + +.. envvar:: MLX_DISABLE_COMPILE + + Disable compilation globally. This variable is enabled by its presence, so + setting it to ``0`` also disables compilation. Calling + :func:`mlx.core.enable_compile` overrides it. + +.. envvar:: MLX_ENABLE_TF32 + + Allow reduced-precision ``float32`` matrix-multiplication family operations + on supported hardware. The default is ``1``. Set it to ``0`` to keep these + operations in full ``float32`` precision. See :doc:`precision`. + +Distributed +----------- + +The variables required to initialize a distributed process depend on the +backend. :doc:`distributed` describes their formats and how ``mlx.launch`` +sets them. + +.. envvar:: MLX_RANK + + The zero-based rank of the current process. This is used by the Ring, + JACCL, and NCCL backends. ``JACCL_RANK`` is accepted as a higher-priority + alias by JACCL. + +.. envvar:: MLX_HOSTFILE + + The path to the JSON host file used by the Ring backend. + +.. envvar:: MLX_RING_VERBOSE + + Enable verbose logging for the Ring backend. This variable is enabled by + its presence. + +.. envvar:: MLX_IBV_DEVICES + + The path to the JSON device-connectivity file used by JACCL. + ``JACCL_IBV_DEVICES`` is accepted as a higher-priority alias. + +.. envvar:: MLX_JACCL_COORDINATOR + + The coordinator address in ``IP:port`` form used to establish JACCL + connections. ``JACCL_COORDINATOR`` is accepted as a higher-priority alias. + +.. envvar:: MLX_JACCL_RING + + Prefer a ring topology for JACCL. This variable is enabled by its presence. + ``JACCL_RING`` is accepted as a higher-priority alias. + +.. envvar:: MLX_WORLD_SIZE + + The total number of processes in an NCCL group. + +.. envvar:: MLX_NCCL_TIMEOUT + + The timeout in milliseconds for establishing NCCL bootstrap connections. + The default is ``300000``. + +.. envvar:: MLX_MPI_LIBNAME + + Override the MPI dynamic-library name. The default is ``libmpi.dylib`` on + macOS and ``libmpi.so`` on other platforms. + +The NCCL backend also requires ``NCCL_HOST_IP`` and ``NCCL_PORT``. Setting +``NCCL_DEBUG=INFO`` enables additional logging while MLX establishes the +bootstrap connection. NCCL itself recognizes additional `NCCL environment +variables `_. +``CUDA_VISIBLE_DEVICES`` selects the local CUDA device for each process and is +handled by the CUDA runtime. + +Metal +----- + +.. envvar:: MLX_METAL_FAST_SYNCH + + Enable the faster Metal CPU/GPU synchronization path. The default is ``0``. + This requires Metal 3.2 or later (macOS 15 or later, or iOS 18 or later). + +Advanced tuning +--------------- + +These variables tune MLX implementation details. They are primarily intended +for development, diagnostics, and performance experiments. The defaults are +selected automatically for the current hardware and are appropriate for most +users. Their behavior may change as the implementation evolves. + +.. envvar:: MLX_BFS_MAX_WIDTH + + Set the breadth-first-search width limit used when constructing an + evaluation tape. The default is ``20``. + +.. envvar:: MLX_MAX_OPS_PER_BUFFER + + Override the maximum number of operations encoded in one Metal command + buffer or CUDA graph. The default depends on the device. + +.. envvar:: MLX_MAX_MB_PER_BUFFER + + Override the approximate memory limit, in megabytes, for one Metal command + buffer or CUDA graph. The default depends on the device. + +.. envvar:: MLX_METAL_GPU_ARCH + + Override the Metal GPU architecture string reported to MLX. This affects + architecture-specific kernel and scheduling choices, but does not change + the capabilities of the physical GPU. Forcing an architecture that does not + match the GPU can select incompatible kernels and produce incorrect results. + +.. envvar:: MLX_SDPA_BLOCKS + + Override the number of reduction blocks used by the Metal scaled + dot-product attention kernel. Positive values are rounded up to a multiple + of ``32``. + +CUDA +---- + +The MLX-prefixed variables in this section are advanced CUDA backend controls. + +.. envvar:: MLX_USE_CUDA_GRAPHS + + Enable CUDA graph capture and replay. The default is ``1``. + +.. envvar:: MLX_SAVE_CUDA_GRAPHS_DOT_FILE + + Use the specified value as the filename prefix when writing captured CUDA + graphs to numbered DOT files. An unset or empty value disables the output. + +.. envvar:: MLX_PTX_CACHE_DIR + + Override the directory used to cache runtime-compiled PTX. By default MLX + uses an ``mlx//ptx`` directory under the system temporary + directory. + +.. envvar:: MLX_CUDA_USE_CUDNN_SDPA + + Allow the CUDA backend to use cuDNN scaled dot-product attention when the + inputs and device are supported. The default is ``1``. + +.. envvar:: MLX_CUDA_CONV_CACHE_SIZE + + Set the CUDA convolution cache capacity. The default is ``128``. + +.. envvar:: MLX_CUDA_FFT_CACHE_SIZE + + Set the CUDA FFT plan cache capacity. The default is ``128``. + +.. envvar:: MLX_CUDA_GRAPH_CACHE_SIZE + + Set the CUDA graph cache capacity. The default is ``400``. + +.. envvar:: MLX_CUDA_SDPA_CACHE_SIZE + + Set the CUDA forward scaled dot-product attention cache capacity. The + default is ``256``. + +.. envvar:: MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE + + Set the CUDA backward scaled dot-product attention cache capacity. The + default is ``64``. + +.. envvar:: MLX_ENABLE_CACHE_THRASHING_CHECK + + Detect repeated CUDA cache misses and raise an error suggesting a larger + cache capacity. The default is ``1``. + +MLX also uses ``CUDA_HOME`` or ``CUDA_PATH`` to locate CUDA headers for runtime +kernel compilation when they cannot be found in the Python environment. diff --git a/docs/src/usage/precision.rst b/docs/src/usage/precision.rst index b5ad9a4f9e..440a1f705d 100644 --- a/docs/src/usage/precision.rst +++ b/docs/src/usage/precision.rst @@ -10,8 +10,8 @@ matrix-multiplication units. Inputs and outputs stay ``float32``, but results can differ from a full-precision reference by several orders of magnitude more than ``float32`` rounding alone would explain. -To keep these operations in full ``float32``, set ``MLX_ENABLE_TF32=0`` -when launching the process: +To keep these operations in full ``float32``, set +:envvar:`MLX_ENABLE_TF32` to ``0`` when launching the process: .. code-block:: shell From 4652b00869b2a0bbe5e6d6ac6663feba6082bcfc Mon Sep 17 00:00:00 2001 From: Jasper Meijerink Date: Thu, 6 Aug 2026 04:11:53 +0200 Subject: [PATCH 068/222] Fix CUDA batched GEMV grid overflow (#3929) Co-authored-by: Cheng --- mlx/backend/cuda/gemms/gemv.cu | 32 ++++++++++++++++++++++++++------ mlx/backend/cuda/matmul.cpp | 4 ++-- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/mlx/backend/cuda/gemms/gemv.cu b/mlx/backend/cuda/gemms/gemv.cu index bb463bf756..e48317daa6 100644 --- a/mlx/backend/cuda/gemms/gemv.cu +++ b/mlx/backend/cuda/gemms/gemv.cu @@ -13,6 +13,14 @@ namespace cg = cooperative_groups; static constexpr int rows_per_block = 8; +// Split batches across grid.y/z because both dimensions are capped at 65,535. +inline dim3 get_gemv_grid_dims(uint32_t num_blocks_x, uint32_t batch_size) { + constexpr uint32_t max_grid_yz_dim = 65535; + uint32_t num_blocks_z = cuda::ceil_div(batch_size, max_grid_yz_dim); + uint32_t num_blocks_y = cuda::ceil_div(batch_size, num_blocks_z); + return {num_blocks_x, num_blocks_y, num_blocks_z}; +} + // Accumulator type selection per input element type T. template struct GemvAccType { @@ -88,12 +96,17 @@ __global__ void gemv_batched( T* out, int rows, int cols, + uint32_t batch_size, const __grid_constant__ Shape batch_shape, const __grid_constant__ Strides mat_batch_strides, const __grid_constant__ Strides vec_batch_strides, int batch_ndim) { - auto block = cg::this_thread_block(); - auto batch_idx = block.group_index().y; + auto grid = cg::this_grid(); + auto batch_idx = + grid.block_index().z * grid.dim_blocks().y + grid.block_index().y; + if (batch_idx >= batch_size) { + return; + } auto [vec_offset, mat_offset] = elem_to_loc( batch_idx, batch_shape.data(), @@ -113,6 +126,7 @@ __global__ void gemv_gather( uint32_t* vec_indices, int rows, int cols, + uint32_t batch_size, const __grid_constant__ Shape mat_batch_shape, const __grid_constant__ Strides mat_batch_strides, int mat_batch_ndim, @@ -123,8 +137,12 @@ __global__ void gemv_gather( const __grid_constant__ Strides mat_index_strides, const __grid_constant__ Strides vec_index_strides, int index_batch_ndim) { - auto block = cg::this_thread_block(); - auto indices_idx = block.group_index().y; + auto grid = cg::this_grid(); + auto indices_idx = + grid.block_index().z * grid.dim_blocks().y + grid.block_index().y; + if (indices_idx >= batch_size) { + return; + } uint32_t index_mat, index_vec; if (index_batch_ndim > 1) { auto [mat_idx_offset, vec_idx_offset] = elem_to_loc( @@ -245,13 +263,14 @@ void gemv( auto kernel = gemv_batched; encoder.add_kernel_node( kernel, - dim3{num_blocks_x, batch_count}, + get_gemv_grid_dims(num_blocks_x, batch_count), block_dims, mat, vec, gpu_ptr(out), rows, cols, + batch_count, const_param(batch_shape), mat_strides, vec_strides, @@ -298,7 +317,7 @@ void gather_mv( auto kernel = gemv_gather; encoder.add_kernel_node( kernel, - dim3{num_blocks_x, batch_size}, + get_gemv_grid_dims(num_blocks_x, batch_size), block_dims, mat, vec, @@ -307,6 +326,7 @@ void gather_mv( gpu_ptr(vec_indices), rows, cols, + batch_size, const_param(mat_.shape()), const_param(mat_.strides()), mat_.ndim() - 2, diff --git a/mlx/backend/cuda/matmul.cpp b/mlx/backend/cuda/matmul.cpp index 90cc46e0d3..7505ea2cba 100644 --- a/mlx/backend/cuda/matmul.cpp +++ b/mlx/backend/cuda/matmul.cpp @@ -374,8 +374,8 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { auto& a_pre = inputs[0]; auto& b_pre = inputs[1]; - // Return 0s if either input is empty. - if (a_pre.size() == 0 || b_pre.size() == 0) { + // Return 0s if the output or either input is empty. + if (out.size() == 0 || a_pre.size() == 0 || b_pre.size() == 0) { array zero(0, a_pre.dtype()); encoder.add_temporary(zero); fill_gpu(zero, out, s); From 26db505342c49ebcba129c064505d15d1b900866 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 5 Aug 2026 20:47:25 -0700 Subject: [PATCH 069/222] chore: Fix all_gather benchmark collapsing its input to a scalar (#4017) --- benchmarks/python/synchronize_bench.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/python/synchronize_bench.py b/benchmarks/python/synchronize_bench.py index db006d70bf..2007ef12c3 100644 --- a/benchmarks/python/synchronize_bench.py +++ b/benchmarks/python/synchronize_bench.py @@ -42,7 +42,7 @@ def all_gather_benchmark(): def fn(x): for _ in range(its_per_eval): - x = mx.distributed.all_gather(x)[0] + x = mx.distributed.all_gather(x)[: a.shape[0]] return x ms = timeit(fn, a) / its_per_eval From f59b34dca3b28336c1ed5d1d8a2e6413e80a8dbb Mon Sep 17 00:00:00 2001 From: Kolja Wawrowsky <3075215+apocryphx@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:48:32 -0700 Subject: [PATCH 070/222] Check threadgroup size in the 1-pass sdpa_vector dispatch (#4018) Co-authored-by: Claude Opus 5 --- mlx/backend/metal/scaled_dot_product_attention.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index aa8213edd6..cb4af8523e 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -380,6 +380,7 @@ void sdpa_vector( // Get the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname, hash_name, func_consts); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.set_compute_pipeline_state(kernel); // Set its arguments From e534e1abc5e7b1f9361e8aa2f8f0baa75a4471d9 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 00:20:16 -0700 Subject: [PATCH 071/222] chore: Align example projects with the Python 3.10 minimum (#4024) --- examples/cmake_project/CMakeLists.txt | 2 +- examples/extensions/setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/cmake_project/CMakeLists.txt b/examples/cmake_project/CMakeLists.txt index 1d1b9d3db8..a955584782 100644 --- a/examples/cmake_project/CMakeLists.txt +++ b/examples/cmake_project/CMakeLists.txt @@ -8,7 +8,7 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) # Comment the following two commands only the MLX C++ library is installed and # set(MLX_ROOT "/path/to/mlx") directly if needed. find_package( - Python 3.9 + Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) execute_process( diff --git a/examples/extensions/setup.py b/examples/extensions/setup.py index 8dd1186ab8..be990a02b6 100644 --- a/examples/extensions/setup.py +++ b/examples/extensions/setup.py @@ -14,5 +14,5 @@ packages=["mlx_sample_extensions"], package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, zip_safe=False, - python_requires=">=3.8", + python_requires=">=3.10", ) From 4462e877fd8f5efaaf7c21e884b7647d306f4d66 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 00:57:15 -0700 Subject: [PATCH 072/222] Fix C++ benchmarks failing to build on overloaded astype (#4025) --- benchmarks/cpp/irregular_strides.cpp | 16 +++++++++++----- benchmarks/cpp/single_ops.cpp | 18 ++++++++++++------ 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/benchmarks/cpp/irregular_strides.cpp b/benchmarks/cpp/irregular_strides.cpp index cc4e975c9c..6da1e82aff 100644 --- a/benchmarks/cpp/irregular_strides.cpp +++ b/benchmarks/cpp/irregular_strides.cpp @@ -9,6 +9,12 @@ namespace mx = mlx::core; +// mx::astype is overloaded, so it cannot be passed directly to the timing +// helpers. Wrap the three argument form instead. +auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); +}; + void time_irregular_binary_ops_1D() { auto device = mx::default_device(); int size = 1000000; @@ -164,7 +170,7 @@ void time_irregular_astype_1D() { int step = 2; auto a = mx::random::uniform({size}); a = slice(a, {0}, {size}, {step}); - TIMEM("1D strided", mx::astype, a, mx::int32, device); + TIMEM("1D strided", astype, a, mx::int32, device); } void time_irregular_astype_2D() { @@ -173,16 +179,16 @@ void time_irregular_astype_2D() { mx::Shape shape = {size, size}; auto a = mx::random::uniform(shape); - TIMEM("2D regular", mx::astype, a, mx::int32, device); + TIMEM("2D regular", astype, a, mx::int32, device); a = mx::transpose(a); - TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device); + TIMEM("2D mx::transpose", astype, a, mx::int32, device); a = mx::broadcast_to(mx::random::uniform({size}), shape); - TIMEM("2D broadcast dim 0", mx::astype, a, mx::int32, device); + TIMEM("2D broadcast dim 0", astype, a, mx::int32, device); a = mx::broadcast_to(mx::random::uniform({size, 1}), shape); - TIMEM("2D broadcast dim 1", mx::astype, a, mx::int32, device); + TIMEM("2D broadcast dim 1", astype, a, mx::int32, device); } int main(int argc, char** argv) { diff --git a/benchmarks/cpp/single_ops.cpp b/benchmarks/cpp/single_ops.cpp index 1f93a78d71..05578269db 100644 --- a/benchmarks/cpp/single_ops.cpp +++ b/benchmarks/cpp/single_ops.cpp @@ -5,6 +5,12 @@ namespace mx = mlx::core; +// mx::astype is overloaded, so it cannot be passed directly to the timing +// helpers. Wrap the three argument form instead. +auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); +}; + void time_creation_ops() { int M = 2000; int N = 500; @@ -28,18 +34,18 @@ void time_type_conversions() { auto a = mx::zeros(shape, mx::float32); mx::eval(a); - TIMEM("mx::float32 to mx::int32", mx::astype, a, mx::int32, device); - TIMEM("mx::float32 to mx::uint32", mx::astype, a, mx::uint32, device); + TIMEM("mx::float32 to mx::int32", astype, a, mx::int32, device); + TIMEM("mx::float32 to mx::uint32", astype, a, mx::uint32, device); a = mx::zeros(shape, mx::int32); mx::eval(a); - TIMEM("mx::int32 to mx::float32", mx::astype, a, mx::float32, device); + TIMEM("mx::int32 to mx::float32", astype, a, mx::float32, device); a = mx::zeros(shape, mx::bool_); mx::eval(a); - TIMEM("bool to mx::float32", mx::astype, a, mx::float32, device); - TIMEM("bool to mx::int32", mx::astype, a, mx::int32, device); - TIMEM("bool to mx::uint32", mx::astype, a, mx::uint32, device); + TIMEM("bool to mx::float32", astype, a, mx::float32, device); + TIMEM("bool to mx::int32", astype, a, mx::int32, device); + TIMEM("bool to mx::uint32", astype, a, mx::uint32, device); } void time_random_generation() { From 074a6bfefa48f6aac42b2af0343bf0fbdc85dcf5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Thu, 6 Aug 2026 13:06:54 +0300 Subject: [PATCH 073/222] Fix conv_transpose maxBufferLength failures on Metal via tiled unfold (#3845) Co-authored-by: Cheng --- mlx/backend/metal/conv.cpp | 247 +++++++++++++++++---------- mlx/backend/metal/kernels/conv.metal | 14 +- python/tests/test_conv_transpose.py | 38 +++++ 3 files changed, 205 insertions(+), 94 deletions(-) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index e3c8462c5b..926f31f05a 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -30,6 +30,23 @@ ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { return result; } +inline int max_unfold_rows(metal::Device& d, size_t row_bytes, int total_rows) { + size_t max_buffer = d.mtl_device()->maxBufferLength(); + size_t max_rows = row_bytes == 0 ? total_rows : max_buffer / row_bytes; + // Force a smaller tile (tests, or to bound the unfold buffer further). + if (int forced = env::get_var("MLX_CONV_UNFOLD_TILE_ROWS", 0); forced > 0) { + max_rows = std::min(max_rows, static_cast(forced)); + } + if (max_rows == 0) { + std::ostringstream msg; + msg << "[conv] A single unfolding row needs " << row_bytes + << " bytes, which exceeds the maximum Metal buffer size of " + << max_buffer << " bytes."; + throw std::runtime_error(msg.str()); + } + return static_cast(std::min(max_rows, static_cast(total_rows))); +} + template void explicit_gemm_conv_ND_gpu( const Stream& s, @@ -42,36 +59,12 @@ void explicit_gemm_conv_ND_gpu( int implicit_M = safe_cast(out.size() / conv_params.O, "conv"); int implicit_K = safe_cast(wt.size() / conv_params.O, "conv"); int implicit_N = conv_params.O; - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate(kname, "naive_unfold_nd_", type_to_name(in_unfolded), "_", N); + concatenate(kname, "naive_unfold_nd_", type_to_name(in), "_", N); auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Reshape weight Shape wt_reshape{implicit_K, implicit_N}; @@ -82,23 +75,72 @@ void explicit_gemm_conv_ND_gpu( wt_flags.col_contiguous = true; wt_reshaped.copy_shared_buffer(wt, wt_restride, wt_flags, wt.data_size()); - // Perform gemm - std::vector copies = {in_unfolded}; - return steel_matmul( - s, - d, - /*a = */ in_unfolded, - /*b = */ wt_reshaped, - /*c = */ out, - /*M = */ implicit_M, - /*N = */ implicit_N, - /*K = */ implicit_K, - /*batch_size_out = */ 1, - /*a_cols = */ implicit_K, - /*b_cols = */ implicit_K, - /*a_transposed = */ false, - /*b_transposed = */ true, - /*copies = */ copies); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{out.strides(-2), out.strides(-1)}; + array out_2d({implicit_M, implicit_N}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = + MTL::Size(conv_params.C, implicit_K / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, implicit_N}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile}; + steel_matmul( + s, + d, + /*a = */ in_tile, + /*b = */ wt_reshaped, + /*c = */ out_tile, + /*M = */ tile_rows, + /*N = */ implicit_N, + /*K = */ implicit_K, + /*batch_size_out = */ 1, + /*a_cols = */ implicit_K, + /*b_cols = */ implicit_K, + /*a_transposed = */ false, + /*b_transposed = */ true, + /*copies = */ copies); + } } template @@ -122,36 +164,12 @@ void explicit_gemm_conv_group_ND_gpu( kernel_size *= conv_params.wS[i]; } - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K * groups}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate( - kname, "naive_unfold_transpose_nd_", type_to_name(in_unfolded), "_", N); + concatenate(kname, "naive_unfold_transpose_nd_", type_to_name(in), "_", N); auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Transpose kernel weights so that we can slice them by contiguous chunks // of channel groups. @@ -163,29 +181,78 @@ void explicit_gemm_conv_group_ND_gpu( // Materialize array wt_transpose = contiguous_copy_gpu(wt_view, s); - // Perform gemm - std::vector copies = {in_unfolded, wt_transpose}; - return steel_matmul_regular( - /* const Stream& s = */ s, - /* Device& d = */ d, - /* const array& a = */ in_unfolded, - /* const array& b = */ wt_transpose, - /* array& c = */ out, - /* int M = */ implicit_M, - /* int N = */ implicit_N, - /* int K = */ implicit_K, - /* int batch_size_out = */ groups, - /* int lda = */ implicit_K * groups, - /* int ldb = */ implicit_K, - /* int ldd = */ implicit_N * groups, - /* bool transpose_a = */ false, - /* bool transpose_b = */ true, - /* std::vector& copies = */ copies, - /* Shape batch_shape = */ {1}, - /* Strides batch_strides = */ {0}, - /* int64_t A_batch_strides = */ int64_t(implicit_K), - /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, - /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{out.strides(-2), out.strides(-1)}; + array out_2d({implicit_M, conv_params.O}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * groups * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = MTL::Size( + conv_params.C, (implicit_K * groups) / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, conv_params.O}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile, wt_transpose}; + steel_matmul_regular( + /* const Stream& s = */ s, + /* Device& d = */ d, + /* const array& a = */ in_tile, + /* const array& b = */ wt_transpose, + /* array& c = */ out_tile, + /* int M = */ tile_rows, + /* int N = */ implicit_N, + /* int K = */ implicit_K, + /* int batch_size_out = */ groups, + /* int lda = */ implicit_K * groups, + /* int ldb = */ implicit_K, + /* int ldd = */ implicit_N * groups, + /* bool transpose_a = */ false, + /* bool transpose_b = */ true, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ {1}, + /* Strides batch_strides = */ {0}, + /* int64_t A_batch_strides = */ int64_t(implicit_K), + /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, + /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + } } void implicit_gemm_conv_2D_gpu( diff --git a/mlx/backend/metal/kernels/conv.metal b/mlx/backend/metal/kernels/conv.metal index 220afc307d..6f3ce0cd43 100644 --- a/mlx/backend/metal/kernels/conv.metal +++ b/mlx/backend/metal/kernels/conv.metal @@ -20,6 +20,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -39,8 +40,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -83,6 +85,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -103,8 +106,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -150,6 +154,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); \ template \ [[host_name("naive_unfold_transpose_nd_" #name "_" #n)]] [[kernel]] void \ @@ -157,6 +162,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); #define instantiate_naive_unfold_nd_dims(name, itype) \ diff --git a/python/tests/test_conv_transpose.py b/python/tests/test_conv_transpose.py index 7289955ed4..e6def081a7 100644 --- a/python/tests/test_conv_transpose.py +++ b/python/tests/test_conv_transpose.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest from itertools import permutations @@ -805,6 +806,43 @@ def run_conv_transpose_3d_output_padding( dtype=dtype, ) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv_transpose_unfold_tiling(self): + # The explicit-GEMM conv path unfolds into one buffer that can exceed + # maxBufferLength for large outputs; it unfolds and runs the gemm in row + # tiles instead (issue #3082). + key = "MLX_CONV_UNFOLD_TILE_ROWS" + prev = os.environ.get(key) + cases = ( + (mx.conv_transpose1d, (2, 9, 4), (5, 3, 4), {"stride": 2}), + (mx.conv_transpose2d, (2, 5, 5, 4), (5, 3, 3, 4), {"stride": 2}), + (mx.conv_transpose3d, (1, 4, 4, 4, 2), (3, 2, 2, 2, 2), {"stride": 2}), + (mx.conv_transpose1d, (2, 9, 4), (6, 3, 2), {"stride": 2, "groups": 2}), + ) + try: + # tile_rows=1 forces uniform tiles; 7 does not divide any output, so + # it exercises a partial final tile too. + for conv, in_shape, wt_shape, kwargs in cases: + for tile_rows in (1, 7): + with self.subTest( + conv=conv.__name__, tile_rows=tile_rows, **kwargs + ): + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + w = mx.array(np.random.normal(size=wt_shape).astype(np.float32)) + os.environ.pop(key, None) + untiled = conv(x, w, **kwargs) + mx.eval(untiled) + os.environ[key] = str(tile_rows) + tiled = conv(x, w, **kwargs) + mx.eval(tiled) + self.assertTrue(np.allclose(untiled, tiled, atol=1e-4)) + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + if __name__ == "__main__": mlx_tests.MLXTestRunner() From d2275d54b7cad52166e7f6ae9164b8347f24feae Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 03:07:35 -0700 Subject: [PATCH 074/222] chore: Give each host a unique rank in Hostfile.from_list (#4027) --- python/mlx/_distributed_utils/common.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/mlx/_distributed_utils/common.py b/python/mlx/_distributed_utils/common.py index d57747a320..c7e0fdd43c 100644 --- a/python/mlx/_distributed_utils/common.py +++ b/python/mlx/_distributed_utils/common.py @@ -90,7 +90,7 @@ def from_file(cls, hostfile): @classmethod def from_list(cls, hostlist, repeats=1): hosts = [] - for i, h in enumerate(hostlist.split(",")): + for h in hostlist.split(","): if h == "": raise ValueError("Hostname cannot be empty") try: @@ -98,8 +98,8 @@ def from_list(cls, hostlist, repeats=1): ips = [h] except ValueError: ips = [] - for i in range(repeats): - hosts.append(Host(i, h, ips, [])) + for _ in range(repeats): + hosts.append(Host(len(hosts), h, ips, [])) return cls(hosts) From 32df751be4bcd22d89f77b641f6f414b0aa6f870 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 03:07:56 -0700 Subject: [PATCH 075/222] Fix crash when reporting partial rings in mlx.distributed_config (#4026) --- python/mlx/_distributed_utils/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/_distributed_utils/config.py b/python/mlx/_distributed_utils/config.py index d9262a0191..0ab31470b0 100644 --- a/python/mlx/_distributed_utils/config.py +++ b/python/mlx/_distributed_utils/config.py @@ -361,7 +361,7 @@ def check_valid_ring(hosts, rings, strict=True): log_error("Try passing --dot to visualize the connectivity") if len(rings) > 0: log_error("Rings found:") - for r in rings: + for r, _ in rings: log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}") sys.exit(1) return has_ring From f1c5bcd9909db89d2392e76eac26cc66c0c70440 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Thu, 6 Aug 2026 06:10:56 -0400 Subject: [PATCH 076/222] docs: Do not pass MLX_METAL_FAST_SYNCH=1 by default (#4005) Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com> Co-authored-by: Cheng Co-authored-by: Cheng --- docs/src/usage/distributed.rst | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 866c627d13..431aef7102 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -335,18 +335,21 @@ of a gigantic model using MLX LM. .. code-block:: - mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json \ - --env MLX_METAL_FAST_SYNCH=1 -- \ # <--- important + mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json -- \ /path/to/remote/python -m mlx_lm chat --model mlx-community/DeepSeek-R1-0528-4bit .. note:: Defining the environment variable :envvar:`MLX_METAL_FAST_SYNCH` to ``1`` - enables a different, faster way of synchronizing between the GPU and the - CPU. It is not specific to the JACCL backend and can be used in all cases - where the CPU and GPU need to collaborate for some computation and is pretty - critical for low-latency communication since the communication is done by - the CPU. + by passing ``--env MLX_METAL_FAST_SYNCH=1`` enables a different, faster way + of synchronizing between the GPU and the CPU. It is not specific to the + JACCL backend and can be used in all cases where the CPU and GPU need to + collaborate for some computation, and it matters for low-latency + communication since the communication is done by the CPU. + + It is however not reliable that can lead to deadlock and leave the GPU wedged + (see `#3142 `_), so it is off + by default and best left unset. Custom side channel ^^^^^^^^^^^^^^^^^^^ From 7e9c3f63db7f03767f1192f78a9f976b8af51812 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Thu, 6 Aug 2026 13:11:21 +0300 Subject: [PATCH 077/222] Fix signed-integer overflow in convolution shape arithmetic (#3938) Co-authored-by: Cheng --- mlx/ops.cpp | 24 +++++++++------ mlx/primitives.cpp | 56 ++++++++++++++++++++++------------- tests/ops_tests.cpp | 72 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 29 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 323c6a9a6e..cdc6e2e7d6 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4344,16 +4344,21 @@ array conv_transpose_general( std::vector padding_lo(padding.size()); std::vector padding_hi(padding.size()); for (int i = 0; i < padding.size(); ++i) { - int wt_size = 1 + dilation[i] * (weight.shape(1 + i) - 1); - padding_lo[i] = wt_size - padding[i] - 1; + int64_t wt_size = + 1 + static_cast(dilation[i]) * (weight.shape(1 + i) - 1); + padding_lo[i] = safe_cast(wt_size - padding[i] - 1, "conv"); - int conv_output_shape = (input.shape(i + 1) - 1) * stride[i] - - 2 * padding[i] + dilation[i] * (weight.shape(i + 1) - 1) + 1; + int64_t conv_output_shape = + static_cast(input.shape(i + 1) - 1) * stride[i] - + 2 * static_cast(padding[i]) + + static_cast(dilation[i]) * (weight.shape(i + 1) - 1) + 1; - int in_size = 1 + (conv_output_shape - 1); - int out_size = 1 + stride[i] * (input.shape(1 + i) - 1); - padding_hi[i] = in_size - out_size + padding[i] + - output_padding[i]; // Adjust with output_padding + int64_t in_size = 1 + (conv_output_shape - 1); + int64_t out_size = + 1 + static_cast(stride[i]) * (input.shape(1 + i) - 1); + // Adjust with output_padding + padding_hi[i] = + safe_cast(in_size - out_size + padding[i] + output_padding[i], "conv"); } auto ndim = stride.size(); @@ -4504,7 +4509,8 @@ array conv_general( for (int i = 0; i < spatial_dims; i++) { if (padding_lo[i] < 0) { - starts[i + 1] -= padding_lo[i]; + starts[i + 1] = safe_cast( + starts[i + 1] - static_cast(padding_lo[i]), "conv"); padding_lo[i] = 0; } diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 67afad54b0..c339bc444c 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1257,9 +1257,11 @@ array conv_weight_backward_patches( // padded shape for (int i = 1; i < in.ndim() - 1; i++) { - in_padded_shape[i] += padding_lo[i - 1] + padding_hi[i - 1]; - padding_ends[i] += padding_lo[i - 1]; - padding_starts[i] += padding_lo[i - 1]; + int64_t lo = padding_lo[i - 1]; + int64_t hi = padding_hi[i - 1]; + in_padded_shape[i] = safe_cast(in_padded_shape[i] + lo + hi, "conv"); + padding_ends[i] = safe_cast(padding_ends[i] + lo, "conv"); + padding_starts[i] = safe_cast(padding_starts[i] + lo, "conv"); } // padded strides (contiguous) @@ -1315,12 +1317,18 @@ array conv_weight_backward_patches( namespace { // Conv helpers -inline int conv_out_axis_size(int in_dim, int wt_dim, int stride, int padding) { +// Computed in 64 bits so extreme but in-range int32 parameters do not overflow. +inline int64_t conv_out_axis_size( + int64_t in_dim, + int64_t wt_dim, + int64_t stride, + int64_t padding) { return ((in_dim + padding - wt_dim) / stride) + 1; } // Conv helpers -inline int dilate_size(int dim, int dil) { +// Computed in 64 bits so extreme but in-range int32 parameters do not overflow. +inline int64_t dilate_size(int64_t dim, int64_t dil) { return 1 + dil * (dim - 1); } @@ -1399,13 +1407,16 @@ Shape Convolution::conv_out_shape( throw std::invalid_argument(msg.str()); } - int kd = dilate_size(wt_shape[i], kernel_dilation[i - 1]); - int id = dilate_size(in_shape[i], input_dilation[i - 1]); + int64_t kd = dilate_size(wt_shape[i], kernel_dilation[i - 1]); + int64_t id = dilate_size(in_shape[i], input_dilation[i - 1]); - out_shape[i] = conv_out_axis_size( - id, kd, strides[i - 1], pads_lo[i - 1] + pads_hi[i - 1]); + int64_t out_size = conv_out_axis_size( + id, + kd, + strides[i - 1], + static_cast(pads_lo[i - 1]) + pads_hi[i - 1]); - if (out_shape[i] <= 0) { + if (out_size <= 0) { std::ostringstream msg; msg << "[conv] Spatial dimensions of input after padding" << " cannot be smaller than weight spatial dimensions." @@ -1414,6 +1425,8 @@ Shape Convolution::conv_out_shape( << ", and weight of shape " << wt_shape << "."; throw std::invalid_argument(msg.str()); } + + out_shape[i] = safe_cast(out_size, "conv"); } out_shape[i] = O; @@ -1457,12 +1470,12 @@ std::vector Convolution::vjp( std::vector padding_hi = padding_hi_; for (int i = 0; i < padding_lo.size(); ++i) { - int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1); - padding_lo[i] = wt_size - padding_lo_[i] - 1; + int64_t wt_size = dilate_size(wt.shape(1 + i), kernel_dilation_[i]); + padding_lo[i] = safe_cast(wt_size - padding_lo_[i] - 1, "conv"); - int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1); - int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1); - padding_hi[i] = in_size - out_size + padding_hi_[i]; + int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); + int64_t out_size = dilate_size(cotan.shape(1 + i), kernel_strides_[i]); + padding_hi[i] = safe_cast(in_size - out_size + padding_hi_[i], "conv"); } // Check for negative padding @@ -1494,7 +1507,8 @@ std::vector Convolution::vjp( for (int i = 0; i < grad.ndim() - 2; i++) { if (padding_lo[i] < 0) { - starts[i + 1] -= padding_lo[i]; + starts[i + 1] = safe_cast( + starts[i + 1] - static_cast(padding_lo[i]), "conv"); } if (padding_hi[i] < 0) { stops[i + 1] += padding_hi[i]; @@ -1522,10 +1536,12 @@ std::vector Convolution::vjp( auto padding_hi = padding_lo_; for (int i = 0; i < padding_hi.size(); ++i) { - int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1); - int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1); - int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1); - padding_hi[i] = out_size - in_size + wt_size - padding_hi[i] - 1; + int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); + int64_t out_size = + dilate_size(cotan.shape(1 + i), kernel_strides_[i]); + int64_t wt_size = dilate_size(wt.shape(1 + i), kernel_dilation_[i]); + padding_hi[i] = safe_cast( + out_size - in_size + wt_size - padding_hi[i] - 1, "conv"); } auto cotan_trans = swapaxes(cotan, 0, -1, stream()); diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 741530aaf7..1f82a9946c 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -4342,6 +4342,78 @@ TEST_CASE("test conv_transpose3d with output_padding") { CHECK(array_equal(out, expected).item()); } +TEST_CASE("test conv shape overflow") { + // Conv shape arithmetic must not overflow (signed-int UB) for large but + // otherwise valid int32 parameters; out-of-range results are rejected + // gracefully. https://github.com/ml-explore/mlx/issues/3611 + const int imax = 2147483647; + const int imin = -2147483647 - 1; + auto in = zeros({1, 8, 8, 1}); + auto wt = zeros({1, 3, 3, 1}); + + // A kernel dilated past the input reports the spatial-size error. + CHECK_THROWS_AS( + conv_general(in, wt, {1, 1}, {0, 0}, {0, 0}, {imax, imax}, {1, 1}), + std::invalid_argument); + + // Padding sums, input dilation, and negating a padding of INT_MIN raise. + CHECK_THROWS_AS( + conv_general(in, wt, {1, 1}, {imax, imax}, {imax, imax}, {1, 1}, {1, 1}), + std::overflow_error); + CHECK_THROWS_AS( + conv_general(in, wt, {1, 1}, {imax, 0}, {0, 0}, {1, 1}, {1, 1}), + std::overflow_error); + CHECK_THROWS_AS( + conv_general(in, wt, {1, 1}, {0, 0}, {0, 0}, {1, 1}, {imax, imax}), + std::overflow_error); + CHECK_THROWS_AS( + conv_general(in, wt, {1, 1}, {imin, imin}, {0, 0}, {1, 1}, {1, 1}), + std::overflow_error); + + // The transposed padding setup runs before conv_general validates it. + auto in_t = zeros({1, 4, 4, 1}); + CHECK_THROWS_AS( + conv_transpose2d(in_t, wt, {1, 1}, {0, 0}, {imax, imax}, {0, 0}), + std::overflow_error); + CHECK_THROWS_AS( + conv_transpose2d(in_t, wt, {1, 1}, {imin, imin}, {1, 1}, {0, 0}), + std::overflow_error); + + // The dilated input and kernel are both near 4e9 and cancel in the forward + // output, so only the gradient's own recompute goes out of range. + auto in_g = zeros({1, 3, 1, 1}); + auto wt_g = zeros({1, 200000, 1, 1}); + auto conv_g = [](const std::vector& primals) { + return std::vector{conv_general( + primals[0], + primals[1], + {1, 1}, + {0, 0}, + {0, 0}, + {20000, 1}, + {2000000000, 1})}; + }; + auto cotan = ones(conv_g({in_g, wt_g})[0].shape()); + CHECK_THROWS_AS(vjp(conv_g, {in_g, wt_g}, {cotan}), std::overflow_error); + + // The weight gradient pads without dividing by the stride. + auto in_w = zeros({1, 8, 8, 1}); + auto conv_w = [&in_w, imax](const std::vector& primals) { + return std::vector{conv_general( + in_w, primals[0], {imax, 1}, {imax, 0}, {imax, 0}, {1, 1}, {1, 1})}; + }; + auto cotan_w = ones(conv_w({wt})[0].shape()); + CHECK_THROWS_AS(vjp(conv_w, {wt}, {cotan_w}), std::overflow_error); + + // In-range parameters still give the same shapes. + CHECK_EQ( + conv_general(in, wt, {1, 1}, {1, 1}, {1, 1}, {2, 2}, {1, 1}).shape(), + Shape{1, 6, 6, 1}); + CHECK_EQ( + conv_transpose2d(in_t, wt, {2, 2}, {1, 1}, {1, 1}, {1, 1}).shape(), + Shape{1, 8, 8, 1}); +} + TEST_CASE("test fp8 conversion") { for (auto t : {float32, float16, bfloat16}) { array in({-1.125, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0}, t); From 39d9a8ac2b5449a49a910429e0edfedc8ff81372 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:58:39 -0700 Subject: [PATCH 078/222] Template Metal C2C FFT scalar lanes (#3969) --- mlx/backend/metal/kernels/fft.h | 111 ++++---- mlx/backend/metal/kernels/fft/radix.h | 302 ++++++++++++---------- mlx/backend/metal/kernels/fft/readwrite.h | 216 +++++++++------- 3 files changed, 354 insertions(+), 275 deletions(-) diff --git a/mlx/backend/metal/kernels/fft.h b/mlx/backend/metal/kernels/fft.h index 3cce29c574..53b48ca7aa 100644 --- a/mlx/backend/metal/kernels/fft.h +++ b/mlx/backend/metal/kernels/fft.h @@ -16,7 +16,7 @@ using namespace metal; #define MAX_RADIX 13 // Reached when elems_per_thread_ = 6, max_radix = 13 -// and some threads have to do 3 radix 6s requiring 18 float2s. +// and some threads have to do 3 radix 6s requiring 18 complex values. #define MAX_OUTPUT_SIZE 18 // Specialize for a particular value of N at runtime @@ -46,17 +46,18 @@ STEEL_CONST int rader_4_steps_ [[function_constant(19)]]; STEEL_CONST int rader_3_steps_ [[function_constant(20)]]; STEEL_CONST int rader_2_steps_ [[function_constant(21)]]; -// See "radix.h" for radix codelets -typedef void (*RadixFunc)(thread float2*, thread float2*); +// See "radix.h" for radix codelets. +template +using RadixFunc = void (*)(thread vec*, thread vec*); // Perform a single radix n butterfly with appropriate twiddles -template +template radix_func> METAL_FUNC void radix_butterfly( int i, int p, - thread float2* x, + thread vec* x, thread short* indices, - thread float2* y) { + thread vec* y) { // i: the index in the overall DFT that we're processing. // p: the size of the DFTs we're merging at this step. // m: how many threads are working on this DFT. @@ -75,14 +76,14 @@ METAL_FUNC void radix_butterfly( // Apply twiddles if (p > 1) { - float2 twiddle_1 = get_twiddle(k, radix * p); - float2 twiddle = twiddle_1; - x[1] = complex_mul(x[1], twiddle); + vec twiddle_1 = get_twiddle(k, radix * p); + vec twiddle = twiddle_1; + x[1] = complex_mul(x[1], twiddle); STEEL_PRAGMA_UNROLL for (int t = 2; t < radix; t++) { - twiddle = complex_mul(twiddle, twiddle_1); - x[t] = complex_mul(x[t], twiddle); + twiddle = complex_mul(twiddle, twiddle_1); + x[t] = complex_mul(x[t], twiddle); } } @@ -96,17 +97,17 @@ METAL_FUNC void radix_butterfly( // Perform all the radix steps required for a // particular radix size n. -template +template radix_func> METAL_FUNC void radix_n_steps( int i, thread int* p, int m, int n, int num_steps, - thread float2* inputs, + thread vec* inputs, thread short* indices, - thread float2* values, - threadgroup float2* buf) { + thread vec* values, + threadgroup vec* buf) { int m_r = n / radix; // When combining different sized radices, we have to do // multiple butterflies in a single thread. @@ -126,7 +127,7 @@ METAL_FUNC void radix_n_steps( for (int r = 0; r < radix; r++) { inputs[r] = buf[index + r * m_r]; } - radix_butterfly( + radix_butterfly( index, *p, inputs, indices + t * radix, values + t * radix); } } @@ -151,15 +152,19 @@ METAL_FUNC void radix_n_steps( } #define RADIX_STEP(radix, radix_func, num_steps) \ - radix_n_steps( \ + radix_n_steps>( \ fft_idx, p, m, n, num_steps, inputs, indices, values, buf); -template -METAL_FUNC void -perform_fft(int fft_idx, thread int* p, int m, int n, threadgroup float2* buf) { - float2 inputs[MAX_RADIX]; +template +METAL_FUNC void perform_fft( + int fft_idx, + thread int* p, + int m, + int n, + threadgroup vec* buf) { + vec inputs[MAX_RADIX]; short indices[MAX_OUTPUT_SIZE]; - float2 values[MAX_OUTPUT_SIZE]; + vec values[MAX_OUTPUT_SIZE]; RADIX_STEP(2, radix2, rader ? rader_2_steps_ : radix_2_steps_); RADIX_STEP(3, radix3, rader ? rader_3_steps_ : radix_3_steps_); @@ -184,7 +189,8 @@ template constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { - threadgroup float2 shared_in[tg_mem_size]; + using scalar_T = typename FFTIOTypeTraits::scalar_T; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -208,9 +214,9 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -219,7 +225,8 @@ template [[kernel]] void rader_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* raders_b_q [[buffer(2)]], + const device vec::scalar_T, 2>* + raders_b_q [[buffer(2)]], const device short* raders_g_q [[buffer(3)]], const device short* raders_g_minus_q [[buffer(4)]], constant const int& n, @@ -227,6 +234,7 @@ template constant const int& rader_n, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Use Rader's algorithm to compute fast FFTs // when a prime factor `p` of `n` is greater than 13 but // has `p - 1` Stockham decomposable into prime factors <= 13. @@ -250,7 +258,7 @@ template // // Rader's uses fewer operations than Bluestein's and so // is more accurate. It's also faster in most cases. - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -275,7 +283,7 @@ template int fft_idx = elem.z; int tg_idx = elem.y * n; - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; // rader_m = n / rader_n; int rader_m = rader_m_; @@ -287,10 +295,10 @@ template // 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 short x_0_index = metal::min(fft_idx * elems_per_thread_ / (rader_n - 1), rader_m - 1); - float2 x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; + vec x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; // Do the Rader permutation in shared memory - float2 temp[MAX_RADIX]; + vec temp[MAX_RADIX]; int max_index = n - rader_m - 1; for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); @@ -309,19 +317,20 @@ template // Rader FFT on x[rader_m:] int p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); // x_1 + ... + x_n is computed for us in the first FFT step so // we save it in the first rader_m indices of the array for later. int x_sum_index = metal::min(fft_idx, rader_m - 1); buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)]; - float2 inv = {1.0f, -1.0f}; + vec inv = {1.0f, -1.0f}; for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); short interleaved_index = index / rader_m + (index % rader_m) * (rader_n - 1); - temp[e] = complex_mul( + temp[e] = complex_mul( buf[rader_m + interleaved_index], raders_b_q[interleaved_index % (rader_n - 1)]); } @@ -337,9 +346,11 @@ template // Rader IFFT on x[rader_m:] p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); - float2 rader_inv_factor = {1.0f / (rader_n - 1), -1.0f / (rader_n - 1)}; + scalar_T rader_inv_r = static_cast(1.0f / (rader_n - 1)); + vec rader_inv_factor = {rader_inv_r, -rader_inv_r}; for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, n - rader_m - 1); @@ -348,7 +359,7 @@ template } // Use the sum of elements that was computed in the first FFT - float2 x_sum = buf[x_0_index] + x_0[0]; + vec x_sum = buf[x_0_index] + x_0[0]; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -365,7 +376,7 @@ template threadgroup_barrier(mem_flags::mem_threadgroup); p = rader_n; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -374,13 +385,16 @@ template [[kernel]] void bluestein_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* w_q [[buffer(2)]], - const device float2* w_k [[buffer(3)]], + const device vec::scalar_T, 2>* w_q + [[buffer(2)]], + const device vec::scalar_T, 2>* w_k + [[buffer(3)]], constant const int& length, constant const int& n, constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Computes arbitrary length FFTs with Bluestein's algorithm // // In numpy: @@ -390,7 +404,7 @@ template // Where w_k and w_q are precomputed on CPU in high precision as: // w_k = np.exp(-1j * np.pi / n * (np.arange(-n + 1, n) ** 2)) // w_q = np.fft.fft(1/w_k[-n:]) - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -414,22 +428,22 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; // fft - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); - float2 inv = float2(1.0f, -1.0f); + vec inv = {1.0f, -1.0f}; for (int t = 0; t < elems_per_thread_; t++) { int index = fft_idx + t * m; - buf[index] = complex_mul(buf[index], w_q[index]) * inv; + buf[index] = complex_mul(buf[index], w_q[index]) * inv; } threadgroup_barrier(mem_flags::mem_threadgroup); // ifft p = 1; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write_padded(length, w_k); } @@ -448,6 +462,7 @@ template < constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Fast four step FFT implementation for powers of 2. int overall_n = n1 * n2; int n = step == 0 ? n1 : n2; @@ -457,8 +472,8 @@ template < int m = grid.z; int fft_idx = elem.z; - threadgroup float2 shared_in[tg_mem_size]; - threadgroup float2* buf = &shared_in[elem.y * n]; + threadgroup vec shared_in[tg_mem_size]; + threadgroup vec* buf = &shared_in[elem.y * n]; using read_writer_t = ReadWriter; read_writer_t read_writer = read_writer_t( @@ -480,7 +495,7 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); int p = 1; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write_strided(stride, overall_n); } diff --git a/mlx/backend/metal/kernels/fft/radix.h b/mlx/backend/metal/kernels/fft/radix.h index bd61eef6d7..c209a4bc37 100644 --- a/mlx/backend/metal/kernels/fft/radix.h +++ b/mlx/backend/metal/kernels/fft/radix.h @@ -16,49 +16,63 @@ them into (n-1)=6,10,12 codelets. */ #include #include -METAL_FUNC float2 complex_mul(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); +// The codelets are templated over the scalar lane type so reduced-precision +// complex paths can reuse them; complex values are plain two-lane vectors. + +template +METAL_FUNC metal::vec complex_mul( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x}; } // Complex mul followed by conjugate -METAL_FUNC float2 complex_mul_conj(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x); +template +METAL_FUNC metal::vec complex_mul_conj( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x}; } // Compute an FFT twiddle factor -METAL_FUNC float2 get_twiddle(int k, int p) { - float theta = -2.0f * k * M_PI_F / p; - - float2 twiddle = {metal::fast::cos(theta), metal::fast::sin(theta)}; - return twiddle; +template +METAL_FUNC metal::vec get_twiddle(int k, int p) { + // Derive phase evaluation precision from the scalar lane and Metal's pi + // constant. Reduced lanes currently promote to float for fast trig. + using phase_T = decltype(M_PI_F * T(0)); + phase_T theta = -phase_T(2) * phase_T(k) * phase_T(M_PI_F) / phase_T(p); + return {metal::fast::cos(theta), metal::fast::sin(theta)}; } -METAL_FUNC void radix2(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix2(thread metal::vec* x, thread metal::vec* y) { y[0] = x[0] + x[1]; y[1] = x[0] - x[1]; } -METAL_FUNC void radix3(thread float2* x, thread float2* y) { - float pi_2_3 = -0.8660254037844387; +template +METAL_FUNC void radix3(thread metal::vec* x, thread metal::vec* y) { + T pi_2_3 = -0.8660254037844387; - float2 a_1 = x[1] + x[2]; - float2 a_2 = x[1] - x[2]; + metal::vec a_1 = x[1] + x[2]; + metal::vec a_2 = x[1] - x[2]; y[0] = x[0] + a_1; - float2 b_1 = x[0] - 0.5 * a_1; - float2 b_2 = pi_2_3 * a_2; + metal::vec b_1 = x[0] - 0.5 * a_1; + metal::vec b_2 = pi_2_3 * a_2; - float2 b_2_j = {-b_2.y, b_2.x}; + metal::vec b_2_j = {-b_2.y, b_2.x}; y[1] = b_1 + b_2_j; y[2] = b_1 - b_2_j; } -METAL_FUNC void radix4(thread float2* x, thread float2* y) { - float2 z_0 = x[0] + x[2]; - float2 z_1 = x[0] - x[2]; - float2 z_2 = x[1] + x[3]; - float2 z_3 = x[1] - x[3]; - float2 z_3_i = {z_3.y, -z_3.x}; +template +METAL_FUNC void radix4(thread metal::vec* x, thread metal::vec* y) { + metal::vec z_0 = x[0] + x[2]; + metal::vec z_1 = x[0] - x[2]; + metal::vec z_2 = x[1] + x[3]; + metal::vec z_3 = x[1] - x[3]; + metal::vec z_3_i = {z_3.y, -z_3.x}; y[0] = z_0 + z_2; y[1] = z_1 + z_3_i; @@ -66,25 +80,26 @@ METAL_FUNC void radix4(thread float2* x, thread float2* y) { y[3] = z_1 - z_3_i; } -METAL_FUNC void radix5(thread float2* x, thread float2* y) { - float2 root_5_4 = 0.5590169943749475; - float2 sin_2pi_5 = 0.9510565162951535; - float2 sin_1pi_5 = 0.5877852522924731; - - float2 a_1 = x[1] + x[4]; - float2 a_2 = x[2] + x[3]; - float2 a_3 = x[1] - x[4]; - float2 a_4 = x[2] - x[3]; - - float2 a_5 = a_1 + a_2; - float2 a_6 = root_5_4 * (a_1 - a_2); - float2 a_7 = x[0] - a_5 / 4; - float2 a_8 = a_7 + a_6; - float2 a_9 = a_7 - a_6; - float2 a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; - float2 a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; - float2 a_10_j = {a_10.y, -a_10.x}; - float2 a_11_j = {a_11.y, -a_11.x}; +template +METAL_FUNC void radix5(thread metal::vec* x, thread metal::vec* y) { + T root_5_4 = 0.5590169943749475; + T sin_2pi_5 = 0.9510565162951535; + T sin_1pi_5 = 0.5877852522924731; + + metal::vec a_1 = x[1] + x[4]; + metal::vec a_2 = x[2] + x[3]; + metal::vec a_3 = x[1] - x[4]; + metal::vec a_4 = x[2] - x[3]; + + metal::vec a_5 = a_1 + a_2; + metal::vec a_6 = root_5_4 * (a_1 - a_2); + metal::vec a_7 = x[0] - a_5 / 4; + metal::vec a_8 = a_7 + a_6; + metal::vec a_9 = a_7 - a_6; + metal::vec a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; + metal::vec a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; + metal::vec a_10_j = {a_10.y, -a_10.x}; + metal::vec a_11_j = {a_11.y, -a_11.x}; y[0] = x[0] + a_5; y[1] = a_8 + a_10_j; @@ -93,23 +108,24 @@ METAL_FUNC void radix5(thread float2* x, thread float2* y) { y[4] = a_8 - a_10_j; } -METAL_FUNC void radix6(thread float2* x, thread float2* y) { - float sin_pi_3 = 0.8660254037844387; - float2 a_1 = x[2] + x[4]; - float2 a_2 = x[0] - a_1 / 2; - float2 a_3 = sin_pi_3 * (x[2] - x[4]); - float2 a_4 = x[5] + x[1]; - float2 a_5 = x[3] - a_4 / 2; - float2 a_6 = sin_pi_3 * (x[5] - x[1]); - float2 a_7 = x[0] + a_1; - - float2 a_3_i = {a_3.y, -a_3.x}; - float2 a_6_i = {a_6.y, -a_6.x}; - float2 a_8 = a_2 + a_3_i; - float2 a_9 = a_2 - a_3_i; - float2 a_10 = x[3] + a_4; - float2 a_11 = a_5 + a_6_i; - float2 a_12 = a_5 - a_6_i; +template +METAL_FUNC void radix6(thread metal::vec* x, thread metal::vec* y) { + T sin_pi_3 = 0.8660254037844387; + metal::vec a_1 = x[2] + x[4]; + metal::vec a_2 = x[0] - a_1 / 2; + metal::vec a_3 = sin_pi_3 * (x[2] - x[4]); + metal::vec a_4 = x[5] + x[1]; + metal::vec a_5 = x[3] - a_4 / 2; + metal::vec a_6 = sin_pi_3 * (x[5] - x[1]); + metal::vec a_7 = x[0] + a_1; + + metal::vec a_3_i = {a_3.y, -a_3.x}; + metal::vec a_6_i = {a_6.y, -a_6.x}; + metal::vec a_8 = a_2 + a_3_i; + metal::vec a_9 = a_2 - a_3_i; + metal::vec a_10 = x[3] + a_4; + metal::vec a_11 = a_5 + a_6_i; + metal::vec a_12 = a_5 - a_6_i; y[0] = a_7 + a_10; y[1] = a_8 - a_11; @@ -119,26 +135,28 @@ METAL_FUNC void radix6(thread float2* x, thread float2* y) { y[5] = a_9 - a_12; } -METAL_FUNC void radix7(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix7(thread metal::vec* x, thread metal::vec* y) { // Rader's algorithm - float2 inv = {1 / 6.0, -1 / 6.0}; + T inv_r = static_cast(1 / 6.0); + metal::vec inv = {inv_r, -inv_r}; // fft - float2 in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; - radix6(in1, y + 1); + metal::vec in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; + radix6(in1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(2.44013336, -1.02261879)); - y[3] = complex_mul_conj(y[3], float2(2.37046941, -1.17510629)); - y[4] = complex_mul_conj(y[4], float2(0, -2.64575131)); - y[5] = complex_mul_conj(y[5], float2(2.37046941, 1.17510629)); - y[6] = complex_mul_conj(y[6], float2(-2.44013336, -1.02261879)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(2.44013336, -1.02261879)); + y[3] = complex_mul_conj(y[3], metal::vec(2.37046941, -1.17510629)); + y[4] = complex_mul_conj(y[4], metal::vec(0, -2.64575131)); + y[5] = complex_mul_conj(y[5], metal::vec(2.37046941, 1.17510629)); + y[6] = complex_mul_conj(y[6], metal::vec(-2.44013336, -1.02261879)); // ifft - radix6(y + 1, x + 1); + radix6(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[5] = x[2] * inv + x[0]; @@ -148,79 +166,87 @@ METAL_FUNC void radix7(thread float2* x, thread float2* y) { y[3] = x[6] * inv + x[0]; } -METAL_FUNC void radix8(thread float2* x, thread float2* y) { - float cos_pi_4 = 0.7071067811865476; - float2 w_0 = {cos_pi_4, -cos_pi_4}; - float2 w_1 = {-cos_pi_4, -cos_pi_4}; - float2 temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; - radix4(temp, x); - radix4(temp + 4, x + 4); +template +METAL_FUNC void radix8(thread metal::vec* x, thread metal::vec* y) { + T cos_pi_4 = 0.7071067811865476; + metal::vec w_0 = {cos_pi_4, -cos_pi_4}; + metal::vec w_1 = {-cos_pi_4, -cos_pi_4}; + metal::vec temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; + radix4(temp, x); + radix4(temp + 4, x + 4); y[0] = x[0] + x[4]; y[4] = x[0] - x[4]; - float2 x_5 = complex_mul(x[5], w_0); + metal::vec x_5 = complex_mul(x[5], w_0); y[1] = x[1] + x_5; y[5] = x[1] - x_5; - float2 x_6 = {x[6].y, -x[6].x}; + metal::vec x_6 = {x[6].y, -x[6].x}; y[2] = x[2] + x_6; y[6] = x[2] - x_6; - float2 x_7 = complex_mul(x[7], w_1); + metal::vec x_7 = complex_mul(x[7], w_1); y[3] = x[3] + x_7; y[7] = x[3] - x_7; } -template -METAL_FUNC void radix10(thread float2* x, thread float2* y) { - float2 w[4]; +template +METAL_FUNC void radix10( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[4]; w[0] = {0.8090169943749475, -0.5877852522924731}; w[1] = {0.30901699437494745, -0.9510565162951535}; w[2] = {-w[1].x, w[1].y}; w[3] = {-w[0].x, w[0].y}; if (raders_perm) { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[3], x[4], x[8], x[2], x[1], x[7], x[9], x[6], x[5]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } else { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[2], x[4], x[6], x[8], x[1], x[3], x[5], x[7], x[9]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } y[0] = x[0] + x[5]; y[5] = x[0] - x[5]; for (int t = 1; t < 5; t++) { - float2 a = complex_mul(x[t + 5], w[t - 1]); + metal::vec a = complex_mul(x[t + 5], w[t - 1]); y[t] = x[t] + a; y[t + 5] = x[t] - a; } } -METAL_FUNC void radix11(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 10.0, -1 / 10.0}; +template +METAL_FUNC void radix11( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 10.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix10(x + 1, y + 1); + radix10(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(0.955301878, -3.17606649)); - y[3] = complex_mul_conj(y[3], float2(2.63610556, 2.01269656)); - y[4] = complex_mul_conj(y[4], float2(2.54127802, 2.13117479)); - y[5] = complex_mul_conj(y[5], float2(2.07016210, 2.59122150)); - y[6] = complex_mul_conj(y[6], float2(0, -3.31662479)); - y[7] = complex_mul_conj(y[7], float2(2.07016210, -2.59122150)); - y[8] = complex_mul_conj(y[8], float2(-2.54127802, 2.13117479)); - y[9] = complex_mul_conj(y[9], float2(2.63610556, -2.01269656)); - y[10] = complex_mul_conj(y[10], float2(-0.955301878, -3.17606649)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(0.955301878, -3.17606649)); + y[3] = complex_mul_conj(y[3], metal::vec(2.63610556, 2.01269656)); + y[4] = complex_mul_conj(y[4], metal::vec(2.54127802, 2.13117479)); + y[5] = complex_mul_conj(y[5], metal::vec(2.07016210, 2.59122150)); + y[6] = complex_mul_conj(y[6], metal::vec(0, -3.31662479)); + y[7] = complex_mul_conj(y[7], metal::vec(2.07016210, -2.59122150)); + y[8] = complex_mul_conj(y[8], metal::vec(-2.54127802, 2.13117479)); + y[9] = complex_mul_conj(y[9], metal::vec(2.63610556, -2.01269656)); + y[10] = + complex_mul_conj(y[10], metal::vec(-0.955301878, -3.17606649)); // ifft - radix10(y + 1, x + 1); + radix10(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[6] = x[2] * inv + x[0]; @@ -234,10 +260,12 @@ METAL_FUNC void radix11(thread float2* x, thread float2* y) { y[2] = x[10] * inv + x[0]; } -template -METAL_FUNC void radix12(thread float2* x, thread float2* y) { - float2 w[6]; - float sin_pi_3 = 0.8660254037844387; +template +METAL_FUNC void radix12( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[6]; + T sin_pi_3 = 0.8660254037844387; w[0] = {sin_pi_3, -0.5}; w[1] = {0.5, -sin_pi_3}; w[2] = {0, -1}; @@ -245,7 +273,7 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { w[4] = {-sin_pi_3, -0.5}; if (raders_perm) { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[3], x[2], @@ -258,10 +286,10 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[10], x[4], x[6]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } else { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[2], x[4], @@ -274,44 +302,50 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[7], x[9], x[11]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } y[0] = x[0] + x[6]; y[6] = x[0] - x[6]; for (int t = 1; t < 6; t++) { - float2 a = complex_mul(x[t + 6], w[t - 1]); + metal::vec a = complex_mul(x[t + 6], w[t - 1]); y[t] = x[t] + a; y[t + 6] = x[t] - a; } } -METAL_FUNC void radix13(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 12.0, -1 / 12.0}; +template +METAL_FUNC void radix13( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 12.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix12(x + 1, y + 1); + radix12(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(3.07497206, -1.88269669)); - y[3] = complex_mul_conj(y[3], float2(3.09912468, 1.84266823)); - y[4] = complex_mul_conj(y[4], float2(3.45084438, -1.04483161)); - y[5] = complex_mul_conj(y[5], float2(0.91083583, 3.48860690)); - y[6] = complex_mul_conj(y[6], float2(-3.60286363, 0.139189267)); - y[7] = complex_mul_conj(y[7], float2(3.60555128, 0)); - y[8] = complex_mul_conj(y[8], float2(3.60286363, 0.139189267)); - y[9] = complex_mul_conj(y[9], float2(0.91083583, -3.48860690)); - y[10] = complex_mul_conj(y[10], float2(-3.45084438, -1.04483161)); - y[11] = complex_mul_conj(y[11], float2(3.09912468, -1.84266823)); - y[12] = complex_mul_conj(y[12], float2(-3.07497206, -1.88269669)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(3.07497206, -1.88269669)); + y[3] = complex_mul_conj(y[3], metal::vec(3.09912468, 1.84266823)); + y[4] = complex_mul_conj(y[4], metal::vec(3.45084438, -1.04483161)); + y[5] = complex_mul_conj(y[5], metal::vec(0.91083583, 3.48860690)); + y[6] = complex_mul_conj(y[6], metal::vec(-3.60286363, 0.139189267)); + y[7] = complex_mul_conj(y[7], metal::vec(3.60555128, 0)); + y[8] = complex_mul_conj(y[8], metal::vec(3.60286363, 0.139189267)); + y[9] = complex_mul_conj(y[9], metal::vec(0.91083583, -3.48860690)); + y[10] = + complex_mul_conj(y[10], metal::vec(-3.45084438, -1.04483161)); + y[11] = complex_mul_conj(y[11], metal::vec(3.09912468, -1.84266823)); + y[12] = + complex_mul_conj(y[12], metal::vec(-3.07497206, -1.88269669)); // ifft - radix12(y + 1, x + 1); + radix12(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[7] = x[2] * inv + x[0]; @@ -325,4 +359,4 @@ METAL_FUNC void radix13(thread float2* x, thread float2* y) { y[8] = x[10] * inv + x[0]; y[4] = x[11] * inv + x[0]; y[2] = x[12] * inv + x[0]; -} \ No newline at end of file +} diff --git a/mlx/backend/metal/kernels/fft/readwrite.h b/mlx/backend/metal/kernels/fft/readwrite.h index 3d1b23f4ed..4e954e29fb 100644 --- a/mlx/backend/metal/kernels/fft/readwrite.h +++ b/mlx/backend/metal/kernels/fft/readwrite.h @@ -27,14 +27,43 @@ Each with support for: using namespace metal; +// Derives the FFT scalar arithmetic type from a storage type. +template +struct FFTStorageTraits { + using scalar_T = storage_T; +}; + +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + +template +struct FFTIOTypeTraits { + using scalar_T = typename FFTStorageTraits::scalar_T; + + static_assert( + metal::is_same_v::scalar_T>, + "FFT input and output storage must share a scalar arithmetic type"); + static_assert( + sizeof(in_T) <= 16 && 16 % sizeof(in_T) == 0, + "FFT input storage must divide the 128-bit sequential access width"); + static_assert( + sizeof(out_T) <= 16 && 16 % sizeof(out_T) == 0, + "FFT output storage must divide the 128-bit sequential access width"); +}; + template < typename in_T, typename out_T, int step = 0, bool four_step_real = false> struct ReadWriter { + using scalar_T = typename FFTIOTypeTraits::scalar_T; + using complex_T = vec; + const device in_T* in; - threadgroup float2* buf; + threadgroup complex_T* buf; device out_T* out; int n; int batch_size; @@ -50,7 +79,7 @@ struct ReadWriter { METAL_FUNC ReadWriter( const device in_T* in_, - threadgroup float2* buf_, + threadgroup complex_T* buf_, device out_T* out_, const short n_, const int batch_size_, @@ -73,21 +102,21 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const thread { - return inv ? float2(elem.x, -elem.y) : elem; + METAL_FUNC complex_T post_in(complex_T elem) const thread { + return inv ? complex_T(elem.x, -elem.y) : elem; } // Handle float case for generic RFFT alg - METAL_FUNC float2 post_in(float elem) const thread { - return float2(elem, 0); + METAL_FUNC complex_T post_in(scalar_T elem) const thread { + return complex_T(elem, 0); } - METAL_FUNC float2 pre_out(float2 elem) const thread { - return inv ? float2(elem.x / n, -elem.y / n) : elem; + METAL_FUNC complex_T pre_out(complex_T elem) const thread { + return inv ? complex_T(elem.x / n, -elem.y / n) : elem; } - METAL_FUNC float2 pre_out(float2 elem, int length) const thread { - return inv ? float2(elem.x / length, -elem.y / length) : elem; + METAL_FUNC complex_T pre_out(complex_T elem, int length) const thread { + return inv ? complex_T(elem.x / length, -elem.y / length) : elem; } METAL_FUNC bool out_of_bounds() const thread { @@ -99,21 +128,22 @@ struct ReadWriter { METAL_FUNC void load() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - // 2 complex64s = 128 bits - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + // Keep each thread's sequential access at 128 bits where possible. + constexpr int read_width = 16 / sizeof(in_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - buf[index] = post_in(in[batch_idx + index]); - buf[index + 1] = post_in(in[batch_idx + index + 1]); + for (short r = 0; r < read_width; r++) { + buf[index + r] = post_in(in[batch_idx + index + r]); + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); buf[index] = post_in(in[batch_idx + index]); } @@ -122,57 +152,60 @@ struct ReadWriter { METAL_FUNC void write() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + constexpr int read_width = 16 / sizeof(out_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - out[batch_idx + index] = pre_out(buf[index]); - out[batch_idx + index + 1] = pre_out(buf[index + 1]); + for (short r = 0; r < read_width; r++) { + out[batch_idx + index + r] = pre_out(buf[index + r]); + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); out[batch_idx + index] = pre_out(buf[index]); } } // Padded IO for Bluestein's algorithm - METAL_FUNC void load_padded(int length, const device float2* w_k) + METAL_FUNC void load_padded(int length, const device complex_T* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = post_in(in[batch_idx + index]); - seq_buf[index] = complex_mul(elem, w_k[index]); + complex_T elem = post_in(in[batch_idx + index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0.0; } } } - METAL_FUNC void write_padded(int length, const device float2* w_k) + METAL_FUNC void write_padded(int length, const device complex_T* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - float2 inv_factor = {1.0f / n, -1.0f / n}; + scalar_T inv_n = static_cast(1.0f / n); + complex_T inv_factor = {inv_n, -inv_n}; - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = seq_buf[index + length - 1] * inv_factor; - out[batch_idx + index] = pre_out(complex_mul(elem, w_k[index]), length); + complex_T elem = seq_buf[index + length - 1] * inv_factor; + out[batch_idx + index] = + pre_out(complex_mul(elem, w_k[index]), length); } } } @@ -199,53 +232,49 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step + // Four-step FFT I/O METAL_FUNC void load_strided(int stride, int overall_n) thread { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - buf[strided_shared_idx + e] = - post_in(in[strided_device_idx + e * stride]); + if constexpr (step == 1 && !four_step_real) { + // Do not invert between C2C four-step passes. + (void)stride; + (void)overall_n; + bool default_inv = inv; + inv = false; + load(); + inv = default_inv; + } else { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + buf[strided_shared_idx + e] = + post_in(in[strided_device_idx + e * stride]); + } } } METAL_FUNC void write_strided(int stride, int overall_n) thread { - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - int combined_idx = (strided_device_idx + e * stride) % overall_n; - int ij = (combined_idx / stride) * (combined_idx % stride); - // Apply four step twiddles at end of first step - float2 twiddle = get_twiddle(ij, overall_n); - out[strided_device_idx + e * stride] = complex_mul(output, twiddle); + if constexpr (step == 1 && !four_step_real) { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + out[strided_device_idx + e * stride] = + pre_out(buf[strided_shared_idx + e], overall_n); + } + } else { + for (int e = 0; e < elems_per_thread; e++) { + complex_T output = buf[strided_shared_idx + e]; + int combined_idx = (strided_device_idx + e * stride) % overall_n; + int ij = (combined_idx / stride) * (combined_idx % stride); + // Apply four step twiddles at end of first step + complex_T twiddle = get_twiddle(ij, overall_n); + out[strided_device_idx + e * stride] = + complex_mul(output, twiddle); + } } } }; -// Four Step FFT Second Step -template <> -METAL_FUNC void ReadWriter::load_strided( - int stride, - int overall_n) thread { - // Silence compiler warnings - (void)stride; - (void)overall_n; - // Don't invert between steps - bool default_inv = inv; - inv = false; - load(); - inv = default_inv; -} - -template <> -METAL_FUNC void ReadWriter::write_strided( - int stride, - int overall_n) thread { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - out[strided_device_idx + e * stride] = pre_out(output, overall_n); - } -} - +// Packed RFFT/IRFFT remains float-specific. This generic foundation covers +// C2C storage; reduced-precision real transforms need separate host dispatch. +// // For RFFT, we interleave batches of two real sequences into one complex one: // // z_k = x_k + j.y_k @@ -310,7 +339,7 @@ METAL_FUNC void ReadWriter::write() const thread { float2 x_n_minus_k = seq_buf[n - index] * conj; out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -335,7 +364,7 @@ METAL_FUNC void ReadWriter::load_padded( if (index < length) { float2 elem = float2(in[batch_idx + index], in[batch_idx + index + next_in]); - seq_buf[index] = complex_mul(elem, w_k[index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0; } @@ -368,18 +397,18 @@ METAL_FUNC void ReadWriter::write_padded( // x_0 = z_0.real // y_0 = z_0.imag if (index == 0) { - float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); out[batch_idx + index] = float2(elem.x, 0); out[batch_idx + index + next_out] = float2(elem.y, 0); } else { - float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); - float2 x_n_minus_k = complex_mul( + float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 x_n_minus_k = complex_mul( w_k[length - index], seq_buf[length - index] * inv_factor); x_n_minus_k *= conj; // w_k should happen before this extraction out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -426,10 +455,10 @@ METAL_FUNC void ReadWriter::load() const thread { x = float2(x.x, 0); y = float2(y.x, 0); } - seq_buf[index] = x + complex_mul(y, plus_j); + seq_buf[index] = x + complex_mul(y, plus_j); seq_buf[index].y = -seq_buf[index].y; if (index > 0 && !last_val) { - seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); + seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[n - index].y = -seq_buf[n - index].y; } } @@ -487,12 +516,12 @@ METAL_FUNC void ReadWriter::load_padded( x = float2(x.x, 0); y = float2(y.x, 0); } - float2 elem1 = x + complex_mul(y, plus_j); - seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); + float2 elem1 = x + complex_mul(y, plus_j); + seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); if (index > 0 && !last_val) { - float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); + float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[length - index] = - complex_mul(elem2 * conj, w_k[length - index]); + complex_mul(elem2 * conj, w_k[length - index]); } } else { short pad_index = metal::min(length + (index - length_over_2) * 2, n - 2); @@ -520,7 +549,8 @@ METAL_FUNC void ReadWriter::write_padded( for (int e = 0; e < elems_per_thread; e++) { int index = fft_idx + e * m; if (index < length) { - float2 output = complex_mul(seq_buf[index] * inv_factor, w_k[index]); + float2 output = + complex_mul(seq_buf[index] * inv_factor, w_k[index]); out[batch_idx + index] = output.x / length; out[batch_idx + index + next_out] = output.y / -length; } From f754f3170555c28b9e95f84e574bf797ea9e01a7 Mon Sep 17 00:00:00 2001 From: Gusanidas <33495733+Gusanidas@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:35:39 +0200 Subject: [PATCH 079/222] Fix segfault in expand_dims for out of bounds negative axes (#4021) --- mlx/ops.cpp | 94 ++++++++++++++++++---------------------- python/tests/test_ops.py | 18 ++++++++ tests/ops_tests.cpp | 18 ++++++++ 3 files changed, 78 insertions(+), 52 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index cdc6e2e7d6..bcfd39aa8e 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -562,26 +562,50 @@ array hadamard_transform( {astype(a, dtype, s)}); } -array squeeze_impl( +namespace { + +std::vector normalize_squeeze_axes( const array& a, - std::vector axes, - StreamOrDevice s /* = {} */) { - for (auto& ax : axes) { - auto new_ax = ax < 0 ? ax + a.ndim() : ax; - if (new_ax < 0 || new_ax >= a.ndim()) { - std::ostringstream msg; - msg << "[squeeze] Invalid axes " << ax << " for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + const std::vector& axes) { + int ndim = a.ndim(); + std::set unique_axes; + for (auto ax : axes) { + int new_ax = normalize_axis_index(ax, ndim, "[squeeze] "); if (a.shape(new_ax) != 1) { std::ostringstream msg; msg << "[squeeze] Cannot squeeze axis " << ax << " with size " - << a.shape(ax) << " which is not equal to 1."; + << a.shape(new_ax) << " which is not equal to 1."; throw std::invalid_argument(msg.str()); } - ax = new_ax; + unique_axes.insert(new_ax); + } + if (unique_axes.size() != axes.size()) { + throw std::invalid_argument("[squeeze] Received duplicate axes."); + } + return std::vector(unique_axes.begin(), unique_axes.end()); +} + +std::vector normalize_expand_dims_axes( + const array& a, + const std::vector& axes) { + int out_ndim = a.ndim() + axes.size(); + std::set unique_axes; + for (auto ax : axes) { + unique_axes.insert(normalize_axis_index(ax, out_ndim, "[expand_dims] ")); } + if (unique_axes.size() != axes.size()) { + throw std::invalid_argument("[expand_dims] Received duplicate axes."); + } + return std::vector(unique_axes.begin(), unique_axes.end()); +} + +} // namespace + +// Assumes the axes are non-negative, sorted, unique, and valid for a. +array squeeze_impl( + const array& a, + std::vector axes, + StreamOrDevice s /* = {} */) { auto shape = Squeeze::output_shape(a, axes); return array( std::move(shape), @@ -597,19 +621,11 @@ array squeeze( if (axes.empty()) { return a; } - std::set unique_axes; - for (auto ax : axes) { - unique_axes.insert(ax < 0 ? ax + a.ndim() : ax); - } - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[squeeze] Received duplicate axes."); - } - std::vector sorted_axes(unique_axes.begin(), unique_axes.end()); - return squeeze_impl(a, std::move(sorted_axes), s); + return squeeze_impl(a, normalize_squeeze_axes(a, axes), s); } array squeeze(const array& a, int axis, StreamOrDevice s /* = {} */) { - return squeeze_impl(a, {axis}, s); + return squeeze_impl(a, normalize_squeeze_axes(a, {axis}), s); } array squeeze(const array& a, StreamOrDevice s /* = {} */) { @@ -622,21 +638,11 @@ array squeeze(const array& a, StreamOrDevice s /* = {} */) { return squeeze_impl(a, std::move(axes), s); } +// Assumes the axes are non-negative, sorted, unique and valid for the output. array expand_dims_impl( const array& a, std::vector axes, StreamOrDevice s /* = {} */) { - auto out_ndim = a.ndim() + axes.size(); - for (auto& ax : axes) { - auto new_ax = ax < 0 ? ax + out_ndim : ax; - if (new_ax < 0 || new_ax >= out_ndim) { - std::ostringstream msg; - msg << "[expand_dims] Invalid axis " << ax << " for output array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - ax = new_ax; - } auto shape = ExpandDims::output_shape(a, axes); return array( std::move(shape), @@ -646,7 +652,7 @@ array expand_dims_impl( } array expand_dims(const array& a, int axis, StreamOrDevice s /* = {} */) { - return expand_dims_impl(a, {axis}, s); + return expand_dims_impl(a, normalize_expand_dims_axes(a, {axis}), s); } array expand_dims( @@ -656,23 +662,7 @@ array expand_dims( if (axes.empty()) { return a; } - { // Check for repeats - std::set unique_axes(axes.begin(), axes.end()); - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[expand_dims] Received duplicate axes."); - } - } - // Check for repeats again - auto out_ndim = a.ndim() + axes.size(); - std::set unique_axes; - for (auto ax : axes) { - unique_axes.insert(ax < 0 ? ax + out_ndim : ax); - } - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[expand_dims] Received duplicate axes."); - } - std::vector sorted_axes(unique_axes.begin(), unique_axes.end()); - return expand_dims_impl(a, std::move(sorted_axes), s); + return expand_dims_impl(a, normalize_expand_dims_axes(a, axes), s); } array flip( diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 214c59f542..67f8b0ae78 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2352,6 +2352,24 @@ def test_squeeze_expand(self): self.assertEqual(mx.expand_dims(a, 0).shape, (1, 2, 2)) self.assertEqual(mx.expand_dims(a, (0, 1)).shape, (1, 1, 2, 2)) self.assertEqual(mx.expand_dims(a, [0, -1]).shape, (1, 2, 2, 1)) + self.assertEqual(mx.expand_dims(a, [-1, -4]).shape, (1, 2, 2, 1)) + + def test_squeeze_expand_invalid_axes(self): + # Out of bounds negative axes must raise instead of wrapping around + a = mx.zeros(()) + self.assertEqual(mx.expand_dims(a, (-2, -1)).shape, (1, 1)) + with self.assertRaises(ValueError): + mx.expand_dims(a, (-3, -2)) + + a = mx.zeros((2, 2)) + for axes in [(-5, -4), (-6, 0), (0, 5)]: + with self.assertRaises(ValueError): + mx.expand_dims(a, axes) + + a = mx.zeros((1, 1, 1)) + for axes in [(-4,), (-5, 0), (0, 4)]: + with self.assertRaises(ValueError): + mx.squeeze(a, axes) def test_sort(self): shape = (6, 4, 10) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 1f82a9946c..0dd9385e14 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -193,6 +193,12 @@ TEST_CASE("test squeeze and expand") { CHECK_THROWS(squeeze(x, {1, 3, 1})); CHECK_THROWS(squeeze(x, {1, 3, -3})); + // Out of bounds negative axes must throw and not wrap around + x = zeros({1, 1, 1}); + CHECK_THROWS(squeeze(x, std::vector{-4})); + CHECK_THROWS(squeeze(x, {-5, 0})); + CHECK_THROWS(squeeze(x, {0, 4})); + x = zeros({2, 2}); CHECK_EQ(expand_dims(x, 0).shape(), Shape{1, 2, 2}); CHECK_EQ(expand_dims(x, -1).shape(), Shape{2, 2, 1}); @@ -206,6 +212,18 @@ TEST_CASE("test squeeze and expand") { CHECK_THROWS(expand_dims(x, -4)); CHECK_THROWS(expand_dims(x, {0, 1, 0})); CHECK_THROWS(expand_dims(x, {0, 1, -4})); + + // Negative axes are resolved against the output shape and sorted + CHECK_EQ(expand_dims(x, {3, -4}).shape(), Shape{1, 2, 2, 1}); + CHECK_EQ(expand_dims(x, {-1, -4}).shape(), Shape{1, 2, 2, 1}); + + // Out of bounds negative axes must throw and not wrap around + CHECK_THROWS(expand_dims(x, {-5, -4})); + CHECK_THROWS(expand_dims(x, {-6, 0})); + + x = zeros({}); + CHECK_EQ(expand_dims(x, {-2, -1}).shape(), Shape{1, 1}); + CHECK_THROWS(expand_dims(x, {-3, -2})); } TEST_CASE("test slice") { From d9bd3c26e4112567bb0331840632c568f7b06881 Mon Sep 17 00:00:00 2001 From: Tanish Jain Date: Fri, 7 Aug 2026 06:05:49 +0530 Subject: [PATCH 080/222] Add optional dtype parameter to zeros_like and ones_like (#4028) --- mlx/ops.cpp | 12 ++++++++++-- mlx/ops.h | 4 ++++ python/src/ops.cpp | 26 ++++++++++++++++++++------ python/tests/test_ops.py | 35 +++++++++++++++++++++++++++++++++++ tests/creations_tests.cpp | 8 ++++++++ 5 files changed, 77 insertions(+), 8 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index bcfd39aa8e..e41434e24d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -363,7 +363,11 @@ array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { } array zeros_like(const array& a, StreamOrDevice s /* = {} */) { - return full_like(a, 0, a.dtype(), to_stream(s)); + return zeros_like(a, a.dtype(), s); +} + +array zeros_like(const array& a, Dtype dtype, StreamOrDevice s /* = {} */) { + return full_like(a, 0, dtype, to_stream(s)); } array ones(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { @@ -371,7 +375,11 @@ array ones(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { } array ones_like(const array& a, StreamOrDevice s /* = {} */) { - return full_like(a, 1, a.dtype(), to_stream(s)); + return ones_like(a, a.dtype(), s); +} + +array ones_like(const array& a, Dtype dtype, StreamOrDevice s /* = {} */) { + return full_like(a, 1, dtype, to_stream(s)); } array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s /* = {} */) { diff --git a/mlx/ops.h b/mlx/ops.h index 20568024f6..8297f77ca8 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -93,14 +93,18 @@ MLX_API array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array zeros(const Shape& shape, StreamOrDevice s = {}) { return zeros(shape, float32, s); } +/** Create an array of zeros with the shape of `a`. */ MLX_API array zeros_like(const array& a, StreamOrDevice s = {}); +MLX_API array zeros_like(const array& a, Dtype dtype, StreamOrDevice s = {}); /** Fill an array of the given shape with ones. */ MLX_API array ones(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array ones(const Shape& shape, StreamOrDevice s = {}) { return ones(shape, float32, s); } +/** Create an array of ones with the shape of `a`. */ MLX_API array ones_like(const array& a, StreamOrDevice s = {}); +MLX_API array ones_like(const array& a, Dtype dtype, StreamOrDevice s = {}); /** Fill an array of the given shape (n,m) with ones in the specified diagonal * k, and zeros everywhere else. */ diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 7057075c9b..4e0373a238 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1959,17 +1959,24 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "zeros_like", - &mx::zeros_like, + [](const mx::array& a, + std::optional dtype, + mx::StreamOrDevice s) { + return mx::zeros_like(a, dtype.value_or(a.dtype()), s); + }, nb::arg(), + "dtype"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def zeros_like(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def zeros_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( An array of zeros like the input. Args: - a (array): The input to take the shape and type from. + a (array): The input to take the shape from. + dtype (Dtype, optional): Output data type. If ``None``, the output + type defaults to the input array's data type. Returns: array: The output array filled with zeros. @@ -2001,17 +2008,24 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "ones_like", - &mx::ones_like, + [](const mx::array& a, + std::optional dtype, + mx::StreamOrDevice s) { + return mx::ones_like(a, dtype.value_or(a.dtype()), s); + }, nb::arg(), + "dtype"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ones_like(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ones_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( An array of ones like the input. Args: - a (array): The input to take the shape and type from. + a (array): The input to take the shape from. + dtype (Dtype, optional): Output data type. If ``None``, the output + type defaults to the input array's data type. Returns: array: The output array filled with ones. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 67f8b0ae78..42565608fd 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3568,6 +3568,41 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals)) self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) + def test_zeros_ones_empty_like_dtype(self): + x = mx.array([1, 2, 3], dtype=mx.int32) + + # Default dtype (should match x) + z = mx.zeros_like(x) + self.assertEqual(z.dtype, mx.int32) + o = mx.ones_like(x) + self.assertEqual(o.dtype, mx.int32) + e = mx.empty_like(x) + self.assertEqual(e.dtype, mx.int32) + + # Positional dtype + z = mx.zeros_like(x, mx.float16) + self.assertEqual(z.dtype, mx.float16) + self.assertTrue(mx.array_equal(z, mx.zeros((3,), mx.float16))) + + o = mx.ones_like(x, mx.float16) + self.assertEqual(o.dtype, mx.float16) + self.assertTrue(mx.array_equal(o, mx.ones((3,), mx.float16))) + + e = mx.empty_like(x, mx.float16) + self.assertEqual(e.dtype, mx.float16) + + # Keyword dtype + z = mx.zeros_like(x, dtype=mx.float32) + self.assertEqual(z.dtype, mx.float32) + self.assertTrue(mx.array_equal(z, mx.zeros((3,), mx.float32))) + + o = mx.ones_like(x, dtype=mx.float32) + self.assertEqual(o.dtype, mx.float32) + self.assertTrue(mx.array_equal(o, mx.ones((3,), mx.float32))) + + e = mx.empty_like(x, dtype=mx.float32) + self.assertEqual(e.dtype, mx.float32) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/tests/creations_tests.cpp b/tests/creations_tests.cpp index ea43bd0e23..b1973cad65 100644 --- a/tests/creations_tests.cpp +++ b/tests/creations_tests.cpp @@ -218,10 +218,18 @@ TEST_CASE("test full") { CHECK_EQ(y.dtype(), int32); CHECK(array_equal(x, y).item()); + auto z = zeros_like(x, float16); + CHECK_EQ(z.dtype(), float16); + CHECK(array_equal(z, zeros({2, 2}, float16)).item()); + x = ones({2, 2}, int32); y = ones_like(x); CHECK_EQ(y.dtype(), int32); CHECK(array_equal(x, y).item()); + + z = ones_like(x, float16); + CHECK_EQ(z.dtype(), float16); + CHECK(array_equal(z, ones({2, 2}, float16)).item()); } // Works for empty shape and empty array From 1070373954e04bd8c7eafb8ccd9024fdf35aff49 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:01:29 +0800 Subject: [PATCH 081/222] Fix mx.longsumexp output shape issue. (#4030) --- mlx/ops.cpp | 2 +- python/tests/test_ops.py | 41 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index e41434e24d..b55d3e93d9 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2862,7 +2862,7 @@ array logsumexp( std::make_shared(to_stream(s)), {astype(a, dtype, s)}); if (!keepdims) { - out = squeeze(out, -1, s); + out = squeeze(out, axes, s); } return out; } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 42565608fd..f093891a15 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -840,6 +840,47 @@ def logsumexp(x, axes=None): x = mx.broadcast_to(mx.random.uniform(shape=(2, 1, 8)), (2, 2, 8)) self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x))) + def test_logsumexp_shape(self): + # A reduction over all axes with keepdims=False must yield a scalar + # array (), consistent with every other reduction (sum, prod, max, ...). + # Regression test: logsumexp's fast path used to squeeze only the last + # axis, so inputs with size-1 leading dims returned (1,) instead of (). + def logsumexp(x, axes=None): + maxs = mx.max(x, axis=axes, keepdims=True) + return mx.log(mx.sum(mx.exp(x - maxs), axis=axes, keepdims=True)) + maxs + + # Full reduction on size-1-leading-dim inputs -> scalar (). + for shape in [(1, 2), (1, 1), (1, 5), (1, 1, 8), (1, 1, 1, 4)]: + x = mx.random.uniform(shape=shape) + out = mx.logsumexp(x) + self.assertEqual(out.shape, (), f"shape {shape}") + self.assertEqual(out.ndim, 0, f"shape {shape}") + # Same shape as mx.sum and same value as the decomposition. + self.assertEqual(out.shape, mx.sum(x).shape) + self.assertTrue(mx.allclose(out, logsumexp(x))) + + # keepdims=True keeps every reduced axis as size 1. + self.assertEqual( + mx.logsumexp(mx.random.uniform(shape=(1, 5)), keepdims=True).shape, (1, 1) + ) + self.assertEqual( + mx.logsumexp(mx.random.uniform(shape=(1, 1, 8)), keepdims=True).shape, + (1, 1, 1), + ) + + # Partial reductions over a contiguous suffix of axes with size-1 + # leading dims must squeeze every reduced axis, not only the last one. + x = mx.random.uniform(shape=(5, 1, 8)) + self.assertEqual(mx.logsumexp(x, axis=[1, 2]).shape, (5,)) + self.assertEqual(mx.logsumexp(x, axis=[1, 2], keepdims=True).shape, (5, 1, 1)) + ref = np.logaddexp.reduce(np.array(x), axis=(1, 2)) + self.assertTrue(np.allclose(np.array(mx.logsumexp(x, axis=[1, 2])), ref)) + + self.assertEqual( + mx.logsumexp(mx.random.uniform(shape=(1, 1, 1, 8)), axis=[1, 2, 3]).shape, + (1,), + ) + def test_mean(self): x = mx.array( [ From 0d5bc64a452fc689fb882bc86866eaf5e8dc78bc Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 20:19:21 -0700 Subject: [PATCH 082/222] Treat backend and envs as optional when parsing a hostfile (#4039) --- python/mlx/_distributed_utils/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mlx/_distributed_utils/common.py b/python/mlx/_distributed_utils/common.py index c7e0fdd43c..6289da0c46 100644 --- a/python/mlx/_distributed_utils/common.py +++ b/python/mlx/_distributed_utils/common.py @@ -69,8 +69,8 @@ def from_file(cls, hostfile): envs = [] hosts = [] if isinstance(data, dict): - backend = data["backend"] - envs = data["envs"] + backend = data.get("backend", backend) + envs = data.get("envs", envs) hosts = data["hosts"] elif isinstance(data, list): hosts = data From d32978b9b2e869438ff183cb4cbbfd6e74bbf2f9 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 6 Aug 2026 20:19:35 -0700 Subject: [PATCH 083/222] chore: Report bad hostfiles as CLI errors in mlx.launch (#4040) --- python/mlx/_distributed_utils/launch.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 3b6af453e2..b052643aaf 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -332,7 +332,7 @@ def launch_ring(parser, hosts, args, command): def launch_nccl(parser, hosts, args, command): if not hosts[0].ips: - raise ValueError("Rank 0 should have an IP reachable from all other ranks") + parser.error("Rank 0 should have an IP reachable from all other ranks") master_host = hosts[0].ips[0] master_port = args.nccl_port @@ -371,13 +371,13 @@ def launch_nccl(parser, hosts, args, command): def launch_jaccl(parser, hosts, args, command): if not hosts[0].ips: - raise ValueError("Rank 0 should have an IP reachable from all other ranks") + parser.error("Rank 0 should have an IP reachable from all other ranks") jaccl_ring = args.backend == "jaccl-ring" have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts) have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts)) if not have_rdmas or not have_nulls: - raise ValueError("Malformed hostfile for jaccl backend") + parser.error("Malformed hostfile for jaccl backend") coordinator = hosts[0].ips[0] env = args.env From 158118bb152f59ad1aef0d8c835c49317ff9637b Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Thu, 6 Aug 2026 20:20:30 -0700 Subject: [PATCH 084/222] chore: Fix ceil error message to say ceil instead of floor (#4042) Co-authored-by: Cheng --- mlx/ops.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index b55d3e93d9..a7034d7595 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3099,7 +3099,7 @@ array floor(const array& a, StreamOrDevice s /* = {} */) { array ceil(const array& a, StreamOrDevice s /* = {} */) { if (a.dtype() == complex64) { - throw std::invalid_argument("[floor] Not supported for complex64."); + throw std::invalid_argument("[ceil] Not supported for complex64."); } return array(a.shape(), a.dtype(), std::make_shared(to_stream(s)), {a}); } From 447bb0f9ef20789894c7786e23358578db2f247b Mon Sep 17 00:00:00 2001 From: yingjiacai <8magino8@gmail.com> Date: Fri, 7 Aug 2026 11:55:52 +0800 Subject: [PATCH 085/222] Fix installed static MLX package on Windows (#3848) Co-authored-by: Yingjia Cai Co-authored-by: Cheng --- CMakeLists.txt | 14 +++++++++++--- mlx.pc.in | 19 +++++++++++++++++++ mlx/backend/cuda/CMakeLists.txt | 2 +- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index ca689e13d1..9d7b5baa30 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -262,7 +262,8 @@ if(WIN32) FetchContent_MakeAvailable(dlfcn-win32) endblock() target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src") - target_link_libraries(mlx PRIVATE dl) + target_sources(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src/dlfcn.c") + target_link_libraries(mlx PRIVATE Psapi.lib) endif() if(MLX_BUILD_CPU) @@ -297,10 +298,15 @@ if(MLX_BUILD_CPU) URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.33/${OPENBLAS_ZIP}" ) FetchContent_MakeAvailable(openblas) - target_link_libraries( - mlx PRIVATE "${openblas_SOURCE_DIR}/lib/${OPENBLAS_LIB}.lib") target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/${OPENBLAS_INC}") + # Make openblas importable by dependencies when built as static library. + set(OPENBLAS_IMPORT_LIBRARY + "${openblas_SOURCE_DIR}/lib/${OPENBLAS_LIB}.lib") + add_library(MLX::OpenBLAS UNKNOWN IMPORTED GLOBAL) + set_target_properties(MLX::OpenBLAS PROPERTIES IMPORTED_LOCATION + "${OPENBLAS_IMPORT_LIBRARY}") + target_link_libraries(mlx PRIVATE MLX::OpenBLAS) # Make sure the DLL file is placed in the same dir with executables. set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/${OPENBLAS_LIB}.dll") add_custom_command( @@ -418,6 +424,8 @@ if(WIN32) if(MLX_BUILD_CPU) # Install OpenBLAS. install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN) + install(FILES ${OPENBLAS_IMPORT_LIBRARY} + DESTINATION ${CMAKE_INSTALL_LIBDIR}) endif() endif() diff --git a/mlx.pc.in b/mlx.pc.in index 2889149db6..a202028b1f 100644 --- a/mlx.pc.in +++ b/mlx.pc.in @@ -11,6 +11,25 @@ @PACKAGE_INIT@ +include(CMakeFindDependencyMacro) + +if(WIN32 AND @MLX_BUILD_CPU@ AND NOT TARGET MLX::OpenBLAS) + add_library(MLX::OpenBLAS UNKNOWN IMPORTED) + set_target_properties(MLX::OpenBLAS PROPERTIES + IMPORTED_LOCATION "@PACKAGE_CMAKE_INSTALL_LIBDIR@/@OPENBLAS_LIB@.lib" + ) +endif() + +if(@MLX_BUILD_CUDA@) + find_dependency(CUDAToolkit REQUIRED) + + set(_MLX_SAVED_CMAKE_MODULE_PATH "${CMAKE_MODULE_PATH}") + list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}") + find_dependency(CUDNN REQUIRED) + set(CMAKE_MODULE_PATH "${_MLX_SAVED_CMAKE_MODULE_PATH}") + unset(_MLX_SAVED_CMAKE_MODULE_PATH) +endif() + include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/MLXTargets.cmake) include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/extension.cmake) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 065220e24e..a82c5ad6e9 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -277,7 +277,7 @@ set(CUDNN_FRONTEND_BUILD_SAMPLES OFF) set(CUDNN_FRONTEND_BUILD_TESTS OFF) set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF) FetchContent_MakeAvailable(cudnn) -target_link_libraries(mlx PRIVATE cudnn_frontend) +target_link_libraries(mlx PRIVATE $) # Link with the actual cuDNN libraries. target_link_libraries(mlx PRIVATE CUDNN::cudnn_all) From 8056817bd1202477b6bed2b1f871da5960d64b28 Mon Sep 17 00:00:00 2001 From: "Jae B." Date: Fri, 7 Aug 2026 02:47:09 -0400 Subject: [PATCH 086/222] Derive the qmv fast path K alignment from bits (#3965) --- mlx/backend/metal/quantized.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 57692dfcd5..c9f6a7e00e 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include "mlx/backend/common/quantized.h" #include "mlx/backend/common/broadcasting.h" #include "mlx/backend/common/compiled.h" #include "mlx/backend/gpu/copy.h" @@ -125,6 +126,12 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } } +// Must match the K step in qmv_fast_impl (kernels/quantized.h): +// pack_factor() * (bits == 2 ? 1 : 2) * SIMD_SIZE +inline int qmv_fast_k_alignment(int bits) { + return get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32; +} + inline int add_strides_and_shapes( CommandEncoder& compute_encoder, bool skip, @@ -460,7 +467,7 @@ void qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % 512 == 0; + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; concatenate( kname, @@ -1300,7 +1307,7 @@ void gather_qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % 512 == 0; + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; concatenate( kname, mode + (fast ? "_gather_qmv_fast_" : "_gather_qmv_"), From 383fe160c72dbbcfa33ed012e6fa7ed8762b2326 Mon Sep 17 00:00:00 2001 From: rohith Date: Sat, 8 Aug 2026 05:11:56 +0530 Subject: [PATCH 087/222] Fix eigh UPLO and zero-size eigh/svd on the CPU (#3834) --- mlx/backend/cpu/eigh.cpp | 17 +++++++++++++--- mlx/backend/cpu/svd.cpp | 29 +++++++++++++++++++++++++++ python/tests/test_linalg.py | 40 +++++++++++++++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/mlx/backend/cpu/eigh.cpp b/mlx/backend/cpu/eigh.cpp index d457c1fd99..813e096320 100644 --- a/mlx/backend/cpu/eigh.cpp +++ b/mlx/backend/cpu/eigh.cpp @@ -155,10 +155,12 @@ void eigh_impl( auto& encoder = cpu::get_command_encoder(stream); encoder.set_output_array(vectors); encoder.set_output_array(values); + // LAPACK reads the row-major input as its (conjugate) transpose, so the + // requested triangle is the opposite one in LAPACK's view. encoder.dispatch([vec_ptr, eig_ptr, jobz, - uplo = uplo[0], + uplo = uplo[0] == 'L' ? 'U' : 'L', N = vectors.shape(-1), size = vectors.size()]() mutable { // Work query @@ -190,12 +192,21 @@ void Eigh::eval_cpu( const auto& a = inputs[0]; auto& values = outputs[0]; + values.set_data(allocator::malloc(values.nbytes())); + + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. The input is + // square, so both outputs are empty and there is nothing to write. + if (a.shape(-1) == 0) { + if (compute_eigenvectors_) { + outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); + } + return; + } + auto vectors = compute_eigenvectors_ ? outputs[1] : array(a.shape(), a.dtype(), nullptr, {}); - values.set_data(allocator::malloc(values.nbytes())); - copy_cpu( a, vectors, diff --git a/mlx/backend/cpu/svd.cpp b/mlx/backend/cpu/svd.cpp index ca01a0a65a..932ededdb4 100644 --- a/mlx/backend/cpu/svd.cpp +++ b/mlx/backend/cpu/svd.cpp @@ -1,5 +1,7 @@ // Copyright © 2024 Apple Inc. +#include + #include "mlx/allocator.h" #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" @@ -206,6 +208,33 @@ void svd_impl( using R = typename SVDWork::R; + // Nothing to decompose when either dimension is zero; LAPACK rejects + // lda = 0. The singular values are empty but the factors are not when only + // one of the dimensions is zero, so fill them with the identity like LAPACK + // does, which keeps them orthonormal. + if (M == 0 || N == 0) { + auto& encoder = cpu::get_command_encoder(stream); + for (auto& o : outputs) { + o.set_data(allocator::malloc(o.nbytes())); + if (o.size() == 0) { + continue; + } + encoder.set_output_array(o); + encoder.dispatch([o = array::unsafe_weak_copy(o)]() mutable { + auto ptr = o.data(); + const size_t size = o.size(); + const int n = o.shape(-1); + std::fill_n(ptr, size, T(0)); + for (size_t i = 0; i < size; i += static_cast(n) * n) { + for (int j = 0; j < n; ++j) { + ptr[i + static_cast(j) * n + j] = T(1); + } + } + }); + } + return; + } + size_t num_matrices = a.size() / (M * N); // lapack clobbers the input, so we have to make a copy. diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index 39f0a3b891..a23de1a523 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -188,6 +188,23 @@ def test_svd_decomposition(self): ) ) + # Zero-size inputs. When only one of the dimensions is zero the + # factors are not empty and hold the identity, like numpy. + for shape in [(0, 4, 4), (3, 0, 0), (2, 5, 0), (2, 0, 5), (5, 0), (0, 5)]: + a_np = np.zeros(shape, dtype=np.float32) + U, S, Vt = mx.linalg.svd(mx.array(a_np), stream=mx.cpu) + mx.eval(U, S, Vt) + U_np, S_np, Vt_np = np.linalg.svd(a_np) + self.assertEqual(U.shape, U_np.shape) + self.assertEqual(S.shape, S_np.shape) + self.assertEqual(Vt.shape, Vt_np.shape) + self.assertTrue(np.array_equal(np.array(U), U_np)) + self.assertTrue(np.array_equal(np.array(Vt), Vt_np)) + + S_only = mx.linalg.svd(mx.array(a_np), compute_uv=False, stream=mx.cpu) + mx.eval(S_only) + self.assertEqual(S_only.shape, S_np.shape) + # Test float64 - use CPU stream since float64 is not supported on GPU with mx.stream(mx.cpu): A_f64 = mx.array( @@ -493,6 +510,29 @@ def check_eigs_and_vecs(A_np, kwargs={}): A_np = A_np + A_np.T.conj() check_eigs_and_vecs(A_np) + # UPLO picks the triangle like numpy; only observable when the two + # triangles disagree + A_np = np.array([[1.0, 999.0], [2.0, 3.0]], dtype=np.float32) + for uplo in ("L", "U"): + w = mx.linalg.eigvalsh(mx.array(A_np), UPLO=uplo, stream=mx.cpu) + w_np = np.linalg.eigvalsh(A_np, UPLO=uplo) + self.assertTrue(np.allclose(w, w_np, atol=1e-5)) + + # Zero-size inputs + for shape in [(0, 4, 4), (3, 0, 0), (0, 0)]: + a_np = np.zeros(shape, dtype=np.float32) + w, v = mx.linalg.eigh(mx.array(a_np), stream=mx.cpu) + mx.eval(w, v) + w_np, v_np = np.linalg.eigh(a_np) + self.assertEqual(w.shape, w_np.shape) + self.assertEqual(v.shape, v_np.shape) + self.assertTrue(np.array_equal(np.array(w), w_np)) + self.assertTrue(np.array_equal(np.array(v), v_np)) + + w_only = mx.linalg.eigvalsh(mx.array(a_np), stream=mx.cpu) + mx.eval(w_only) + self.assertEqual(w_only.shape, w_np.shape) + # Test error cases with self.assertRaises(ValueError): mx.linalg.eigh(mx.array([1.0, 2.0])) # 1D array From 3dc6e9b57e0b89c98949cec2f31468408b4edbda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Sat, 8 Aug 2026 04:02:29 +0300 Subject: [PATCH 088/222] Add metadata to exported functions (#3855) --- docs/src/usage/export.rst | 29 +++++++++++++ mlx/export.cpp | 66 ++++++++++++++++++++++-------- mlx/export.h | 18 +++++--- mlx/export_impl.h | 19 +++++++-- python/src/export.cpp | 62 +++++++++++++++++++++++----- python/tests/test_export_import.py | 37 +++++++++++++++++ tests/export_import_tests.cpp | 28 +++++++++++++ 7 files changed, 221 insertions(+), 38 deletions(-) diff --git a/docs/src/usage/export.rst b/docs/src/usage/export.rst index ac1ec218bb..9900ed6d24 100644 --- a/docs/src/usage/export.rst +++ b/docs/src/usage/export.rst @@ -109,6 +109,35 @@ keyword arguments when calling the imported function. out, = imported_fun(x, z=y) +Saving Metadata +--------------- + +You can save metadata, such as a model configuration, alongside an exported +function. The metadata is a string, so structured data can be encoded with +JSON: + +.. code-block:: python + + import json + + def fun(x, y): + return x + y + + x = mx.array(1.0) + y = mx.array(1.0) + config = {"description": "adds two arrays", "version": 1} + mx.export_function("add.mlxfn", fun, x, y, metadata=json.dumps(config)) + +Pass ``return_metadata=True`` to read the metadata back when importing: + +.. code-block:: python + + imported_fun, metadata = mx.import_function("add.mlxfn", return_metadata=True) + + # Prints: adds two arrays + print(json.loads(metadata)["description"]) + + Exporting Modules ----------------- diff --git a/mlx/export.cpp b/mlx/export.cpp index 4777fd3048..de5b78fd61 100644 --- a/mlx/export.cpp +++ b/mlx/export.cpp @@ -121,6 +121,9 @@ void serialize(Writer& os, T v) { } else if constexpr (std::is_enum_v) { serialize(os, static_cast(v)); } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v) { + serialize(os, static_cast(v.size())); + os.write(v.data(), v.size()); } else if constexpr (is_iterable) { serialize(os, static_cast(v.size())); for (const auto& t : v) { @@ -156,6 +159,19 @@ T deserialize(Reader& is) { return static_cast(deserialize(is)); } else if constexpr (std::is_same_v) { return nullptr; + } else if constexpr (std::is_same_v) { + auto size = deserialize(is); + // Bound the allocation by what the file can still contain + auto pos = is.tell(); + is.seek(0, std::ios_base::end); + auto remaining = static_cast(is.tell() - pos); + is.seek(pos); + if (size > remaining) { + throw std::runtime_error("[import_function] Invalid string size."); + } + T v(size, '\0'); + is.read(v.data(), size); + return v; } else if constexpr (is_iterable) { T v; auto size = deserialize(is); @@ -521,10 +537,15 @@ struct PrimitiveFactory { }; }; -void write_header(Writer& os, int count, bool shapeless) { +void write_header( + Writer& os, + int count, + bool shapeless, + const std::string& metadata) { serialize(os, std::string(version())); serialize(os, count); serialize(os, shapeless); + serialize(os, metadata); } // A struct to hold and retrieve the graphs that are exported / imported @@ -674,14 +695,16 @@ FunctionTable::Function* FunctionTable::find( FunctionExporter::FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless) + bool shapeless, + std::string metadata) : os(file), fun(std::move(fun)), - ftable(std::make_shared(shapeless)) { + ftable(std::make_shared(shapeless)), + metadata_(std::move(metadata)) { if (!os.is_open()) { throw std::runtime_error("[export_function] Failed to open " + file); } - write_header(os, count, shapeless); + write_header(os, count, shapeless, metadata_); } FunctionExporter::FunctionExporter( @@ -819,7 +842,7 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) { // Update the header auto pos = os.tell(); os.seek(0); - write_header(os, count, ftable->shapeless); + write_header(os, count, ftable->shapeless, metadata_); os.seek(pos); serialize(os, kwarg_keys); @@ -898,44 +921,51 @@ void FunctionExporter::operator()(const Args& args, const Kwargs& kwargs) { FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { return FunctionExporter{ file, [fun](const Args& args, const Kwargs&) { return fun(args); }, - shapeless}; + shapeless, + metadata}; } FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { return exporter( file, [fun](const Args&, const Kwargs kwargs) { return fun(kwargs); }, - shapeless); + shapeless, + metadata); } FunctionExporter exporter( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless /* = false */) { - return FunctionExporter{file, fun, shapeless}; + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + return FunctionExporter{file, fun, shapeless, metadata}; } void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(args); } void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(kwargs); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(kwargs); } void export_function( @@ -943,8 +973,9 @@ void export_function( const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args, kwargs); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(args, kwargs); } FunctionExporter exporter( @@ -1054,6 +1085,7 @@ ImportedFunction::ImportedFunction(const std::string& file) auto mlx_version = deserialize(is); auto function_count = deserialize(is); ftable->shapeless = deserialize(is); + metadata_ = deserialize(is); std::unordered_map constants; auto import_one = [&]() { diff --git a/mlx/export.h b/mlx/export.h index 5532f7c818..8135919b26 100644 --- a/mlx/export.h +++ b/mlx/export.h @@ -50,17 +50,20 @@ struct FunctionExporter; MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API FunctionExporter exporter( const std::string& path, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); /** * Export a function to a file. @@ -69,20 +72,23 @@ MLX_API void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API void export_function( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); struct ImportedFunction; diff --git a/mlx/export_impl.h b/mlx/export_impl.h index 467a5f0d6c..7fd053a900 100644 --- a/mlx/export_impl.h +++ b/mlx/export_impl.h @@ -27,17 +27,20 @@ struct MLX_API FunctionExporter { friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&, const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const ExportCallback&, @@ -57,7 +60,8 @@ struct MLX_API FunctionExporter { FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless); + bool shapeless, + std::string metadata); FunctionExporter( const ExportCallback& callback, @@ -77,6 +81,7 @@ struct MLX_API FunctionExporter { int count{0}; bool closed{false}; std::shared_ptr ftable; + std::string metadata_; }; struct MLX_API ImportedFunction { @@ -88,12 +93,18 @@ struct MLX_API ImportedFunction { std::vector operator()(const Kwargs& kwargs) const; std::vector operator()(const Args& args, const Kwargs& kwargs) const; + // The metadata saved with the function when it was exported. + const std::string& metadata() const { + return metadata_; + } + private: ImportedFunction(const std::string& file); friend MLX_API ImportedFunction import_function(const std::string&); ImportedFunction(); std::shared_ptr ftable; + std::string metadata_; }; } // namespace mlx::core diff --git a/python/src/export.cpp b/python/src/export.cpp index b4c2998a03..30d48bcae4 100644 --- a/python/src/export.cpp +++ b/python/src/export.cpp @@ -138,6 +138,7 @@ void init_export(nb::module_& m) { const nb::callable& fun, const nb::args& args, bool shapeless, + const std::optional& metadata, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs(args, kwargs, "[export_function]"); @@ -147,8 +148,14 @@ void init_export(nb::module_& m) { wrap_export_function(fun), args_, kwargs_, - shapeless); + shapeless, + metadata.value_or("")); } else { + if (metadata && !metadata->empty()) { + throw std::invalid_argument( + "[export_function] The metadata argument is only supported " + "when exporting to a file, not when using a callback."); + } auto callback = nb::cast(file_or_callback); auto wrapped_callback = [callback](const mx::ExportCallbackInput& input) { @@ -163,9 +170,10 @@ void init_export(nb::module_& m) { "args"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), "kwargs"_a, nb::sig( - "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, **kwargs) -> None"), + "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, metadata: Optional[str] = None, **kwargs) -> None"), R"pbdoc( Export an MLX function. @@ -187,9 +195,16 @@ void init_export(nb::module_& m) { *args (array): Example array inputs to the function. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (str, optional): A string to save alongside the + function, for example a JSON encoded model configuration. Only + supported when exporting to a file. Read it back with + :func:`import_function`. Default: ``None``. **kwargs (array): Additional example keyword array inputs to the function. + Raises: + ValueError: If ``metadata`` is given when exporting with a callback. + Example: .. code-block:: python @@ -203,17 +218,25 @@ void init_export(nb::module_& m) { )pbdoc"); m.def( "import_function", - [](const std::string& file) { - return nb::cpp_function( - [fn = mx::import_function(file)]( + [](const std::string& file, bool return_metadata) -> nb::object { + auto imported = mx::import_function(file); + auto metadata = imported.metadata(); + auto fn = nb::cpp_function( + [imported = std::move(imported)]( const nb::args& args, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs( args, kwargs, "[import_function::call]"); - return nb::tuple(nb::cast(fn(args_, kwargs_))); + return nb::tuple(nb::cast(imported(args_, kwargs_))); }); + if (return_metadata) { + return nb::make_tuple(fn, nb::cast(metadata)); + } + return fn; }, "file"_a, - nb::sig("def import_function(file: str) -> Callable"), + "return_metadata"_a = false, + nb::sig( + "def import_function(file: str, return_metadata: bool = False) -> Union[Callable, tuple[Callable, str]]"), R"pbdoc( Import a function from a file. @@ -230,15 +253,20 @@ void init_export(nb::module_& m) { Args: file (str): The file path to import the function from. + return_metadata (bool, optional): If ``True`` also return the + metadata string saved with the function. Default: ``False``. Returns: - Callable: The imported function. + Callable or tuple: + The imported function. If ``return_metadata`` is ``True`` a + tuple of the imported function and the metadata string is + returned instead. Example: >>> fn = mx.import_function("function.mlxfn") >>> out = fn(a, b, x=x, y=y)[0] >>> - >>> out = fn((a, b), {"x": x, "y": y}[0] + >>> out = fn((a, b), {"x": x, "y": y})[0] )pbdoc"); nb::class_( @@ -274,14 +302,23 @@ void init_export(nb::module_& m) { m.def( "exporter", - [](const std::string& file, nb::callable fun, bool shapeless) { + [](const std::string& file, + nb::callable fun, + bool shapeless, + const std::optional& metadata) { return PyFunctionExporter{ - mx::exporter(file, wrap_export_function(fun), shapeless), fun}; + mx::exporter( + file, + wrap_export_function(fun), + shapeless, + metadata.value_or("")), + fun}; }, "file"_a, "fun"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), R"pbdoc( Make a callable object to export multiple traces of a function to a file. @@ -295,6 +332,9 @@ void init_export(nb::module_& m) { file (str): File path to export the function to. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (str, optional): A string to save alongside the + function, for example a JSON encoded model configuration. Read + it back with :func:`import_function`. Default: ``None``. Example: diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index 6b5f0ab146..15d3b8c4c9 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -1,6 +1,7 @@ # Copyright © 2024 Apple Inc. import gc +import json import os import tempfile import unittest @@ -748,6 +749,42 @@ def fn(x): self.assertEqual(y.shape, (B, seq_len, H)) self.assertTrue(mx.allclose(y, expected)) + def test_export_import_metadata(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + def fun(x): + return mx.abs(x) + + x = mx.array([1.0, -2.0, 3.0]) + metadata = json.dumps({"name": "model", "params": 7_000_000_000, "lr": 0.1}) + + mx.export_function(path, fun, x, metadata=metadata) + + imported = mx.import_function(path) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + imported, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + self.assertEqual(json.loads(imported_metadata)["params"], 7_000_000_000) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + mx.export_function(path, fun, x) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, "") + + # Metadata survives the per-trace header rewrite of a multi-trace export + with mx.exporter(path, fun, metadata=metadata) as exporter: + exporter(mx.array([1.0])) + exporter(mx.array([1.0, 2.0])) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + + with self.assertRaises(TypeError): + mx.export_function(path, fun, x, metadata={"name": "model"}) + + with self.assertRaises(ValueError): + mx.export_function(lambda x: None, fun, x, metadata=metadata) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/tests/export_import_tests.cpp b/tests/export_import_tests.cpp index ef6a18e199..c601c40b80 100644 --- a/tests/export_import_tests.cpp +++ b/tests/export_import_tests.cpp @@ -161,3 +161,31 @@ TEST_CASE("test export function on different stream") { // Should make a new stream that we can run computation on eval(import_function(file_path)({array({0, 1, 2})})); } + +TEST_CASE("test export import with metadata") { + std::string file_path = get_temp_file("model.mlxfn"); + + auto fun = [](const std::vector& args) -> std::vector { + return {abs(args[0])}; + }; + + std::string metadata = "{\"name\": \"model\", \"params\": 7000000000}"; + + export_function(file_path, fun, {array({0, 1, 2})}, false, metadata); + + auto imported = import_function(file_path); + CHECK(imported.metadata() == metadata); + eval(imported({array({0, 1, 2})})); + + // No metadata gives an empty string + export_function(file_path, fun, {array({0, 1, 2})}); + CHECK(import_function(file_path).metadata().empty()); + + // Metadata survives the per-trace header rewrite of a multi-trace export + { + auto fn_exporter = exporter(file_path, fun, false, metadata); + fn_exporter({array({0, 1})}); + fn_exporter({array({0, 1, 2})}); + } + CHECK(import_function(file_path).metadata() == metadata); +} From ed116d24ba016c200917ea0d4779d68c8859c4fe Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+prady0t@users.noreply.github.com> Date: Sat, 8 Aug 2026 06:32:58 +0530 Subject: [PATCH 089/222] Fix empty matrix case in cholesky (#4033) Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- mlx/backend/cpu/cholesky.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx/backend/cpu/cholesky.cpp b/mlx/backend/cpu/cholesky.cpp index 3c5bbbc93d..28cb87dd66 100644 --- a/mlx/backend/cpu/cholesky.cpp +++ b/mlx/backend/cpu/cholesky.cpp @@ -26,6 +26,9 @@ void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) { a.flags().row_contiguous ? CopyType::Vector : CopyType::General, stream); + if (a.shape(-1) == 0) { + return; + } auto& encoder = cpu::get_command_encoder(stream); encoder.set_output_array(factor); encoder.dispatch([matrix = factor.data(), From 2b64179d46d6749e2d00388618b763e7790e586a Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 7 Aug 2026 20:13:15 -0700 Subject: [PATCH 090/222] chore: Reject a negative max_norm in clip_grad_norm (#4058) --- python/mlx/optimizers/optimizers.py | 3 +++ python/tests/test_optimizers.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index 3f167fd36b..65efab222d 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/python/mlx/optimizers/optimizers.py @@ -968,6 +968,9 @@ def clip_grad_norm(grads, max_norm): (dict, float): The possibly rescaled gradients and the original gradient norm. """ + if max_norm < 0: + raise ValueError(f"max_norm should be >=0, {max_norm} was provided instead") + norm_squared = tree_reduce(lambda acc, g: acc + g.square().sum(), grads, 0.0) total_norm = mx.sqrt(norm_squared) normalizer = mx.minimum(max_norm / (total_norm + 1e-6), 1.0) diff --git a/python/tests/test_optimizers.py b/python/tests/test_optimizers.py index 68a88d1820..fec3d70a21 100644 --- a/python/tests/test_optimizers.py +++ b/python/tests/test_optimizers.py @@ -562,6 +562,9 @@ def test_clip_grad_norm(self): "Gradients were not scaled correctly during clipping.", ) + with self.assertRaises(ValueError): + opt.clip_grad_norm(small_grads, -1.0) + def test_init_from_state(self): class Model(nn.Module): def __init__(self): From 5139a8643df334ab2f51a8b8d376d34c651924ef Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 7 Aug 2026 20:13:30 -0700 Subject: [PATCH 091/222] Build extensions with the interpreter running the build (#4057) --- python/mlx/extension.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/mlx/extension.py b/python/mlx/extension.py index c426d59531..d794fbe28f 100644 --- a/python/mlx/extension.py +++ b/python/mlx/extension.py @@ -30,13 +30,14 @@ def build_extension(self, ext: CMakeExtension) -> None: debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug cfg = "Debug" if debug else "Release" - # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON - # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code - # from Python. + # Point CMake at the interpreter running the build. Otherwise + # find_package(Python) picks whichever interpreter it finds first, + # which is not the one nanobind and mlx are installed into. cmake_args = [ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}", f"-DCMAKE_BUILD_TYPE={cfg}", "-DBUILD_SHARED_LIBS=ON", + f"-DPython_EXECUTABLE={sys.executable}", ] build_args = [] # Adding CMake arguments set as environment variable From 11fa2c839723386361aade3ac33365200c137d28 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 7 Aug 2026 20:13:56 -0700 Subject: [PATCH 092/222] chore: Reject negative dimensions in broadcast_to and random shapes (#4046) --- mlx/ops.cpp | 9 +++++++++ mlx/random.cpp | 5 +++++ python/tests/test_ops.py | 6 ++++++ python/tests/test_random.py | 3 +++ 4 files changed, 23 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index a7034d7595..b7b9a84291 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -304,6 +304,10 @@ array as_strided( Strides strides, size_t offset, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument( + "[as_strided] Negative dimensions not allowed."); + } auto copied_shape = shape; // |shape| will be moved auto dtype = a.dtype(); // |a| will be moved return array( @@ -1692,6 +1696,11 @@ array broadcast_to( const array& a, const Shape& shape, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument( + "[broadcast_to] Negative dimensions not allowed."); + } + if (a.shape() == shape) { return a; } diff --git a/mlx/random.cpp b/mlx/random.cpp index def3169cb5..4885382c7e 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include #include #include @@ -38,6 +39,10 @@ array bits( int width /* 4 */, const std::optional& key_ /*= nullopt */, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument("[bits] Negative dimensions not allowed."); + } + auto key = key_ ? *key_ : KeySequence::default_().next(); if (key.dtype() != uint32) { std::ostringstream msg; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index f093891a15..e8555b8747 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -810,6 +810,9 @@ def test_broadcast(self): self.assertListEqual(list(b_npy.shape), list(b_mlx.shape)) self.assertTrue(np.array_equal(b_npy, b_mlx)) + with self.assertRaises(ValueError): + mx.broadcast_to(a_mlx, (-1, 10, 20)) + def test_logsumexp(self): def logsumexp(x, axes=None): maxs = mx.max(x, axis=axes, keepdims=True) @@ -2224,6 +2227,9 @@ def test_as_strided(self): y = mx.as_strided(x, (x.size,), (-1,), x.size - 1) self.assertTrue(mx.array_equal(y, x[::-1])) + with self.assertRaises(ValueError): + mx.as_strided(x, (-2, 3), (3, 1), 0) + def test_logcumsumexp(self): npop = np.logaddexp.accumulate mxop = mx.logcumsumexp diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 551c32993c..8cb1cbb43b 100644 --- a/python/tests/test_random.py +++ b/python/tests/test_random.py @@ -64,6 +64,9 @@ def test_uniform(self): self.assertEqual(mx.random.uniform().dtype, mx.random.uniform(dtype=None).dtype) + with self.assertRaises(ValueError): + mx.random.uniform(shape=(2, -3)) + def test_normal_and_laplace(self): # Same tests for normal and laplace. for distribution_sampler in [mx.random.normal, mx.random.laplace]: From f3d0c1ad4b17d7fe38c50cd67410d7ffde1e4830 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 7 Aug 2026 20:14:10 -0700 Subject: [PATCH 093/222] chore: Report bad host arguments as CLI errors instead of tracebacks (#4045) --- python/mlx/_distributed_utils/config.py | 11 +++++++---- python/mlx/_distributed_utils/launch.py | 11 +++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/python/mlx/_distributed_utils/config.py b/python/mlx/_distributed_utils/config.py index 0ab31470b0..a4dabed629 100644 --- a/python/mlx/_distributed_utils/config.py +++ b/python/mlx/_distributed_utils/config.py @@ -615,10 +615,13 @@ def main(): ) args = parser.parse_args() - if args.hostfile is not None: - hosts = Hostfile.from_file(args.hostfile).hosts - else: - hosts = Hostfile.from_list(args.hosts).hosts + try: + if args.hostfile is not None: + hosts = Hostfile.from_file(args.hostfile).hosts + else: + hosts = Hostfile.from_list(args.hosts).hosts + except ValueError as e: + parser.error(str(e)) # Check that we can ssh log( diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index b052643aaf..464bff7e70 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -537,10 +537,13 @@ def main(): rest.pop(0) # Try to extract a list of hosts and corresponding ips - if args.hostfile is not None: - hostfile = Hostfile.from_file(args.hostfile) - else: - hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts) + try: + if args.hostfile is not None: + hostfile = Hostfile.from_file(args.hostfile) + else: + hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts) + except ValueError as e: + parser.error(str(e)) # Extract extra arguments from the hostfile if hostfile.backend != "" and args.backend is None: From 47bbfe8fa473d6d19037a8d97f1f7d30514e4cf6 Mon Sep 17 00:00:00 2001 From: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:14:47 -0700 Subject: [PATCH 094/222] Fix Metal FFT for sizes above 2**20 (#4013) --- mlx/backend/metal/fft.cpp | 27 ++++++++++++++++++++++----- python/tests/test_fft.py | 15 +++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/fft.cpp b/mlx/backend/metal/fft.cpp index 61eb02dac9..6c7a931390 100644 --- a/mlx/backend/metal/fft.cpp +++ b/mlx/backend/metal/fft.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "mlx/3rdparty/pocketfft.h" #include "mlx/backend/common/utils.h" @@ -127,6 +128,21 @@ FFTPlan plan_fft(int n) { // Rough heuristic for choosing faster powers of two when we can plan.n2 = n > 65536 ? 1024 : 64; plan.n1 = n / plan.n2; + // Each step is a single threadgroup FFT, so neither factor can be larger + // than the largest Stockham size. The no transpose kernels cannot + // decompose a step any further, so grow n2 instead of nesting four step. + if (plan.n1 > MAX_STOCKHAM_FFT_SIZE) { + plan.n1 = MAX_STOCKHAM_FFT_SIZE; + plan.n2 = n / plan.n1; + } + if (plan.n2 > MAX_STOCKHAM_FFT_SIZE) { + std::ostringstream msg; + msg << "[FFT] GPU FFT is not supported for size " << n + << ", the largest supported size is " + << MAX_STOCKHAM_FFT_SIZE * MAX_STOCKHAM_FFT_SIZE + << ". Run the FFT on the CPU stream instead."; + throw std::runtime_error(msg.str()); + } return plan; } else if (n > MAX_STOCKHAM_FFT_SIZE) { // Otherwise we use a multi-upload Bluestein's @@ -638,13 +654,14 @@ void fft_op( // We batch among threadgroups for improved efficiency when n is small int threadgroup_batch_size = std::max(MIN_THREADGROUP_MEM_SIZE / fft_size, 1); if (four_step_params.required) { - // Require a threadgroup batch size of at least 4 for four step FFT - // so we can coalesce the memory accesses. - threadgroup_batch_size = - std::max(threadgroup_batch_size, MIN_COALESCE_WIDTH); + // Batch the four step FFT so we can coalesce the memory accesses, but + // never past what fits in threadgroup memory. + threadgroup_batch_size = std::max( + threadgroup_batch_size, + std::min(MIN_COALESCE_WIDTH, MAX_STOCKHAM_FFT_SIZE / fft_size)); } int threadgroup_mem_size = next_power_of_2(threadgroup_batch_size * fft_size); - // FFTs up to 2^20 are currently supported + // The plan keeps every step within the threadgroup memory limit assert(threadgroup_mem_size <= MAX_STOCKHAM_FFT_SIZE); // ceil divide diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 0bb190bd99..26f473f2ca 100644 --- a/python/tests/test_fft.py +++ b/python/tests/test_fft.py @@ -185,6 +185,21 @@ def test_fft_big_powers_of_two(self): for k in range(17, 20): self._run_ffts((3, 2**k), atol=1e-2) + # Past 2**20 the four step plan has to grow n2 to keep both factors + # inside threadgroup memory + for k in range(20, 25): + self._run_ffts((1, 2**k), atol=1e-2, rtol=1e-3) + + @unittest.skipIf( + not mx.metal.is_available(), "the size limit is specific to the Metal FFT plan" + ) + def test_fft_too_large(self): + # Larger than the four step plan can decompose, so it has to throw + # rather than run a kernel that silently returns the wrong answer. + # CUDA hands this to cuFFT instead and has no such limit. + with self.assertRaises(RuntimeError): + mx.eval(mx.fft.fft(mx.zeros(2**25))) + def test_fft_large_numbers(self): numbers = [ 1037, # prime > 2048 From e190d04592f2d1c71bb0369213340c6297e2f2a8 Mon Sep 17 00:00:00 2001 From: Magnus Lundstedt Date: Sat, 8 Aug 2026 07:57:14 +0200 Subject: [PATCH 095/222] Request MSL 4.1 from the runtime compiler on macOS 27 (#4052) Co-authored-by: Cheng --- mlx/backend/metal/device.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 197b919b1d..a7448dfea8 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -64,7 +64,10 @@ void set_compile_options( auto get_metal_version() { auto get_metal_version_ = []() { - if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) { + if (__builtin_available(macOS 27, iOS 27, tvOS 27, visionOS 27, *)) { + // TODO: Use MTL::LanguageVersion4_1 after metal-cpp_27 is released. + return static_cast((4 << 16) + 1); + } else if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) { return MTL::LanguageVersion4_0; } else if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) { return MTL::LanguageVersion3_2; From f599c0209f606ed7faad5c78a6c761e222da4e63 Mon Sep 17 00:00:00 2001 From: kitty <914384274@qq.com> Date: Sat, 8 Aug 2026 13:57:39 +0800 Subject: [PATCH 096/222] Fix concurrent Metal kernel cache lookup (#4043) Co-authored-by: hezz Co-authored-by: Cheng --- mlx/backend/metal/device.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index a7448dfea8..cd601dd319 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -886,9 +886,12 @@ MTL::ComputePipelineState* Device::get_kernel( std::shared_lock lock(kernel_mtx_); // Look for cached kernel - auto& kernel_map_ = library_kernels_[mtl_lib]; - if (auto it = kernel_map_.find(kname); it != kernel_map_.end()) { - return it->second.get(); + auto library_it = library_kernels_.find(mtl_lib); + if (library_it != library_kernels_.end()) { + auto kernel_it = library_it->second.find(kname); + if (kernel_it != library_it->second.end()) { + return kernel_it->second.get(); + } } } return get_kernel_(base_name, mtl_lib, kname, func_consts, linked_functions); From 2ab4d27cd760bf2429a46ab3c3da6971157e3559 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 00:10:26 -0700 Subject: [PATCH 097/222] chore: Use dispatch_all_types in ArgReduce::eval_cpu (#4066) --- mlx/backend/cpu/arg_reduce.cpp | 49 ++++------------------------------ 1 file changed, 5 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/arg_reduce.cpp b/mlx/backend/cpu/arg_reduce.cpp index 66468912d1..d93f74cf0a 100644 --- a/mlx/backend/cpu/arg_reduce.cpp +++ b/mlx/backend/cpu/arg_reduce.cpp @@ -4,6 +4,7 @@ #include "mlx/backend/common/utils.h" #include "mlx/backend/cpu/encoder.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -74,50 +75,10 @@ void ArgReduce::eval_cpu(const std::vector& inputs, array& out) { out = array::unsafe_weak_copy(out), reduce_type_ = reduce_type_, axis_ = axis_]() mutable { - switch (in.dtype()) { - case bool_: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint8: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int8: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case bfloat16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case complex64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + arg_reduce_dispatch(in, out, reduce_type_, axis_); + }); }); } From 361883c37897982834b31973c3226b801418dfa2 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 00:11:11 -0700 Subject: [PATCH 098/222] chore: Use dispatch_all_types in NumberOfElements::eval (#4064) --- mlx/backend/common/common.cpp | 49 ++++------------------------------- 1 file changed, 5 insertions(+), 44 deletions(-) diff --git a/mlx/backend/common/common.cpp b/mlx/backend/common/common.cpp index cbc90ed27e..5426301a35 100644 --- a/mlx/backend/common/common.cpp +++ b/mlx/backend/common/common.cpp @@ -3,6 +3,7 @@ #include "mlx/backend/common/broadcasting.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -99,50 +100,10 @@ void NumberOfElements::eval(const std::vector& inputs, array& out) { numel = 1.0 / numel; } - switch (out.dtype()) { - case bool_: - *out.data() = static_cast(numel); - break; - case uint8: - *out.data() = static_cast(numel); - break; - case uint16: - *out.data() = static_cast(numel); - break; - case uint32: - *out.data() = static_cast(numel); - break; - case uint64: - *out.data() = static_cast(numel); - break; - case int8: - *out.data() = static_cast(numel); - break; - case int16: - *out.data() = static_cast(numel); - break; - case int32: - *out.data() = static_cast(numel); - break; - case int64: - *out.data() = static_cast(numel); - break; - case float16: - *out.data() = static_cast(numel); - break; - case float32: - *out.data() = static_cast(numel); - break; - case bfloat16: - *out.data() = static_cast(numel); - break; - case float64: - *out.data() = static_cast(numel); - break; - case complex64: - *out.data() = static_cast(numel); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + *out.data() = static_cast(numel); + }); } std::pair prepare_reshape(const array& in, const array& out) { From 456216e38d2fd730a218ff9a1a2aa901c8320d68 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 00:12:06 -0700 Subject: [PATCH 099/222] chore: Use dispatch_all_types in CPU select_op (#4065) --- mlx/backend/cpu/select.cpp | 52 ++++---------------------------------- 1 file changed, 5 insertions(+), 47 deletions(-) diff --git a/mlx/backend/cpu/select.cpp b/mlx/backend/cpu/select.cpp index bf6a9b8259..89164f9ddb 100644 --- a/mlx/backend/cpu/select.cpp +++ b/mlx/backend/cpu/select.cpp @@ -4,6 +4,7 @@ #include "mlx/backend/cpu/binary_ops.h" #include "mlx/backend/cpu/ternary.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -32,53 +33,10 @@ void select_op( out = array::unsafe_weak_copy(out), op, topt]() mutable { - switch (out.dtype()) { - case bool_: - ternary_op(a, b, c, out, op, topt); - break; - case uint8: - ternary_op(a, b, c, out, op, topt); - break; - case uint16: - ternary_op(a, b, c, out, op, topt); - break; - case uint32: - ternary_op(a, b, c, out, op, topt); - break; - case uint64: - ternary_op(a, b, c, out, op, topt); - break; - case int8: - ternary_op(a, b, c, out, op, topt); - break; - case int16: - ternary_op(a, b, c, out, op, topt); - break; - case int32: - ternary_op(a, b, c, out, op, topt); - break; - case int64: - ternary_op(a, b, c, out, op, topt); - break; - case float16: - ternary_op( - a, b, c, out, op, topt); - break; - case float32: - ternary_op(a, b, c, out, op, topt); - break; - case float64: - ternary_op(a, b, c, out, op, topt); - break; - case bfloat16: - ternary_op( - a, b, c, out, op, topt); - break; - case complex64: - ternary_op( - a, b, c, out, op, topt); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + ternary_op(a, b, c, out, op, topt); + }); }); } From e7838d5e386d1159192eb959bee266dd65d27bb3 Mon Sep 17 00:00:00 2001 From: pierre427 Date: Sat, 8 Aug 2026 03:12:45 -0400 Subject: [PATCH 100/222] Raise qmv batch limit for large matrices on M5-class GPUs (#3791) Co-authored-by: Pierre Lamy Co-authored-by: Cheng --- mlx/backend/metal/quantized.cpp | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index c9f6a7e00e..78d8cf0df5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -85,7 +85,23 @@ inline array ensure_row_contiguous_matrix( inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { auto arch_size = d.get_architecture().back(); auto arch_gen = d.get_architecture_gen(); - if (arch_gen == 13 || arch_gen == 14) { + if (arch_gen >= 17 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 33; + } else if (D <= 4096 && O <= 4096) { + return 25; + } else { + return 13; + } + } else if (arch_gen >= 15 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 13; + } else if (D <= 4096 && O <= 4096) { + return 15; + } else { + return 13; + } + } else if (arch_gen >= 13) { switch (arch_size) { case 'd': if (D <= 2048 && O <= 2048) { From a88d454d3f8b591c925bf66dd1095b0f015f40a3 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:32:05 -0700 Subject: [PATCH 101/222] Enable half-precision complex Metal FFT kernels (#3981) Co-authored-by: Cheng --- mlx/backend/metal/kernels/complex.h | 25 +++++++- mlx/backend/metal/kernels/fft.h | 31 +++++++--- mlx/backend/metal/kernels/fft.metal | 16 +++-- mlx/backend/metal/kernels/fft/radix.h | 9 +-- mlx/backend/metal/kernels/fft/readwrite.h | 72 +++++++++++++++++++---- 5 files changed, 124 insertions(+), 29 deletions(-) diff --git a/mlx/backend/metal/kernels/complex.h b/mlx/backend/metal/kernels/complex.h index 06a3050e8a..fe78e0050f 100644 --- a/mlx/backend/metal/kernels/complex.h +++ b/mlx/backend/metal/kernels/complex.h @@ -95,6 +95,28 @@ struct complex_t { constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), imag(static_cast(x.imag)) {} + // Conversions to and from two-lane vectors (the FFT lane representation) + constexpr complex_t(vec v) thread : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) threadgroup : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) device : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) constant : real(v.x), imag(v.y) {}; + + constexpr operator vec() const thread { + return vec(real, imag); + } + + constexpr operator vec() const threadgroup { + return vec(real, imag); + } + + constexpr operator vec() const device { + return vec(real, imag); + } + + constexpr operator vec() const constant { + return vec(real, imag); + } + // Conversions to scalar types template < typename U, @@ -129,10 +151,11 @@ struct complex_t { } }; +using complex32_t = complex_t; using complex64_t = complex_t; +static_assert(sizeof(complex32_t) == 2 * sizeof(half)); static_assert(sizeof(complex64_t) == 2 * sizeof(float)); -static_assert(sizeof(complex_t) == 2 * sizeof(half)); static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); template diff --git a/mlx/backend/metal/kernels/fft.h b/mlx/backend/metal/kernels/fft.h index 53b48ca7aa..408fc1a6d2 100644 --- a/mlx/backend/metal/kernels/fft.h +++ b/mlx/backend/metal/kernels/fft.h @@ -325,7 +325,18 @@ template int x_sum_index = metal::min(fft_idx, rader_m - 1); buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)]; - vec inv = {1.0f, -1.0f}; + // The convolution is conjugated and scaled by 1/(rader_n - 1) exactly once + // each. Float lanes apply the scale after the inverse FFT, preserving the + // existing float code generation; reduced lanes apply it before, so the + // intermediates cannot overflow the narrow lane range. + scalar_T rader_inv_r = static_cast(1.0f / (rader_n - 1)); + vec convolution_factor = {1.0f, -1.0f}; + vec convolution_inv_factor = {rader_inv_r, -rader_inv_r}; + if constexpr (!metal::is_same_v) { + convolution_factor = convolution_inv_factor; + convolution_inv_factor = {1.0f, -1.0f}; + } + for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); short interleaved_index = @@ -339,7 +350,7 @@ template for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); - buf[rader_m + index] = temp[e] * inv; + buf[rader_m + index] = temp[e] * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -349,13 +360,10 @@ template perform_fft( fft_idx, &p, m, n - rader_m, buf + rader_m); - scalar_T rader_inv_r = static_cast(1.0f / (rader_n - 1)); - vec rader_inv_factor = {rader_inv_r, -rader_inv_r}; - for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, n - rader_m - 1); short diff_index = index / (rader_n - 1) - x_0_index; - temp[e] = buf[rader_m + index] * rader_inv_factor + x_0[diff_index]; + temp[e] = buf[rader_m + index] * convolution_inv_factor + x_0[diff_index]; } // Use the sum of elements that was computed in the first FFT @@ -433,10 +441,17 @@ template // fft perform_fft(fft_idx, &p, m, n, buf); - vec inv = {1.0f, -1.0f}; + vec convolution_factor = {1.0f, -1.0f}; + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize the convolution before the inverse FFT so the + // intermediates cannot overflow; the write path then only conjugates. + scalar_T inv_n = static_cast(1.0f / n); + convolution_factor = {inv_n, -inv_n}; + } for (int t = 0; t < elems_per_thread_; t++) { int index = fft_idx + t * m; - buf[index] = complex_mul(buf[index], w_q[index]) * inv; + buf[index] = + complex_mul(buf[index], w_q[index]) * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); diff --git a/mlx/backend/metal/kernels/fft.metal b/mlx/backend/metal/kernels/fft.metal index 590b558efc..4d457928a7 100644 --- a/mlx/backend/metal/kernels/fft.metal +++ b/mlx/backend/metal/kernels/fft.metal @@ -38,18 +38,22 @@ real) // clang-format off -#define instantiate_ffts(tg_mem_size) \ - instantiate_fft(tg_mem_size, float2, float2) \ +#define instantiate_c2c_ffts(tg_mem_size, complex_T) \ + instantiate_fft(tg_mem_size, complex_T, complex_T) \ + instantiate_rader(tg_mem_size, complex_T, complex_T) \ + instantiate_bluestein(tg_mem_size, complex_T, complex_T) \ + instantiate_four_step(tg_mem_size, complex_T, complex_T, 0, /*real=*/false) \ + instantiate_four_step(tg_mem_size, complex_T, complex_T, 1, /*real=*/false) + +#define instantiate_ffts(tg_mem_size) \ + instantiate_c2c_ffts(tg_mem_size, float2) \ + instantiate_c2c_ffts(tg_mem_size, complex32_t) \ instantiate_fft(tg_mem_size, float, float2) \ instantiate_fft(tg_mem_size, float2, float) \ - instantiate_rader(tg_mem_size, float2, float2) \ instantiate_rader(tg_mem_size, float, float2) \ instantiate_rader(tg_mem_size, float2, float) \ - instantiate_bluestein(tg_mem_size, float2, float2) \ instantiate_bluestein(tg_mem_size, float, float2) \ instantiate_bluestein(tg_mem_size, float2, float) \ - instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/false) \ - instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/false) \ instantiate_four_step(tg_mem_size, float, float2, 0, /*real=*/true) \ instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/true) \ instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/true) \ diff --git a/mlx/backend/metal/kernels/fft/radix.h b/mlx/backend/metal/kernels/fft/radix.h index c209a4bc37..9d742eb33c 100644 --- a/mlx/backend/metal/kernels/fft/radix.h +++ b/mlx/backend/metal/kernels/fft/radix.h @@ -38,10 +38,11 @@ METAL_FUNC metal::vec complex_mul_conj( template METAL_FUNC metal::vec get_twiddle(int k, int p) { // Derive phase evaluation precision from the scalar lane and Metal's pi - // constant. Reduced lanes currently promote to float for fast trig. + // constant. Reduced lanes currently promote to float for fast trig, then + // narrow the result to the lane type explicitly. using phase_T = decltype(M_PI_F * T(0)); phase_T theta = -phase_T(2) * phase_T(k) * phase_T(M_PI_F) / phase_T(p); - return {metal::fast::cos(theta), metal::fast::sin(theta)}; + return {T(metal::fast::cos(theta)), T(metal::fast::sin(theta))}; } template @@ -193,8 +194,8 @@ METAL_FUNC void radix10( thread metal::vec* x, thread metal::vec* y) { metal::vec w[4]; - w[0] = {0.8090169943749475, -0.5877852522924731}; - w[1] = {0.30901699437494745, -0.9510565162951535}; + w[0] = {T(0.8090169943749475), T(-0.5877852522924731)}; + w[1] = {T(0.30901699437494745), T(-0.9510565162951535)}; w[2] = {-w[1].x, w[1].y}; w[3] = {-w[0].x, w[0].y}; diff --git a/mlx/backend/metal/kernels/fft/readwrite.h b/mlx/backend/metal/kernels/fft/readwrite.h index 4e954e29fb..86ba7df1fd 100644 --- a/mlx/backend/metal/kernels/fft/readwrite.h +++ b/mlx/backend/metal/kernels/fft/readwrite.h @@ -2,6 +2,7 @@ #include +#include "mlx/backend/metal/kernels/complex.h" #include "mlx/backend/metal/kernels/fft/radix.h" /* FFT helpers for reading and writing from/to device memory. @@ -38,6 +39,11 @@ struct FFTStorageTraits> { using scalar_T = T; }; +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + template struct FFTIOTypeTraits { using scalar_T = typename FFTStorageTraits::scalar_T; @@ -106,17 +112,31 @@ struct ReadWriter { return inv ? complex_T(elem.x, -elem.y) : elem; } + // Handle packed complex_t storage + METAL_FUNC complex_T post_in(complex_t elem) const thread { + return post_in(complex_T(elem.real, elem.imag)); + } + // Handle float case for generic RFFT alg METAL_FUNC complex_T post_in(scalar_T elem) const thread { return complex_T(elem, 0); } METAL_FUNC complex_T pre_out(complex_T elem) const thread { - return inv ? complex_T(elem.x / n, -elem.y / n) : elem; + return pre_out(elem, n); } METAL_FUNC complex_T pre_out(complex_T elem, int length) const thread { - return inv ? complex_T(elem.x / length, -elem.y / length) : elem; + if (!inv) { + return elem; + } + if constexpr (metal::is_same_v) { + return complex_T(elem.x / length, -elem.y / length); + } else { + // Compute the reciprocal in float before narrowing to the lane. + scalar_T inv_length = static_cast(1.0f / length); + return complex_T(elem.x * inv_length, -elem.y * inv_length); + } } METAL_FUNC bool out_of_bounds() const thread { @@ -125,7 +145,7 @@ struct ReadWriter { return grid_index >= batch_size; } - METAL_FUNC void load() const thread { + METAL_FUNC void load(scalar_T scale) const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; // Keep each thread's sequential access at 128 bits where possible. @@ -137,7 +157,7 @@ struct ReadWriter { index = metal::min(index, max_full_index); // vectorized reads for (short r = 0; r < read_width; r++) { - buf[index + r] = post_in(in[batch_idx + index + r]); + buf[index + r] = post_in(in[batch_idx + index + r]) * scale; } } short max_index = grid.y * n - 1; @@ -145,10 +165,14 @@ struct ReadWriter { short index = tg_idx + r * threads_per_tg + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); - buf[index] = post_in(in[batch_idx + index]); + buf[index] = post_in(in[batch_idx + index]) * scale; } } + METAL_FUNC void load() const thread { + load(static_cast(1)); + } + METAL_FUNC void write() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; @@ -198,6 +222,11 @@ struct ReadWriter { int m = grid.z; scalar_T inv_n = static_cast(1.0f / n); complex_T inv_factor = {inv_n, -inv_n}; + if constexpr (!metal::is_same_v) { + // Reduced lanes applied the 1/n scale with the convolution, so only + // the conjugation remains here. + inv_factor = {scalar_T(1), -scalar_T(1)}; + } threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { @@ -240,13 +269,28 @@ struct ReadWriter { (void)overall_n; bool default_inv = inv; inv = false; - load(); + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize inverse transforms on input so + // intermediates stay in range; see write_strided. + if (default_inv) { + load(static_cast(1.0f / n)); + } else { + load(); + } + } else { + load(); + } inv = default_inv; } else { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { - buf[strided_shared_idx + e] = - post_in(in[strided_device_idx + e * stride]); + complex_T input = post_in(in[strided_device_idx + e * stride]); + if constexpr (!metal::is_same_v) { + if (step == 0 && !four_step_real && inv) { + input *= static_cast(1.0f / n); + } + } + buf[strided_shared_idx + e] = input; } } } @@ -255,8 +299,16 @@ struct ReadWriter { if constexpr (step == 1 && !four_step_real) { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { - out[strided_device_idx + e * stride] = - pre_out(buf[strided_shared_idx + e], overall_n); + if constexpr (!metal::is_same_v) { + // Inverse C2C passes were normalized on input, so only the + // conjugation of the inverse transform remains here. + complex_T value = buf[strided_shared_idx + e]; + out[strided_device_idx + e * stride] = + inv ? complex_T(value.x, -value.y) : value; + } else { + out[strided_device_idx + e * stride] = + pre_out(buf[strided_shared_idx + e], overall_n); + } } } else { for (int e = 0; e < elems_per_thread; e++) { From 9700e9d33ab4c0ce00f324a178c3bf422f6b8875 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Sat, 8 Aug 2026 17:56:23 +0800 Subject: [PATCH 102/222] Fix mx.distributed.sum_scatter crashes on a scalar (#4071) --- mlx/distributed/ops.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/distributed/ops.cpp b/mlx/distributed/ops.cpp index 1762f0e6bc..583a371edd 100644 --- a/mlx/distributed/ops.cpp +++ b/mlx/distributed/ops.cpp @@ -165,11 +165,11 @@ array sum_scatter( if (group.size() == 1) { return x; } - if (x.shape()[0] % group.size() != 0) { + if (x.ndim() == 0 || x.shape()[0] % group.size() != 0) { std::ostringstream msg; msg << "[sum_scatter] Invalid shape=" << x.shape() << " for a group of size " << group.size() - << ". The first dimension (axis 0) must be divisible by the group size."; + << ". The first dimension (axis 0) must be divisible by the group size and input mst have at least one."; throw std::invalid_argument(msg.str()); } From b24b79fc90dd4096796535f44610517f579e49de Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sat, 8 Aug 2026 03:05:44 -0700 Subject: [PATCH 103/222] Return an empty result from pinv for zero-size inputs (#4069) --- mlx/linalg.cpp | 9 +++++++++ python/tests/test_linalg.py | 7 +++++++ 2 files changed, 16 insertions(+) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index 75a662b1e6..159d2468bc 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -362,6 +362,15 @@ array pinv(const array& a, StreamOrDevice s /* = {} */) { throw std::invalid_argument(msg.str()); } + // The cutoff below reduces over the singular values, which cannot run on an + // empty array. Nothing needs computing anyway, and the result is the shape + // of the transposed input, like numpy. + if (a.size() == 0) { + auto out_shape = a.shape(); + std::swap(out_shape[a.ndim() - 1], out_shape[a.ndim() - 2]); + return zeros(std::move(out_shape), a.dtype(), s); + } + int m = a.shape(-2); int n = a.shape(-1); int k = std::min(m, n); diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index a23de1a523..5284f99baf 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -310,6 +310,13 @@ def test_pseudo_inverse(self): A_plus = mx.linalg.pinv(A, stream=mx.cpu) self.assertTrue(mx.allclose(A @ A_plus @ A, A)) + # Zero-size inputs. The result takes the shape of the transposed input. + for shape in [(0, 0), (0, 3), (3, 0), (0, 2, 2), (2, 0, 0), (0, 4, 3)]: + A_np = np.zeros(shape, dtype=np.float32) + A_plus = mx.linalg.pinv(mx.array(A_np), stream=mx.cpu) + mx.eval(A_plus) + self.assertEqual(A_plus.shape, np.linalg.pinv(A_np).shape) + def test_cholesky_inv(self): mx.random.seed(7) From 6539d180781112bc4e4f07c7b03c434bc0b46db8 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sat, 8 Aug 2026 03:58:08 -0700 Subject: [PATCH 104/222] Handle zero-size inputs in eig and qr on the CPU (#4068) Co-authored-by: Cheng --- mlx/backend/cpu/eig.cpp | 12 ++++++++++-- mlx/backend/cpu/eigh.cpp | 3 +-- mlx/backend/cpu/qrf.cpp | 8 ++++++++ python/tests/test_linalg.py | 23 +++++++++++++++++++++++ 4 files changed, 42 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cpu/eig.cpp b/mlx/backend/cpu/eig.cpp index fb63eefdb9..7bf674cfda 100644 --- a/mlx/backend/cpu/eig.cpp +++ b/mlx/backend/cpu/eig.cpp @@ -228,6 +228,16 @@ void Eig::eval_cpu( const auto& a = inputs[0]; auto& values = outputs[0]; + values.set_data(allocator::malloc(values.nbytes())); + + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. + if (a.shape(-1) == 0) { + if (compute_eigenvectors_) { + outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); + } + return; + } + auto vectors = compute_eigenvectors_ ? outputs[1] : array(a.shape(), complex64, nullptr, {}); @@ -239,8 +249,6 @@ void Eig::eval_cpu( a.flags().row_contiguous ? CopyType::Vector : CopyType::General, stream()); - values.set_data(allocator::malloc(values.nbytes())); - if (compute_eigenvectors_) { // Set the strides and flags so the eigenvectors // are in the columns of the output diff --git a/mlx/backend/cpu/eigh.cpp b/mlx/backend/cpu/eigh.cpp index 813e096320..ce624453ba 100644 --- a/mlx/backend/cpu/eigh.cpp +++ b/mlx/backend/cpu/eigh.cpp @@ -194,8 +194,7 @@ void Eigh::eval_cpu( values.set_data(allocator::malloc(values.nbytes())); - // Nothing to decompose for n = 0; LAPACK rejects lda = 0. The input is - // square, so both outputs are empty and there is nothing to write. + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. if (a.shape(-1) == 0) { if (compute_eigenvectors_) { outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); diff --git a/mlx/backend/cpu/qrf.cpp b/mlx/backend/cpu/qrf.cpp index 13c7e11321..f438d492ae 100644 --- a/mlx/backend/cpu/qrf.cpp +++ b/mlx/backend/cpu/qrf.cpp @@ -12,6 +12,14 @@ template void qrf_impl(const array& a, array& q, array& r, Stream stream) { const int M = a.shape(-2); const int N = a.shape(-1); + + // Nothing to factorize when either dimension is zero; LAPACK rejects lda = 0. + if (M == 0 || N == 0) { + q.set_data(allocator::malloc(q.nbytes())); + r.set_data(allocator::malloc(r.nbytes())); + return; + } + const int lda = M; size_t num_matrices = a.size() / (M * N); diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index 5284f99baf..9b3859976f 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -154,6 +154,16 @@ def test_qr_factorization(self): mx.allclose(out, mx.eye(min(A.shape)), rtol=1e-4, atol=1e-6) ) + # Zero-size inputs. Both factors carry min(M, N) as a dimension, so + # both are empty whichever dimension is zero. + for shape in [(0, 0), (3, 0, 0), (0, 4, 4), (5, 0), (0, 5)]: + A_np = np.zeros(shape, dtype=np.float32) + Q, R = mx.linalg.qr(mx.array(A_np), stream=mx.cpu) + mx.eval(Q, R) + Q_np, R_np = np.linalg.qr(A_np) + self.assertEqual(Q.shape, Q_np.shape) + self.assertEqual(R.shape, R_np.shape) + def test_svd_decomposition(self): A = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float32) U, S, Vt = mx.linalg.svd(A, compute_uv=True, stream=mx.cpu) @@ -459,6 +469,19 @@ def check_eigs_and_vecs(A_np, kwargs={}): self.assertEqual(eig_vals_c64.dtype, mx.complex64) self.assertEqual(eig_vecs_c64.dtype, mx.complex64) + # Zero-size inputs. The input is square, so both outputs are empty. + for shape in [(0, 0), (3, 0, 0), (0, 4, 4)]: + A_np = np.zeros(shape, dtype=np.float32) + eig_vals, eig_vecs = mx.linalg.eig(mx.array(A_np), stream=mx.cpu) + mx.eval(eig_vals, eig_vecs) + vals_np, vecs_np = np.linalg.eig(A_np) + self.assertEqual(eig_vals.shape, vals_np.shape) + self.assertEqual(eig_vecs.shape, vecs_np.shape) + + vals_only = mx.linalg.eigvals(mx.array(A_np), stream=mx.cpu) + mx.eval(vals_only) + self.assertEqual(vals_only.shape, vals_np.shape) + # Test error cases with self.assertRaises(ValueError): mx.linalg.eig(mx.array([1.0, 2.0])) # 1D array From 8d666298652fdac2e7727ecdcf507b1d199bba16 Mon Sep 17 00:00:00 2001 From: "Duhyeon, Kim" <49020301+dudududukim@users.noreply.github.com> Date: Sat, 8 Aug 2026 20:17:31 +0900 Subject: [PATCH 105/222] Normalize biases before encoding in gather_qmm_rhs (#4056) Co-authored-by: Cheng --- mlx/backend/metal/quantized.cpp | 18 ++++++++++++------ python/tests/test_quantized.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 78d8cf0df5..7c461bc1b5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1476,6 +1476,10 @@ void gather_qmm_rhs_nax( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 64, bn = 64, bk = 64; @@ -1554,9 +1558,8 @@ void gather_qmm_rhs_nax( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); @@ -1627,6 +1630,10 @@ void gather_qmm_rhs( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 16, bn = 32, bk = 32; @@ -1704,9 +1711,8 @@ void gather_qmm_rhs( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..0b756b39d5 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1403,6 +1403,34 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y3, atol=tol)) self.assertTrue(mx.allclose(y1, y4, atol=tol)) + @unittest.skipIf(mx.cuda.is_available(), "Not implemented for CUDA") + def test_gather_qmm_sorted_sliced_weight(self): + E, R, D, N = 8, 64, 256, 64 + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + mx.random.seed(0) + w = (mx.random.normal((E, 2 * R, D)) * 0.05).astype(dtype) + qw, s, b = mx.quantize(w, group_size=64, bits=4) + x = (mx.random.normal((N, 1, D)) * 0.5).astype(dtype) + indices = mx.sort(mx.random.randint(0, E, (N,)).astype(mx.uint32)) + + for sl in (slice(0, R), slice(R, 2 * R)): + view = (qw[:, sl], s[:, sl], b[:, sl]) + copy = tuple(mx.contiguous(a) for a in view) + kwargs = dict( + rhs_indices=indices, + transpose=True, + group_size=64, + bits=4, + sorted_indices=True, + ) + self.assertTrue( + mx.allclose( + mx.gather_qmm(x, *view, **kwargs), + mx.gather_qmm(x, *copy, **kwargs), + atol=1e-4, + ) + ) + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: From 5bc46282785fe57a19fae4846a740f10099798f7 Mon Sep 17 00:00:00 2001 From: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com> Date: Sat, 8 Aug 2026 20:32:24 -0700 Subject: [PATCH 106/222] Keep randint samples inside [low, high) (#4012) Co-authored-by: Cheng --- mlx/random.cpp | 7 ++++++- python/src/random.cpp | 6 ++++++ python/tests/test_random.py | 20 ++++++++++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) diff --git a/mlx/random.cpp b/mlx/random.cpp index 4885382c7e..1937743793 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -285,7 +285,12 @@ array randint( "[randint] randint only accepts integer dtypes and bool."); } auto u = uniform(low, high, shape, float32, key, s); - return astype(maximum(u, low, s), dtype, s); + // astype truncates decimal parts so -1.7 becomes -1, use floor. + auto out = astype(floor(u, s), dtype, s); + // low/high may not be representable in float32 so actual range of uniform + // may be larger than [low, high) and we have to clamp to [low, high - 1]. + auto hi = astype(subtract(high, array(1, high.dtype()), s), dtype, s); + return maximum(minimum(out, hi, s), astype(low, dtype, s), s); } array bernoulli( diff --git a/python/src/random.cpp b/python/src/random.cpp index ceebb52e7c..032c99b9ad 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -322,6 +322,12 @@ void init_random(nb::module_& parent_module) { half-open interval ``[low, high)``. The lower and upper bound can be scalars or arrays and must be broadcastable to ``shape``. + .. note:: + The samples are drawn from a ``float32`` uniform and clamped to + ``[low, high - 1]``, so not every integer in the range is reachable + once the bounds or the width of the interval go beyond the + ``2**24`` integer resolution of ``float32``. + Args: low (scalar or array): Lower bound of the interval. high (scalar or array): Upper bound of the interval. diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 8cb1cbb43b..238db5971b 100644 --- a/python/tests/test_random.py +++ b/python/tests/test_random.py @@ -233,6 +233,26 @@ def test_randint(self): a = mx.random.randint(10, -10, [1000, 1000]) self.assertTrue(mx.all(a == 10).item()) + # Bounds hold when the interval is not exactly representable in float32 + for dtype, low, high in [ + (mx.int32, 2**24, 2**24 + 2), + (mx.uint32, 2**24, 2**24 + 2), + (mx.int64, 2**40, 2**40 + 1024), + ]: + a = mx.random.randint(low, high, [10000], dtype=dtype, key=key) + self.assertTrue(mx.all(a >= low).item()) + self.assertTrue(mx.all(a < high).item()) + + # The lower bound is reachable when the interval spans negative values + a = mx.random.randint(-5, 5, [20000], key=key) + self.assertEqual(sorted(set(a.tolist())), list(range(-5, 5))) + + # Booleans use the whole interval + a = mx.random.randint(0, 2, [1000], dtype=mx.bool_, key=key) + self.assertEqual(sorted(set(a.tolist())), [False, True]) + a = mx.random.randint(0, 1, [1000], dtype=mx.bool_, key=key) + self.assertFalse(mx.any(a).item()) + self.assertEqual( mx.random.randint(0, 1).dtype, mx.random.randint(0, 1, dtype=None).dtype ) From 1f9fed2d04b3f19a390a01897a0af4e4611bd876 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sat, 8 Aug 2026 20:33:12 -0700 Subject: [PATCH 107/222] chore: Allow saving empty arrays to npy and safetensors (#4080) --- mlx/io/load.cpp | 10 +++++----- mlx/io/safetensors.cpp | 13 +++++-------- python/tests/test_load.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/mlx/io/load.cpp b/mlx/io/load.cpp index de3bb60ba3..a9379c3079 100644 --- a/mlx/io/load.cpp +++ b/mlx/io/load.cpp @@ -148,10 +148,6 @@ void save(std::shared_ptr out_stream, array a) { a = contiguous(a, true); a.eval(); - if (a.nbytes() == 0) { - throw std::invalid_argument("[save] cannot serialize an empty array"); - } - //////////////////////////////////////////////////////// // Check file if (!out_stream->good() || !out_stream->is_open()) { @@ -212,7 +208,11 @@ void save(std::shared_ptr out_stream, array a) { out_stream->write(magic_ver_len.str().c_str(), magic_ver_len.str().length()); out_stream->write(header.str().c_str(), header.str().length()); - out_stream->write(a.data(), a.nbytes()); + // An empty array has no data to write, and asking for its pointer is not + // meaningful, so stop after the header. + if (a.nbytes() > 0) { + out_stream->write(a.data(), a.nbytes()); + } } /** Save array to file in .npy format */ diff --git a/mlx/io/safetensors.cpp b/mlx/io/safetensors.cpp index a8d8627bac..d1b8440a06 100644 --- a/mlx/io/safetensors.cpp +++ b/mlx/io/safetensors.cpp @@ -250,13 +250,6 @@ void save_safetensors( size_t offset = 0; for (auto& [key, arr] : a) { - if (arr.nbytes() == 0) { - std::ostringstream msg; - msg << "[save_safetensors] Cannot serialize an empty array ('" << key - << "')"; - throw std::invalid_argument(msg.str()); - } - json child; child["dtype"] = dtype_to_safetensor_str(arr.dtype()); child["shape"] = arr.shape(); @@ -270,7 +263,11 @@ void save_safetensors( out_stream->write(reinterpret_cast(&header_len), 8); out_stream->write(header.c_str(), header_len); for (auto& [key, arr] : a) { - out_stream->write(arr.data(), arr.nbytes()); + // An empty tensor contributes a zero length span and has no data pointer + // worth asking for. + if (arr.nbytes() > 0) { + out_stream->write(arr.data(), arr.nbytes()); + } } } diff --git a/python/tests/test_load.py b/python/tests/test_load.py index 10fb63ea63..1c52f333a6 100644 --- a/python/tests/test_load.py +++ b/python/tests/test_load.py @@ -524,6 +524,34 @@ def test_load_donation(self): self.assertEqual(load_only, load_with_binary) + def test_save_and_load_empty(self): + for i, shape in enumerate([(0,), (0, 3), (3, 0), (2, 0, 4)]): + with self.subTest(shape=shape): + save_arr = mx.zeros(shape) + + npy_file = os.path.join(self.test_dir, f"empty_{i}.npy") + mx.save(npy_file, save_arr) + self.assertEqual(mx.load(npy_file).shape, shape) + # numpy can read what we wrote + self.assertEqual(np.load(npy_file).shape, shape) + + st_file = os.path.join(self.test_dir, f"empty_{i}.safetensors") + mx.save_safetensors(st_file, {"x": save_arr}) + self.assertEqual(mx.load(st_file)["x"].shape, shape) + + # An empty array alongside a normal one round trips both + npz_file = os.path.join(self.test_dir, "empty.npz") + mx.savez(npz_file, x=mx.zeros((0, 3)), y=mx.ones((2, 2))) + loaded = mx.load(npz_file) + self.assertEqual(loaded["x"].shape, (0, 3)) + self.assertTrue(mx.array_equal(loaded["y"], mx.ones((2, 2)))) + + st_file = os.path.join(self.test_dir, "empty_mixed.safetensors") + mx.save_safetensors(st_file, {"x": mx.zeros((0, 3)), "y": mx.ones((2, 2))}) + loaded = mx.load(st_file) + self.assertEqual(loaded["x"].shape, (0, 3)) + self.assertTrue(mx.array_equal(loaded["y"], mx.ones((2, 2)))) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 6a0dd0fd2478e013c51e67ef7a89443043619673 Mon Sep 17 00:00:00 2001 From: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:35:28 +0800 Subject: [PATCH 108/222] Fix DeviceType annotations for device and stream arguments (#4059) --- mlx/backend/common/metal_kernel.cpp | 2 + mlx/utils.cpp | 6 + mlx/utils.h | 8 +- python/mlx/_stub_patterns.txt | 30 +- python/src/array.cpp | 2 +- python/src/device.cpp | 4 + python/src/distributed.cpp | 18 +- python/src/export.cpp | 4 +- python/src/fast.cpp | 12 +- python/src/linalg.cpp | 37 ++- python/src/ops.cpp | 431 ++++++++++++++-------------- python/src/random.cpp | 22 +- python/src/stream.cpp | 6 + python/src/transforms.cpp | 6 +- 14 files changed, 295 insertions(+), 293 deletions(-) diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index 4795f08778..43ed8c2047 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -37,9 +37,11 @@ Stream resolve_metal_kernel_stream(StreamOrDevice s) { // recorded in the graph on a placeholder GPU stream. The importing process // remaps it to one of its own streams. auto* device = std::get_if(&s); + auto* device_type = std::get_if(&s); auto* stream = std::get_if(&s); auto* tl_stream = std::get_if(&s); if ((device && *device != Device::gpu) || + (device_type && *device_type != Device::gpu) || (stream && stream->device != Device::gpu) || (tl_stream && tl_stream->device != Device::gpu)) { throw std::invalid_argument("[metal_kernel] Only supports the GPU."); diff --git a/mlx/utils.cpp b/mlx/utils.cpp index ffe1e23173..4e29e8847b 100644 --- a/mlx/utils.cpp +++ b/mlx/utils.cpp @@ -18,6 +18,8 @@ Stream to_stream(StreamOrDevice s) { return default_stream(default_device()); } else if (std::holds_alternative(s)) { return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return default_stream(std::get(s)); } else if (std::holds_alternative(s)) { return stream_from_thread_local_stream(std::get(s)); } else { @@ -30,6 +32,10 @@ Stream to_stream(StreamOrDevice s, Device default_) { return default_stream(default_); } else if (std::holds_alternative(s)) { return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return stream_from_thread_local_stream(std::get(s)); } else { return std::get(s); } diff --git a/mlx/utils.h b/mlx/utils.h index 9fa332fd2e..b5b516d89c 100644 --- a/mlx/utils.h +++ b/mlx/utils.h @@ -17,8 +17,12 @@ namespace mlx::core { -using StreamOrDevice = - std::variant; +using StreamOrDevice = std::variant< + std::monostate, + Stream, + ThreadLocalStream, + Device, + Device::DeviceType>; MLX_API Stream to_stream(StreamOrDevice s); MLX_API Stream to_stream(StreamOrDevice s, Device default_); diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt index c05f879d13..2a3fbce133 100644 --- a/python/mlx/_stub_patterns.txt +++ b/python/mlx/_stub_patterns.txt @@ -1,10 +1,5 @@ mlx.core.__prefix__: - from typing import Any, Callable, Dict, List, Optional, Protocol, Sequence, Tuple, Union, ParamSpec, TypeVar - import sys - if sys.version_info >= (3, 10): - from typing import TypeAlias - else: - from typing_extensions import TypeAlias + from typing import Any, ParamSpec, Protocol, TypeAlias, TypeVar P = ParamSpec("P") R = TypeVar("R") class DLPackCompatible(Protocol): @@ -12,28 +7,27 @@ mlx.core.__prefix__: __dlpack_device__: Callable[..., Any] mlx.core.__suffix__: - from typing import Union - scalar: TypeAlias = Union[int, float, bool] - list_or_scalar: TypeAlias = Union[scalar, list["list_or_scalar"]] + scalar: TypeAlias = int | float | bool + list_or_scalar: TypeAlias = scalar | list["list_or_scalar"] + StreamOrDevice: TypeAlias = Stream | ThreadLocalStream | Device | DeviceType | None bool_: Dtype = ... mlx.core.distributed.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar + from mlx.core import array, Dtype, StreamOrDevice, scalar from mlx.core.distributed import Group - from typing import Sequence, Optional, Union + from collections.abc import Sequence mlx.core.fast.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Union + from mlx.core import array, Dtype, StreamOrDevice, scalar mlx.core.linalg.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Tuple, Union + from mlx.core import array, Dtype, StreamOrDevice, scalar + from collections.abc import Sequence mlx.core.metal.__prefix__: from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Union + from collections.abc import Sequence mlx.core.random.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar, float32, int32 - from typing import Sequence, Optional, Union + from mlx.core import array, Dtype, StreamOrDevice, scalar, float32, int32 + from collections.abc import Sequence diff --git a/python/src/array.cpp b/python/src/array.cpp index 653b4a4c5b..6f65e14cde 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -312,7 +312,7 @@ void init_array(nb::module_& m) { "val"_a, "dtype"_a = nb::none(), nb::sig( - "def __init__(self: array, val: Union[scalar, list, tuple, DLPackCompatible, array], dtype: Optional[Dtype] = None)")) + "def __init__(self: array, val: scalar | list | tuple | DLPackCompatible | array, dtype: Dtype | None = None)")) .def_prop_ro( "size", &mx::array::size, diff --git a/python/src/device.cpp b/python/src/device.cpp index e70b69bd34..83d32d1cb6 100644 --- a/python/src/device.cpp +++ b/python/src/device.cpp @@ -61,11 +61,13 @@ void init_device(nb::module_& m) { "set_default_device", &mx::set_default_device, "device"_a, + nb::sig("def set_default_device(device: Device | DeviceType) -> None"), R"pbdoc(Set the default device.)pbdoc"); m.def( "is_available", &mx::is_available, "device"_a, + nb::sig("def is_available(device: Device | DeviceType) -> bool"), R"pbdoc(Check if a back-end is available for the given device.)pbdoc"); m.def( "device_count", @@ -86,6 +88,8 @@ void init_device(nb::module_& m) { return mx::device_info(d.value_or(mx::default_device())); }, "d"_a = nb::none(), + nb::sig( + "def device_info(d: None | Device | DeviceType = None) -> dict[str, str | int]"), R"pbdoc( Get information about a device. diff --git a/python/src/distributed.cpp b/python/src/distributed.cpp index ed80001df5..0ad91c6e48 100644 --- a/python/src/distributed.cpp +++ b/python/src/distributed.cpp @@ -125,7 +125,7 @@ void init_distributed(nb::module_& parent_module) { nb::kw_only(), "all_gather_factory"_a = nb::none(), nb::sig( - "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Optional[Callable[[int, int], Callable[[bytes, int], bytes]]] = None) -> Group"), + "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Callable[[int, int], Callable[[bytes, int], bytes]] | None = None) -> Group"), R"pbdoc( Initialize the communication backend and create the global communication group. @@ -170,7 +170,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_sum(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_sum(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce sum. @@ -199,7 +199,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_max(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_max(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce max. @@ -228,7 +228,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_min(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_min(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce min. @@ -257,7 +257,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_gather(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_gather(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Gather arrays from all processes. @@ -290,7 +290,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def send(x: array, dst: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def send(x: array, dst: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Send an array from the current process to the process that has rank ``dst`` in the group. @@ -318,7 +318,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Recv an array with shape ``shape`` and dtype ``dtype`` from process with rank ``src``. @@ -351,7 +351,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def recv_like(x: array, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def recv_like(x: array, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Recv an array with shape and type like ``x`` from process with rank ``src``. @@ -384,7 +384,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def sum_scatter(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def sum_scatter(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sum ``x`` across all processes in the group and shard the result along the first axis across ranks. ``x.shape[0]`` must be divisible by the group size. diff --git a/python/src/export.cpp b/python/src/export.cpp index 30d48bcae4..3596e4f846 100644 --- a/python/src/export.cpp +++ b/python/src/export.cpp @@ -173,7 +173,7 @@ void init_export(nb::module_& m) { "metadata"_a = nb::none(), "kwargs"_a, nb::sig( - "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, metadata: Optional[str] = None, **kwargs) -> None"), + "def export_function(file_or_callback: str | Callable, fun: Callable, *args, shapeless: bool = False, metadata: str | None = None, **kwargs) -> None"), R"pbdoc( Export an MLX function. @@ -236,7 +236,7 @@ void init_export(nb::module_& m) { "file"_a, "return_metadata"_a = false, nb::sig( - "def import_function(file: str, return_metadata: bool = False) -> Union[Callable, tuple[Callable, str]]"), + "def import_function(file: str, return_metadata: bool = False) -> Callable | tuple[Callable, str]"), R"pbdoc( Import a function from a file. diff --git a/python/src/fast.cpp b/python/src/fast.cpp index cd30b0bacd..e59357bc33 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -127,7 +127,7 @@ void init_fast(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), + "def rms_norm(x: array, weight: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Root Mean Square normalization (RMS norm). @@ -154,7 +154,7 @@ void init_fast(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), + "def layer_norm(x: array, weight: array | None, bias: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Layer normalization. @@ -197,7 +197,7 @@ void init_fast(nb::module_& parent_module) { "freqs"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def rope(a: array, dims: int, *, traditional: bool, base: Optional[float], scale: float, offset: Union[int, array], freqs: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def rope(a: array, dims: int, *, traditional: bool, base: float | None, scale: float, offset: int | array, freqs: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Apply rotary positional encoding to the input. @@ -271,7 +271,7 @@ void init_fast(nb::module_& parent_module) { "sinks"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: Union[None, str, array] = None, sinks: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: None | str | array = None, sinks: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``. @@ -365,7 +365,7 @@ void init_fast(nb::module_& parent_module) { "verbose"_a = false, "stream"_a = nb::none(), nb::sig( - "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"), + "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"), R"pbdoc( Run the kernel. @@ -489,7 +489,7 @@ void init_fast(nb::module_& parent_module) { "verbose"_a = false, "stream"_a = nb::none(), nb::sig( - "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"), + "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"), R"pbdoc( Run the kernel. diff --git a/python/src/linalg.cpp b/python/src/linalg.cpp index 0bf7b6f12c..a5e82ae853 100644 --- a/python/src/linalg.cpp +++ b/python/src/linalg.cpp @@ -55,7 +55,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def norm(a: array, /, ord: Union[None, int, float, str] = None, axis: Union[None, int, list[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def norm(a: array, /, ord: None | int | float | str = None, axis: None | int | list[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix or vector norm. @@ -177,7 +177,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def qr(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def qr(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( The QR factorization of the input matrix. @@ -220,7 +220,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def svd(a: array, compute_uv: bool = True, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"), + "def svd(a: array, compute_uv: bool = True, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( The Singular Value Decomposition (SVD) of the input matrix. @@ -246,8 +246,7 @@ void init_linalg(nb::module_& parent_module) { "a"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def inv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def inv(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a square matrix. @@ -271,7 +270,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tri_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tri_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a triangular square matrix. @@ -296,7 +295,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cholesky(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cholesky(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the Cholesky decomposition of a real symmetric positive semi-definite matrix. @@ -326,7 +325,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cholesky_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cholesky_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition. @@ -363,8 +362,7 @@ void init_linalg(nb::module_& parent_module) { "a"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def pinv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def pinv(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the (Moore-Penrose) pseudo-inverse of a matrix. @@ -390,7 +388,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cross(a: array, b: array, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cross(a: array, b: array, axis: int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the cross product of two arrays along a specified axis. @@ -449,7 +447,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eig(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def eig(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Compute the eigenvalues and eigenvectors of a square matrix. @@ -527,7 +525,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eigh(a: array, UPLO: str = 'L', *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def eigh(a: array, UPLO: str = 'L', *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Compute the eigenvalues and eigenvectors of a complex Hermitian or real symmetric matrix. @@ -573,7 +571,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def lu(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"), + "def lu(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( Compute the LU factorization of the given matrix ``A``. @@ -604,7 +602,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def lu_factor(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def lu_factor(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Computes a compact representation of the LU factorization. @@ -624,7 +622,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def solve(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def solve(a: array, b: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the solution to a system of linear equations ``AX = B``. @@ -646,7 +644,7 @@ void init_linalg(nb::module_& parent_module) { "upper"_a = false, "stream"_a = nb::none(), nb::sig( - "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Computes the solution of a triangular system of linear equations ``AX = B``. @@ -668,8 +666,7 @@ void init_linalg(nb::module_& parent_module) { "a"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def det(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def det(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the determinant of a square matrix. @@ -701,7 +698,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def slogdet(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def slogdet(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Compute the sign and natural log of the absolute value of the determinant of a square matrix. diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 4e0373a238..8ddb3e4f96 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -63,7 +63,7 @@ void init_ops(nb::module_& m) { "stream"_a = nb::none(), nb::sig( "def reshape(a: array, /, shape: Sequence[int], *, stream: " - "Union[None, Stream, Device] = None) -> array"), + "StreamOrDevice = None) -> array"), R"pbdoc( Reshape an array while preserving the size. @@ -91,7 +91,7 @@ void init_ops(nb::module_& m) { "stream"_a = nb::none(), nb::sig( "def flatten(a: array, /, start_axis: int = 0, end_axis: int = " - "-1, *, stream: Union[None, Stream, Device] = None) -> array"), + "-1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Flatten an array. @@ -127,7 +127,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Unflatten an axis of an array to a shape. @@ -164,8 +164,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def squeeze(a: array, /, axis: Union[None, int, Sequence[int]] = " - "None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def squeeze(a: array, /, axis: None | int | Sequence[int] = " + "None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Remove length one axes from an array. @@ -193,8 +193,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def flip(a: array, /, axis: Union[None, int, Sequence[int]] = None, " - "*, stream: Union[None, Stream, Device] = None) -> array"), + "def flip(a: array, /, axis: None | int | Sequence[int] = None, " + "*, stream: StreamOrDevice = None) -> array"), R"pbdoc( Reverse the order of elements along the given axis. @@ -216,8 +216,7 @@ void init_ops(nb::module_& m) { "axis"_a = 0, "stream"_a = nb::none(), nb::sig( - "def unstack(x: array, /, *, axis: int = 0, stream: Union[None, " - "Stream, Device] = None) -> list[array]"), + "def unstack(x: array, /, *, axis: int = 0, stream: StreamOrDevice = None) -> list[array]"), R"pbdoc( Split an array into a sequence of arrays along the given axis. @@ -247,8 +246,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def expand_dims(a: array, /, axis: Union[int, Sequence[int]], " - "*, stream: Union[None, Stream, Device] = None) -> array"), + "def expand_dims(a: array, /, axis: int | Sequence[int], " + "*, stream: StreamOrDevice = None) -> array"), R"pbdoc( Add a size one dimension at the given axis. @@ -268,7 +267,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def abs(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def abs(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise absolute value. @@ -287,7 +286,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sign(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sign(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise sign. @@ -304,7 +303,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def positive(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def positive(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise unary plus. Returns a copy of the input. @@ -323,7 +322,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def negative(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def negative(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise negation. @@ -346,7 +345,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def add(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def add(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise addition. @@ -373,7 +372,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def subtract(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def subtract(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise subtraction. @@ -400,7 +399,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def multiply(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def multiply(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise multiplication. @@ -427,7 +426,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise division. @@ -454,7 +453,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def divmod(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def divmod(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise quotient and remainder. @@ -482,7 +481,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def floor_divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def floor_divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise integer division. @@ -509,7 +508,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def remainder(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def remainder(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise remainder of division. @@ -537,7 +536,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise equality. @@ -564,7 +563,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def not_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def not_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise not equal. @@ -591,7 +590,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def less(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def less(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise less than. @@ -618,7 +617,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def less_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def less_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise less than or equal. @@ -645,7 +644,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def greater(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def greater(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise greater than. @@ -672,7 +671,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def greater_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def greater_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise greater or equal. @@ -701,7 +700,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def array_equal(a: Union[scalar, array], b: Union[scalar, array], equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def array_equal(a: scalar | array, b: scalar | array, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Array equality check. @@ -726,7 +725,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def matmul(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def matmul(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication. @@ -757,7 +756,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def trunc(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def trunc(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise truncation towards zero. @@ -776,7 +775,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def square(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def square(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise square. @@ -795,7 +794,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise square root. @@ -814,7 +813,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def rsqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def rsqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise reciprocal and square root. @@ -833,7 +832,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def reciprocal(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def reciprocal(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise reciprocal. @@ -852,7 +851,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_not(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_not(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical not. @@ -872,7 +871,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_and(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_and(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical and. @@ -894,7 +893,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_or(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_or(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical or. @@ -915,7 +914,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_xor(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_xor(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical exclusive or. @@ -939,7 +938,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logaddexp(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logaddexp(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise log-add-exp. @@ -964,7 +963,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def exp(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def exp(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise exponential. @@ -983,7 +982,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def expm1(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def expm1(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise exponential minus 1. @@ -1004,7 +1003,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def erf(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def erf(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise error function. @@ -1026,7 +1025,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def erfinv(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def erfinv(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse of :func:`erf`. @@ -1045,7 +1044,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sin(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise sine. @@ -1064,7 +1063,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cos(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise cosine. @@ -1083,7 +1082,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tan(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise tangent. @@ -1102,7 +1101,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arcsin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arcsin(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse sine. @@ -1121,7 +1120,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arccos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arccos(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse cosine. @@ -1140,7 +1139,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctan(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse tangent. @@ -1158,7 +1157,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctan2(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctan2(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse tangent of the ratio of two arrays. @@ -1178,7 +1177,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sinh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic sine. @@ -1197,7 +1196,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cosh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic cosine. @@ -1216,7 +1215,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tanh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic tangent. @@ -1235,7 +1234,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arcsinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arcsinh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic sine. @@ -1254,7 +1253,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arccosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arccosh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic cosine. @@ -1273,7 +1272,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctanh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic tangent. @@ -1292,7 +1291,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def degrees(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def degrees(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from radians to degrees. @@ -1311,7 +1310,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def radians(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def radians(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from degrees to radians. @@ -1330,7 +1329,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise natural logarithm. @@ -1349,7 +1348,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log2(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log2(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise base-2 logarithm. @@ -1368,7 +1367,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log10(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log10(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise base-10 logarithm. @@ -1387,7 +1386,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log1p(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log1p(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise natural log of one plus the array. @@ -1404,7 +1403,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def stop_gradient(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def stop_gradient(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Stop gradients from being computed. @@ -1428,7 +1427,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sigmoid(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sigmoid(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logistic sigmoid. @@ -1456,7 +1455,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def power(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def power(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise power operation. @@ -1503,7 +1502,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arange(start : Union[int, float], stop : Union[None, int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arange(start : int | float, stop : None | int | float, step : None | int | float, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generates ranges of numbers. @@ -1548,7 +1547,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arange(stop : Union[int, float], step : Union[None, int, float] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array")); + "def arange(stop : int | float, step : None | int | float = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array")); m.def( "bartlett", &mlx::core::bartlett, @@ -1599,8 +1598,7 @@ void init_ops(nb::module_& m) { "M"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def hamming(M: int, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def hamming(M: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the Hamming window. @@ -1624,7 +1622,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def blackman(M: int, *, stream: Union[None, Stream, Device] = None) -> array"), // <--- J'ai rajouté ça + "def blackman(M: int, *, stream: StreamOrDevice = None) -> array"), // <--- J'ai rajouté ça R"pbdoc( Return the Blackman window. @@ -1661,7 +1659,7 @@ void init_ops(nb::module_& m) { "dtype"_a.none() = mx::float32, "stream"_a = nb::none(), nb::sig( - "def linspace(start: scalar, stop: scalar, num: Optional[int] = 50, dtype: Optional[Dtype] = float32, stream: Union[None, Stream, Device] = None) -> array"), + "def linspace(start: scalar, stop: scalar, num: int | None = 50, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. @@ -1683,14 +1681,14 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def kron(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def kron(a: array, b: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the Kronecker product of two arrays ``a`` and ``b``. Args: a (array): The first input array. b (array): The second input array. - stream (Union[None, Stream, Device], optional): Optional stream or + stream (StreamOrDevice, optional): Optional stream or device for execution. Default: ``None``. Returns: @@ -1727,7 +1725,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def take(a: array, /, indices: Union[int, array], axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def take(a: array, /, indices: int | array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Take elements along an axis. @@ -1764,7 +1762,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def take_along_axis(a: array, /, indices: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def take_along_axis(a: array, /, indices: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Take values along an axis at the specified indices. @@ -1803,7 +1801,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def put_along_axis(a: array, /, indices: array, values: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def put_along_axis(a: array, /, indices: array, values: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Put values along an axis at the specified indices. @@ -1835,7 +1833,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def full(shape: Union[int, Sequence[int]], vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def full(shape: int | Sequence[int], vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array with the given value. @@ -1866,7 +1864,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def full_like(a: array, vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def full_like(a: array, vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array filled with ``vals`` with the same shape as the input. @@ -1892,7 +1890,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def zeros(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def zeros(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array of zeros. @@ -1914,8 +1912,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "copy"_a = nb::none(), nb::sig( - "def asarray(a: Union[scalar, array, Sequence, DLPackCompatible], dtype: " - "Optional[Dtype] = None, *, copy: Optional[bool] = None) -> array"), + "def asarray(a: scalar | array | Sequence | DLPackCompatible, dtype: " + "Dtype | None = None, *, copy: bool | None = None) -> array"), R"pbdoc( Convert the input to an array. @@ -1942,7 +1940,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "copy"_a = nb::none(), nb::sig( - "def from_dlpack(x: DLPackCompatible, /, *, copy: Optional[bool] = None) -> array"), + "def from_dlpack(x: DLPackCompatible, /, *, copy: bool | None = None) -> array"), R"pbdoc( Create an array from an object that supports DLPack. @@ -1969,7 +1967,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def zeros_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def zeros_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array of zeros like the input. @@ -1994,7 +1992,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ones(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ones(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array of ones. @@ -2018,7 +2016,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ones_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ones_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array of ones like the input. @@ -2046,7 +2044,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eye(n: int, m: Optional[int] = None, k: int = 0, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def eye(n: int, m: int | None = None, k: int = 0, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create an identity matrix or a general diagonal matrix. @@ -2070,7 +2068,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def identity(n: int, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def identity(n: int, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create a square identity matrix. @@ -2098,7 +2096,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tri(n: int, m: int, k: int, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tri(n: int, m: int, k: int, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array with ones at and below the given diagonal and zeros elsewhere. @@ -2120,7 +2118,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tril(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tril(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array above the given diagonal. @@ -2140,7 +2138,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def triu(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def triu(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array below the given diagonal. @@ -2163,7 +2161,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Approximate comparison of two arrays. @@ -2200,7 +2198,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a boolean array where two arrays are element-wise equal within a tolerance. @@ -2241,7 +2239,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def all(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def all(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An `and` reduction over the given axes. @@ -2270,7 +2268,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def any(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def any(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An `or` reduction over the given axes. @@ -2298,7 +2296,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def minimum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def minimum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise minimum. @@ -2325,7 +2323,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def maximum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def maximum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise maximum. @@ -2348,7 +2346,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def floor(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def floor(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise floor. @@ -2367,7 +2365,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ceil(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ceil(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise ceil. @@ -2385,8 +2383,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isnan(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isnan(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are NaN. @@ -2404,8 +2401,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isinf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isinf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are +/- inifnity. @@ -2423,8 +2419,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isfinite(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isfinite(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are finite. @@ -2444,14 +2439,13 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isposinf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isposinf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are positive infinity. Args: a (array): Input array. - stream (Union[None, Stream, Device]): Optional stream or device. + stream (StreamOrDevice): Optional stream or device. Returns: array: The boolean array indicating which elements are positive infinity. @@ -2464,14 +2458,13 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isneginf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isneginf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are negative infinity. Args: a (array): Input array. - stream (Union[None, Stream, Device]): Optional stream or device. + stream (StreamOrDevice): Optional stream or device. Returns: array: The boolean array indicating which elements are negative infinity. @@ -2485,7 +2478,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def moveaxis(a: array, /, source: int, destination: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def moveaxis(a: array, /, source: int, destination: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Move an axis to a new position. @@ -2506,7 +2499,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Swap two axes of an array. @@ -2534,7 +2527,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def transpose(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def transpose(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Transpose the dimensions of the array. @@ -2562,7 +2555,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def permute_dims(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def permute_dims(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( See :func:`transpose`. )pbdoc"); @@ -2580,7 +2573,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sum(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sum(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sum reduce the array over the given axes. @@ -2616,7 +2609,7 @@ void init_ops(nb::module_& m) { "keepdims"_a = false, "stream"_a = nb::none(), nb::sig( - "def count_nonzero(a: array, /, *, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def count_nonzero(a: array, /, *, axis: None | int | Sequence[int] = None, keepdims: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Count the number of non-zero elements along the given axis. @@ -2644,7 +2637,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def prod(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def prod(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An product reduction over the given axes. @@ -2673,7 +2666,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def min(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def min(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `min` reduction over the given axes. @@ -2702,7 +2695,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def max(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def max(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `max` reduction over the given axes. @@ -2738,7 +2731,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def logcumsumexp(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def logcumsumexp(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative logsumexp of the elements along the given axis. @@ -2768,7 +2761,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logsumexp(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logsumexp(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `log-sum-exp` reduction over the given axes. @@ -2803,7 +2796,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def mean(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def mean(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the mean(s) over the given axes. @@ -2832,7 +2825,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def median(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def median(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the median(s) over the given axes. @@ -2863,7 +2856,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def var(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the variance(s) over the given axes. @@ -2896,7 +2889,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def std(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the standard deviation(s) over the given axes. @@ -2932,7 +2925,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def split(a: array, /, indices_or_sections: Union[int, Sequence[int]], axis: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def split(a: array, /, indices_or_sections: int | Sequence[int], axis: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Split an array along a given axis. @@ -2976,7 +2969,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argmin(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argmin(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Indices of the minimum values along the axis. @@ -3008,7 +3001,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argmax(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argmax(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Indices of the maximum values along the axis. @@ -3036,7 +3029,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a sorted copy of the array. @@ -3066,7 +3059,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argsort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argsort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the indices that sort the array. @@ -3100,7 +3093,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def partition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def partition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a partitioned copy of the array such that the smaller ``kth`` elements are first. @@ -3138,7 +3131,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argpartition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argpartition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the indices that partition the array. @@ -3177,7 +3170,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def topk(a: array, /, k: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def topk(a: array, /, k: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the ``k`` largest elements from the input along a given axis. @@ -3203,7 +3196,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def broadcast_to(a: Union[scalar, array], /, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def broadcast_to(a: scalar | array, /, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Broadcast an array to the given shape. @@ -3225,7 +3218,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def broadcast_arrays(*arrays: array, stream: Union[None, Stream, Device] = None) -> Tuple[array, ...]"), + "def broadcast_arrays(*arrays: array, stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Broadcast arrays against one another. @@ -3251,7 +3244,7 @@ void init_ops(nb::module_& m) { "precise"_a = false, "stream"_a = nb::none(), nb::sig( - "def softmax(a: array, /, axis: Union[None, int, Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def softmax(a: array, /, axis: None | int | Sequence[int] = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the softmax along the given axis. @@ -3286,7 +3279,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def concatenate(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def concatenate(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Concatenate the arrays along the given axis. @@ -3314,7 +3307,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def concat(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def concat(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( See :func:`concatenate`. )pbdoc"); @@ -3334,7 +3327,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def stack(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def stack(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Stacks the arrays along a new axis. @@ -3362,7 +3355,7 @@ void init_ops(nb::module_& m) { "indexing"_a = "xy", "stream"_a = nb::none(), nb::sig( - "def meshgrid(*arrays: array, sparse: Optional[bool] = False, indexing: Optional[str] = 'xy', stream: Union[None, Stream, Device] = None) -> array"), + "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate multidimensional coordinate grids from 1-D coordinate arrays @@ -3395,7 +3388,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def repeat(array: array, repeats: int, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Repeat an array along a specified axis. @@ -3432,7 +3425,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def clip(a: array, /, a_min: Union[scalar, array, None], a_max: Union[scalar, array, None], *, stream: Union[None, Stream, Device] = None) -> array"), + "def clip(a: array, /, a_min: scalar | array | None, a_max: scalar | array | None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Clip the values of the array between the given minimum and maximum. @@ -3482,7 +3475,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def pad(a: array, pad_width: Union[int, tuple[int], tuple[int, int], list[tuple[int, int]]], mode: Literal['constant', 'edge'] = 'constant', constant_values: Union[scalar, array] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Pad an array with a constant value @@ -3529,7 +3522,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def as_strided(a: array, /, shape: Optional[Sequence[int]] = None, strides: Optional[Sequence[int]] = None, offset: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def as_strided(a: array, /, shape: Sequence[int] | None = None, strides: Sequence[int] | None = None, offset: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create a view into the array with the given shape and strides. @@ -3566,7 +3559,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def astype(a: array, dtype: Dtype, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def astype(a: array, dtype: Dtype, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Cast the array to a specified type. @@ -3598,7 +3591,7 @@ void init_ops(nb::module_& m) { "dtype"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def cumsum(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def cumsum(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative sum of the elements along the given axis. @@ -3636,7 +3629,7 @@ void init_ops(nb::module_& m) { "dtype"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def cumprod(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def cumprod(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative product of the elements along the given axis. @@ -3673,7 +3666,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def cummax(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cummax(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative maximum of the elements along the given axis. @@ -3709,7 +3702,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def cummin(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cummin(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative minimum of the elements along the given axis. @@ -3734,7 +3727,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( The n-th discrete difference along the given axis. @@ -3755,8 +3748,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def conj(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def conj(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the elementwise complex conjugate of the input. Alias for `mx.conjugate`. @@ -3776,7 +3768,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conjugate(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conjugate(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the elementwise complex conjugate of the input. Alias for `mx.conj`. @@ -3850,7 +3842,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: Union[None, Stream, Device] = None) -> array)"), + R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: StreamOrDevice = None) -> array)"), R"pbdoc( The discrete convolution of 1D arrays. @@ -3877,7 +3869,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 1D convolution over an input with several channels @@ -3935,7 +3927,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv2d(input: array, weight: array, /, stride: Union[int, tuple[int, int]] = 1, padding: Union[int, tuple[int, int]] = 0, dilation: Union[int, tuple[int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 2D convolution over an input with several channels @@ -4005,7 +3997,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv3d(input: array, weight: array, /, stride: Union[int, tuple[int, int, int]] = 1, padding: Union[int, tuple[int, int, int]] = 0, dilation: Union[int, tuple[int, int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 3D convolution over an input with several channels @@ -4041,7 +4033,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 1D transposed convolution over an input with several channels @@ -4116,7 +4108,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose2d(input: array, weight: array, /, stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 0, dilation: Union[int, Tuple[int, int]] = 1, output_padding: Union[int, Tuple[int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, output_padding: int | tuple[int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 2D transposed convolution over an input with several channels @@ -4202,7 +4194,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose3d(input: array, weight: array, /, stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, dilation: Union[int, Tuple[int, int, int]] = 1, output_padding: Union[int, Tuple[int, int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, output_padding: int | tuple[int, int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 3D transposed convolution over an input with several channels @@ -4304,7 +4296,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_general(input: array, weight: array, /, stride: Union[int, Sequence[int]] = 1, padding: Union[int, Sequence[int], tuple[Sequence[int], Sequence[int]]] = 0, kernel_dilation: Union[int, Sequence[int]] = 1, input_dilation: Union[int, Sequence[int]] = 1, groups: int = 1, flip: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_general(input: array, weight: array, /, stride: int | Sequence[int] = 1, padding: int | Sequence[int] | tuple[Sequence[int], Sequence[int]] = 0, kernel_dilation: int | Sequence[int] = 1, input_dilation: int | Sequence[int] = 1, groups: int = 1, flip: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( General convolution over an input with several channels @@ -4337,8 +4329,7 @@ void init_ops(nb::module_& m) { &mlx_save_helper, "file"_a, "arr"_a, - nb::sig( - "def save(file: Union[file, str, pathlib.Path], arr: array) -> None"), + nb::sig("def save(file: file | str | pathlib.Path, arr: array) -> None"), R"pbdoc( Save the array to a binary file in ``.npy`` format. @@ -4354,8 +4345,7 @@ void init_ops(nb::module_& m) { "file"_a, "args"_a, "kwargs"_a, - nb::sig( - "def savez(file: Union[file, str, pathlib.Path], *args, **kwargs)"), + nb::sig("def savez(file: file | str | pathlib.Path, *args, **kwargs)"), R"pbdoc( Save several arrays to a binary file in uncompressed ``.npz`` format. @@ -4389,7 +4379,7 @@ void init_ops(nb::module_& m) { "args"_a, "kwargs"_a, nb::sig( - "def savez_compressed(file: Union[file, str, pathlib.Path], *args, **kwargs)"), + "def savez_compressed(file: file | str | pathlib.Path, *args, **kwargs)"), R"pbdoc( Save several arrays to a binary file in compressed ``.npz`` format. @@ -4408,7 +4398,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def load(file: Union[file, str, pathlib.Path], /, format: Optional[str] = None, return_metadata: bool = False, *, stream: Union[None, Stream, Device] = None) -> Union[array, dict[str, array], Tuple[dict[str, array], dict[str, Any]]]"), + "def load(file: file | str | pathlib.Path, /, format: str | None = None, return_metadata: bool = False, *, stream: StreamOrDevice = None) -> array | dict[str, array] | tuple[dict[str, array], dict[str, Any]]"), R"pbdoc( Load array(s) from a binary file. @@ -4443,7 +4433,7 @@ void init_ops(nb::module_& m) { "arrays"_a, "metadata"_a = nb::none(), nb::sig( - "def save_safetensors(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: Optional[dict[str, str]] = None)"), + "def save_safetensors(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, str] | None = None)"), R"pbdoc( Save array(s) to a binary file in ``.safetensors`` format. @@ -4465,7 +4455,7 @@ void init_ops(nb::module_& m) { "arrays"_a, "metadata"_a = nb::none(), nb::sig( - "def save_gguf(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: dict[str, Union[array, str, list[str]]])"), + "def save_gguf(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, array | str | list[str]])"), R"pbdoc( Save array(s) to a binary file in ``.gguf`` format. @@ -4496,7 +4486,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def where(condition: Union[scalar, array], x: Union[scalar, array], y: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def where(condition: scalar | array, x: scalar | array, y: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Select from ``x`` or ``y`` according to ``condition``. @@ -4528,7 +4518,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def nan_to_num(a: Union[scalar, array], nan: float = 0, posinf: Optional[float] = None, neginf: Optional[float] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def nan_to_num(a: scalar | array, nan: float = 0, posinf: float | None = None, neginf: float | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Replace NaN and Inf values with finite numbers. @@ -4555,7 +4545,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def round(a: array, /, decimals: int = 0, stream: Union[None, Stream, Device] = None) -> array"), + "def round(a: array, /, decimals: int = 0, stream: StreamOrDevice = None) -> array"), R"pbdoc( Round to the given number of decimals. @@ -4588,7 +4578,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def quantized_matmul(x: array, w: array, /, scales: array, biases: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, stream: Union[None, Stream, Device] = None) -> array"), + "def quantized_matmul(x: array, w: array, /, scales: array, biases: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the matrix multiplication with the quantized matrix ``w``. The quantization uses one floating point scale and bias per ``group_size`` of @@ -4626,7 +4616,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def quantize(w: array, /, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, global_scale: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> tuple[array, array, array]"), + "def quantize(w: array, /, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, global_scale: array | None = None, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( Quantize the array ``w``. @@ -4727,7 +4717,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', global_scale: Optional[array] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def dequantize(w: array, /, scales: array, biases: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', global_scale: array | None = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Dequantize the matrix ``w`` using quantization parameters. @@ -4782,7 +4772,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_qmm(x: array, w: array, /, scales: array, biases: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def gather_qmm(x: array, w: array, /, scales: array, biases: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform quantized matrix multiplication with matrix-level gather. @@ -4836,7 +4826,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_qqmm(x: array, w: array, /, scales: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def gather_qqmm(x: array, w: array, /, scales: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Fused :func:`qqmm` with matrix-level gather. @@ -4880,7 +4870,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def segmented_mm(a: array, b: array, /, segments: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def segmented_mm(a: array, b: array, /, segments: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform a matrix multiplication but segment the inner dimension and save the result for each segment separately. @@ -4916,7 +4906,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tensordot(a: array, b: array, /, axes: Union[int, list[Sequence[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tensordot(a: array, b: array, /, axes: int | list[Sequence[int]] = 2, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the tensor dot product along the specified axes. @@ -4940,7 +4930,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def inner(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def inner(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes. @@ -4960,7 +4950,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: Union[None, Stream, Device] = None) -> array"), + "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the vector dot product of two arrays along an axis. @@ -4980,7 +4970,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def outer(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def outer(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand. @@ -5007,7 +4997,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tile(a: array, reps: Union[int, Sequence[int]], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tile(a: array, reps: int | Sequence[int], /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array by repeating ``a`` the number of times given by ``reps``. @@ -5029,7 +5019,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with addition and optional scaling. @@ -5059,7 +5049,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: Optional[array] = None, mask_lhs: Optional[array] = None, mask_rhs: Optional[array] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: array | None = None, mask_lhs: array | None = None, mask_rhs: array | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with block masking. @@ -5098,7 +5088,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with matrix-level gather. @@ -5140,7 +5130,7 @@ void init_ops(nb::module_& m) { "axis2"_a = 1, "stream"_a = nb::none(), nb::sig( - "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: Union[None, Stream, Device] = None) -> array"), + "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return specified diagonals. @@ -5172,7 +5162,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def diag(a: array, /, k: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def diag(a: array, /, k: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Extract a diagonal or construct a diagonal matrix. If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the @@ -5208,7 +5198,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the sum along a specified diagonal in the given array. @@ -5238,13 +5228,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_1d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_1d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least one dimension. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least one dimension. @@ -5261,13 +5251,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_2d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_2d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least two dimensions. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least two dimensions. @@ -5284,13 +5274,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_3d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_3d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least three dimensions. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least three dimensions. @@ -5320,7 +5310,7 @@ void init_ops(nb::module_& m) { ""_a, ""_a, nb::sig( - "def issubdtype(arg1: Union[Dtype, DtypeCategory], arg2: Union[Dtype, DtypeCategory]) -> bool"), + "def issubdtype(arg1: Dtype | DtypeCategory, arg2: Dtype | DtypeCategory) -> bool"), R"pbdoc( Check if a :obj:`Dtype` or :obj:`DtypeCategory` is a subtype of another. @@ -5392,8 +5382,7 @@ void init_ops(nb::module_& m) { } return t; }, - nb::sig( - "def result_type(*arrays_and_dtypes: Union[array, Dtype]) -> Dtype"), + nb::sig("def result_type(*arrays_and_dtypes: array | Dtype) -> Dtype"), R"pbdoc( The type that results from applying type promotion to the inputs. @@ -5420,7 +5409,7 @@ void init_ops(nb::module_& m) { }, "from_"_a, "to"_a, - nb::sig("def can_cast(from_: Union[array, Dtype], to: Dtype) -> bool"), + nb::sig("def can_cast(from_: array | Dtype, to: Dtype) -> bool"), R"pbdoc( Determine if one data type can be cast to another according to type promotion rules. @@ -5481,7 +5470,7 @@ void init_ops(nb::module_& m) { "dtype"_a, "kind"_a, nb::sig( - "def isdtype(dtype: Dtype, kind: Union[Dtype, str, tuple[Union[Dtype, str], ...]]) -> bool"), + "def isdtype(dtype: Dtype, kind: Dtype | str | tuple[Dtype | str, ...]) -> bool"), R"pbdoc( Test whether a dtype belongs to one or more data type kinds. @@ -5509,7 +5498,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_and(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_and(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise and. @@ -5536,7 +5525,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_or(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_or(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise or. @@ -5563,7 +5552,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_xor(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_xor(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise xor. @@ -5591,7 +5580,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def left_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def left_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise left shift. @@ -5619,7 +5608,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def right_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def right_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise right shift. @@ -5644,7 +5633,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_invert(a: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_invert(a: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise inverse. @@ -5666,7 +5655,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def view(a: Union[scalar, array], dtype: Dtype, stream: Union[None, Stream, Device] = None) -> array"), + "def view(a: scalar | array, dtype: Dtype, stream: StreamOrDevice = None) -> array"), R"pbdoc( View the array as a different type. @@ -5692,7 +5681,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def hadamard_transform(a: array, scale: Optional[float] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def hadamard_transform(a: array, scale: float | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the Walsh-Hadamard transform along the final axis. @@ -5756,7 +5745,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def einsum(subscripts: str, *operands, stream: Union[None, Stream, Device] = None) -> array"), + "def einsum(subscripts: str, *operands, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the Einstein summation convention on the operands. @@ -5791,7 +5780,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def roll(a: array, shift: Union[int, Tuple[int]], axis: Union[None, int, Tuple[int]] = None, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def roll(a: array, shift: int | tuple[int], axis: None | int | tuple[int] = None, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Roll array elements along a given axis. @@ -5819,7 +5808,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def real(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def real(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the real part of a complex array. @@ -5838,7 +5827,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def imag(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def imag(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the imaginary part of a complex array. @@ -5865,7 +5854,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Extract a sub-array from the input array. @@ -5904,7 +5893,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Update a sub-array of the input array. @@ -5933,7 +5922,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Force an array to be row contiguous. Copy if necessary. @@ -5962,7 +5951,7 @@ void init_ops(nb::module_& m) { return nb::tuple(nb::cast(result)); }, - nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"), + nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int]"), R"pbdoc( Broadcast shapes. @@ -6013,7 +6002,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::arg(), nb::sig( - "def depends(inputs: Union[array, Sequence[array]], dependencies: Union[array, Sequence[array]])"), + "def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array])"), R"pbdoc( Insert dependencies between arrays in the graph. The outputs are identical to ``inputs`` but with dependencies on ``dependencies``. @@ -6040,7 +6029,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def qqmm(x: array, w: array, scales: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def qqmm(x: array, w: array, scales: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform a matrix multiplication using a possibly quantized weight matrix ``w`` and a non-quantized input ``x``. The input ``x`` is quantized on the @@ -6090,7 +6079,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"), + "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert the array from fp8 (e4m3) to another floating-point type. @@ -6108,7 +6097,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def to_fp8(x: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert the array to fp8 (e4m3) from another floating-point type. diff --git a/python/src/random.cpp b/python/src/random.cpp index 032c99b9ad..8485faea41 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -149,7 +149,7 @@ void init_random(nb::module_& parent_module) { "num"_a = 2, "stream"_a = nb::none(), nb::sig( - "def split(key: array, num: int = 2, stream: Union[None, Stream, Device] = None) -> array"), + "def split(key: array, num: int = 2, stream: StreamOrDevice = None) -> array"), R"pbdoc( Split a PRNG key into sub keys. @@ -184,7 +184,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def uniform(low: Union[scalar, array] = 0, high: Union[scalar, array] = 1, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def uniform(low: scalar | array = 0, high: scalar | array = 1, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate uniformly distributed random numbers. @@ -227,7 +227,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Union[scalar, array, None] = None, scale: Union[scalar, array, None] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def normal(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: scalar | array | None = None, scale: scalar | array | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate normally distributed random numbers. @@ -267,7 +267,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate jointly-normal random samples given a mean and covariance. @@ -314,7 +314,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def randint(low: Union[scalar, array], high: Union[scalar, array], shape: Sequence[int] = [], dtype: Optional[Dtype] = int32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def randint(low: scalar | array, high: scalar | array, shape: Sequence[int] = [], dtype: Dtype | None = int32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate random integers from the given interval. @@ -357,7 +357,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def bernoulli(p: Union[scalar, array] = 0.5, shape: Optional[Sequence[int]] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def bernoulli(p: scalar | array = 0.5, shape: Sequence[int] | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate Bernoulli random values. @@ -401,7 +401,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def truncated_normal(lower: Union[scalar, array], upper: Union[scalar, array], shape: Optional[Sequence[int]] = None, dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def truncated_normal(lower: scalar | array, upper: scalar | array, shape: Sequence[int] | None = None, dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate values from a truncated normal distribution. @@ -435,7 +435,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def gumbel(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def gumbel(shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample from the standard Gumbel distribution. @@ -481,7 +481,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def categorical(logits: array, axis: int = -1, shape: Optional[Sequence[int]] = None, num_samples: Optional[int] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def categorical(logits: array, axis: int = -1, shape: Sequence[int] | None = None, num_samples: int | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample from a categorical distribution. @@ -524,7 +524,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def laplace(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: float = 0.0, scale: float = 1.0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def laplace(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: float = 0.0, scale: float = 1.0, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample numbers from a Laplace distribution. @@ -557,7 +557,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def permutation(x: Union[int, array], axis: int = 0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def permutation(x: int | array, axis: int = 0, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate a random permutation or permute the entries of an array. diff --git a/python/src/stream.cpp b/python/src/stream.cpp index 76fe427696..004301a45c 100644 --- a/python/src/stream.cpp +++ b/python/src/stream.cpp @@ -88,6 +88,7 @@ void init_stream(nb::module_& m) { "default_stream", &mx::default_stream, "device"_a, + nb::sig("def default_stream(device: Device | DeviceType) -> Stream"), R"pbdoc(Get the device's default stream.)pbdoc"); m.def( "set_default_stream", @@ -106,6 +107,7 @@ void init_stream(nb::module_& m) { "new_stream", &mx::new_stream, "device"_a, + nb::sig("def new_stream(device: Device | DeviceType) -> Stream"), R"pbdoc( Make a new stream on the given device. @@ -116,6 +118,8 @@ void init_stream(nb::module_& m) { "new_thread_unsafe_stream", &mx::new_thread_unsafe_stream, "device"_a, + nb::sig( + "def new_thread_unsafe_stream(device: Device | DeviceType) -> Stream"), R"pbdoc( Make a new stream that can be used in any thread. @@ -128,6 +132,8 @@ void init_stream(nb::module_& m) { "new_thread_local_stream", &mx::new_thread_local_stream, "device"_a, + nb::sig( + "def new_thread_local_stream(device: Device | DeviceType) -> ThreadLocalStream"), R"pbdoc(Make a new stream that will be unique per thread.)pbdoc"); m.def( "clear_streams", diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 6e3fda8882..1d7aa8b9b1 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -1352,7 +1352,7 @@ void init_transforms(nb::module_& m) { "argnums"_a = nb::none(), "argnames"_a = std::vector{}, nb::sig( - "def value_and_grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Tuple[R, Any]]"), + "def value_and_grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, tuple[R, Any]]"), R"pbdoc( Returns a function which computes the value and gradient of ``fun``. @@ -1421,7 +1421,7 @@ void init_transforms(nb::module_& m) { "argnums"_a = nb::none(), "argnames"_a = std::vector{}, nb::sig( - "def grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Any]"), + "def grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, Any]"), R"pbdoc( Returns a function which computes the gradient of ``fun``. @@ -1491,7 +1491,7 @@ void init_transforms(nb::module_& m) { "outputs"_a = nb::none(), "shapeless"_a = false, nb::sig( - "def compile(fun: Callable[P, R], inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable[P, R]"), + "def compile(fun: Callable[P, R], inputs: object | None = None, outputs: object | None = None, shapeless: bool = False) -> Callable[P, R]"), R"pbdoc( Returns a compiled function which produces the same output as ``fun``. From aadcba1a20f81046c63dd64bd1c65cafbf26e9bd Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:55:26 -0700 Subject: [PATCH 109/222] chore: Use dispatch_all_types in DivMod::eval_cpu (#4088) --- mlx/backend/cpu/binary.cpp | 55 ++++++++------------------------------ 1 file changed, 11 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/binary.cpp b/mlx/backend/cpu/binary.cpp index 5c64f03129..7f61316e57 100644 --- a/mlx/backend/cpu/binary.cpp +++ b/mlx/backend/cpu/binary.cpp @@ -3,12 +3,14 @@ #include #include #include +#include #include "mlx/allocator.h" #include "mlx/backend/cpu/binary.h" #include "mlx/backend/cpu/binary_ops.h" #include "mlx/backend/cpu/binary_two.h" #include "mlx/backend/cpu/encoder.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -51,51 +53,16 @@ void DivMod::eval_cpu( return std::make_pair(std::trunc(x / y), std::fmod(x, y)); }; - switch (out_a.dtype()) { - case bool_: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint8: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint16: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint32: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint64: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int8: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int16: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int32: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int64: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case float16: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case float32: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case float64: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case bfloat16: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case complex64: - // Should never get here + dispatch_all_types(out_a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("[DivMod] Complex type not supported"); - break; - } + } else if constexpr (std::is_integral_v) { + binary_op(a, b, out_a, out_b, integral_op, bopt); + } else { + binary_op(a, b, out_a, out_b, float_op, bopt); + } + }); }); } From 62e2508266c102a3186e522a323144bc236558f8 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:55:58 -0700 Subject: [PATCH 110/222] chore: Use dispatch_all_types in SegmentedMM::eval_cpu (#4087) --- mlx/backend/cpu/masked_mm.cpp | 98 ++++++++--------------------------- 1 file changed, 21 insertions(+), 77 deletions(-) diff --git a/mlx/backend/cpu/masked_mm.cpp b/mlx/backend/cpu/masked_mm.cpp index 688479c602..db38f39424 100644 --- a/mlx/backend/cpu/masked_mm.cpp +++ b/mlx/backend/cpu/masked_mm.cpp @@ -8,6 +8,7 @@ #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/gemm.h" #include "mlx/backend/cpu/lapack.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -525,83 +526,26 @@ void SegmentedMM::eval_cpu(const std::vector& inputs, array& out) { b_transposed = b_transposed, lda = lda, ldb = ldb]() { - switch (a.dtype()) { - case float64: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case float32: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case float16: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case bfloat16: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - default: - throw std::invalid_argument( - "Segmented mm supports only real float types."); - } + dispatch_float_types( + a.dtype(), "[SegmentedMM::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + segmented_mm( + a.data(), + b.data(), + segments.data(), + static_cast(out_ptr), + a_transposed, + b_transposed, + lda, + ldb, + a.shape(), + a.strides(), + b.shape(), + b.strides(), + segments.size() / 2, + segments.shape(), + segments.strides()); + }); }); } From 0c53491b0ad7e7f3643c5512a774866125d804c7 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:56:29 -0700 Subject: [PATCH 111/222] chore: Use dispatch_all_types in Arange::eval_cpu (#4085) --- mlx/backend/cpu/primitives.cpp | 52 ++++++---------------------------- 1 file changed, 9 insertions(+), 43 deletions(-) diff --git a/mlx/backend/cpu/primitives.cpp b/mlx/backend/cpu/primitives.cpp index f1d83dd306..3299b74717 100644 --- a/mlx/backend/cpu/primitives.cpp +++ b/mlx/backend/cpu/primitives.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/allocator.h" #include "mlx/backend/common/slicing.h" @@ -13,6 +14,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/threefry.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -125,50 +127,14 @@ void Transpose::eval_cpu(const std::vector& inputs, array& out) { void Arange::eval_cpu(const std::vector& inputs, array& out) { assert(inputs.size() == 0); out.set_data(allocator::malloc(out.nbytes())); - switch (out.dtype()) { - case bool_: + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("Bool type unsupported for arange."); - break; - case uint8: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int8: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case bfloat16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case complex64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - } + } else { + arange(start_, start_ + step_, out, out.size(), stream()); + } + }); } void AsType::eval_cpu(const std::vector& inputs, array& out) { From e2aa0d096a50ea2d08089613c4f53878264b22a3 Mon Sep 17 00:00:00 2001 From: Jeremy Gu <145739220+wgu9@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:56:49 -0700 Subject: [PATCH 112/222] Fix Device and Stream lexicographic ordering (#4086) --- mlx/device.h | 3 ++- mlx/stream.h | 3 ++- tests/device_tests.cpp | 45 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/mlx/device.h b/mlx/device.h index 9ddd1b3f5b..ff75fe0cca 100644 --- a/mlx/device.h +++ b/mlx/device.h @@ -5,6 +5,7 @@ #include "mlx/api.h" #include +#include #include #include @@ -27,7 +28,7 @@ struct MLX_API Device { // TODO: Use default three-way comparison when it gets supported in XCode. bool operator==(const Device&) const = default; bool operator<(const Device& rhs) const { - return type < rhs.type || index < rhs.index; + return std::tie(type, index) < std::tie(rhs.type, rhs.index); } }; diff --git a/mlx/stream.h b/mlx/stream.h index 8243369876..3099ee2682 100644 --- a/mlx/stream.h +++ b/mlx/stream.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include "mlx/api.h" @@ -17,7 +18,7 @@ struct MLX_API Stream { // TODO: Use default three-way comparison when it gets supported in XCode. bool operator==(const Stream&) const = default; bool operator<(const Stream& rhs) const { - return device < rhs.device || index < rhs.index; + return std::tie(device, index) < std::tie(rhs.device, rhs.index); } }; diff --git a/tests/device_tests.cpp b/tests/device_tests.cpp index 7bae91d88d..b9e26aed21 100644 --- a/tests/device_tests.cpp +++ b/tests/device_tests.cpp @@ -2,12 +2,57 @@ #include "doctest/doctest.h" +#include +#include #include +#include #include "mlx/mlx.h" using namespace mlx::core; +template +void check_strict_weak_ordering(const std::array& ordered_values) { + for (const auto& value : ordered_values) { + CHECK_FALSE(value < value); + } + + for (const auto& lhs : ordered_values) { + for (const auto& rhs : ordered_values) { + if (lhs < rhs) { + CHECK_FALSE(rhs < lhs); + } + for (const auto& other : ordered_values) { + if (lhs < rhs && rhs < other) { + CHECK(lhs < other); + } + } + } + } + + std::set values(ordered_values.begin(), ordered_values.end()); + REQUIRE_EQ(values.size(), ordered_values.size()); + CHECK(std::equal(values.begin(), values.end(), ordered_values.begin())); +} + +TEST_CASE("test device and stream ordering") { + const std::array devices{ + Device(Device::cpu, 0), + Device(Device::cpu, 3), + Device(Device::gpu, 0), + Device(Device::gpu, 1), + Device(Device::gpu, 2)}; + check_strict_weak_ordering(devices); + + const std::array streams{ + Stream(0, Device::cpu), + Stream(3, Device::cpu), + Stream(0, Device::gpu), + Stream(1, Device::gpu), + Stream(2, Device::gpu)}; + check_strict_weak_ordering(streams); +} + TEST_CASE("test device placement") { auto device = default_device(); Device d = gpu::is_available() ? Device::gpu : Device::cpu; From b0c35627e85b237c3b653f55ee2d2c0ea28a4175 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sat, 8 Aug 2026 23:28:52 -0700 Subject: [PATCH 113/222] Enable complex64 scatter addition on GPU (#4078) --- mlx/backend/metal/indexing.cpp | 3 ++- mlx/backend/metal/kernels/reduction/ops.h | 9 +++++++ mlx/ops.cpp | 3 ++- python/tests/test_array.py | 31 +++++++++++++++++++++++ 4 files changed, 44 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/indexing.cpp b/mlx/backend/metal/indexing.cpp index a3df0e6e28..d562fa6f22 100644 --- a/mlx/backend/metal/indexing.cpp +++ b/mlx/backend/metal/indexing.cpp @@ -230,7 +230,8 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { } void Scatter::eval_gpu(const std::vector& inputs, array& out) { - if (size_of(out.dtype()) == 8) { + if (size_of(out.dtype()) == 8 && + !(out.dtype() == complex64 && reduce_type_ == Scatter::Sum)) { std::ostringstream msg; msg << "[Scatter::eval_gpu] Does not support " << out.dtype(); throw std::invalid_argument(msg.str()); diff --git a/mlx/backend/metal/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h index b7a9cacb39..f875e51c92 100644 --- a/mlx/backend/metal/kernels/reduction/ops.h +++ b/mlx/backend/metal/kernels/reduction/ops.h @@ -133,6 +133,15 @@ struct Sum { mlx_atomic_fetch_add_explicit(out, val, offset); } + void atomic_update( + device mlx_atomic* out, + complex64_t val, + size_t offset = 0) thread { + auto out_lanes = reinterpret_cast*>(out); + mlx_atomic_fetch_add_explicit(out_lanes, val.real, 2 * offset); + mlx_atomic_fetch_add_explicit(out_lanes, val.imag, 2 * offset + 1); + } + // Operator U operator()(U a, U b) thread { return a + b; diff --git a/mlx/ops.cpp b/mlx/ops.cpp index b7b9a84291..f95dede289 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3794,7 +3794,8 @@ array scatter( } // TODO, remove when scatter supports 64-bit outputs - if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8) { + if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8 && + !(a.dtype() == complex64 && mode == Scatter::Sum)) { std::ostringstream msg; msg << "[scatter] GPU scatter does not yet support " << a.dtype() << " for the input or updates."; diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 75bab7da1f..5c6db31482 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1548,6 +1548,37 @@ def test_array_at(self): a = a.at[1:3, :, 0].minimum(update) self.assertEqualArray(a[1:3, :, 0], mx.minimum(a[1:3, :, 0], update)) + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") + def test_array_at_complex_add_gpu(self): + n = 4096 + base = [1 + 10j, 2 + 20j, 3 + 30j, 4 + 40j] + + with mx.stream(mx.gpu): + a = mx.array(base, dtype=mx.complex64) + update_indices = mx.full((n,), 3, dtype=mx.int32) + updates = mx.full((n,), 1 + 3j, dtype=mx.complex64) + out = a.at[update_indices].add(updates) + mx.eval(out) + + indices = mx.array([1, 1, 3]) + x = mx.array([1 + 0j, 3 + 4j, 6 + 8j, 5 + 12j], dtype=mx.complex64) + + def loss(z): + return mx.square(mx.abs(z[indices])).sum() + + _, gradient = mx.value_and_grad(loss)(x) + mx.eval(gradient) + + expected = base.copy() + expected[-1] += n * (1 + 3j) + self.assertEqual(out.tolist(), expected) + np.testing.assert_allclose( + np.array(gradient), + np.array([0, 12 + 16j, 0, 10 + 24j], dtype=np.complex64), + rtol=0, + atol=1e-5, + ) + def test_array_at_slice_update_extensive(self): # Test with transposed inputs a = mx.zeros((4, 5)) From 5a1e44c3bb991dab753cee394b0b1d889e2eb9a7 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Sun, 9 Aug 2026 00:10:06 -0700 Subject: [PATCH 114/222] Optimize large NVFP4 QMV on M5 Max (#3961) Co-authored-by: Cheng --- mlx/backend/metal/kernels/fp_quantized.h | 20 +++++++--- mlx/backend/metal/kernels/fp_quantized.metal | 41 ++++++++++++++++---- mlx/backend/metal/quantized.cpp | 16 ++++++-- python/tests/test_quantized.py | 37 +++++++++++++++++- 4 files changed, 96 insertions(+), 18 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index f3aa24bf7a..6e77569f56 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -321,7 +321,12 @@ METAL_FUNC void fp_qmv_quad_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( const device uint32_t* w, const device uint8_t* scales, @@ -335,7 +340,6 @@ METAL_FUNC void fp_qmv_fast_impl( uint simd_lid [[thread_index_in_simdgroup]]) { constexpr int packs_per_thread = 2; constexpr int num_simdgroups = 2; - constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; @@ -1167,7 +1171,8 @@ template < int group_size, int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( const device uint32_t* w, const device uint8_t* scales, @@ -1203,7 +1208,12 @@ template < s_strides, tid); } - fp_qmv_fast_impl( + fp_qmv_fast_impl< + T, + group_size, + bits, + has_global_scale, + results_per_simdgroup>( w, scales, global_scale, diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 7404f2023f..d8d462288a 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" @@ -57,6 +57,30 @@ aligned, \ batched) +#define instantiate_quantized_qmv_fast(mode, type, results, batched, group_size, bits) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched, \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + false, \ + results) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched "_hgs", \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + true, \ + results) + +#define instantiate_quantized_qmv_fast_r2(mode, type, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 1, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 0, group_size, bits) + #define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \ @@ -185,13 +209,14 @@ instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_modes(type, nvfp4, 16, 4) \ - instantiate_quantized_modes(type, mxfp8, 32, 8) \ - instantiate_quantized_modes(type, mxfp4, 32, 4) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ - instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ - instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ + instantiate_quantized_modes(type, nvfp4, 16, 4) \ + instantiate_quantized_modes(type, mxfp8, 32, 8) \ + instantiate_quantized_modes(type, mxfp4, 32, 4) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ + instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ + instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ + instantiate_quantized_qmv_fast_r2(nvfp4, type, 16, 4) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 7c461bc1b5..cb42e817d5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/backend/common/quantized.h" #include "mlx/backend/common/broadcasting.h" @@ -477,13 +477,19 @@ void qmv( int bn = 8; int bk = 32; - MTL::Size group_dims(bk, 2, 1); - MTL::Size grid_dims(M, (N + bn - 1) / bn, B); std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + // A narrower output tile reduces register pressure for large + // floating-point quantized matrix-vector products on M5 Max GPUs. + bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && + d.get_architecture().back() == 's' && mode == "nvfp4"; + int results_per_simdgroup = use_narrow_qmv ? 2 : 4; + bn = 2 * results_per_simdgroup; + MTL::Size group_dims(bk, 2, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); concatenate( kname, @@ -493,6 +499,7 @@ void qmv( group_size, "_b_", bits, + use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( @@ -504,7 +511,8 @@ void qmv( group_size, bits, B > 1, - global_scale.has_value()); + global_scale.has_value(), + results_per_simdgroup); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 0b756b39d5..806784cfc7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,4 +1,4 @@ -# Copyright © 2023 Apple Inc. +# Copyright © 2023-2026 Apple Inc. import platform import subprocess @@ -516,6 +516,41 @@ def test_fp_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_fp_qmv_large_output(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + K = 512 + N = 4096 + + for B in [1, 2]: + with self.subTest(B=B, N=N, K=K): + x_shape = (1, K) if B == 1 else (B, 1, K) + w_shape = (N, K) if B == 1 else (B, N, K) + x = mx.random.normal(shape=x_shape, key=k1) / K**0.5 + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode="nvfp4") + + dtypes = ( + [mx.float16, mx.bfloat16, mx.float32] + if mx.default_device() == mx.gpu + else [mx.float32] + ) + for dtype in dtypes: + with self.subTest(dtype=dtype): + x_t = x.astype(dtype) + w_hat = mx.dequantize(w_q, scales, mode="nvfp4", dtype=dtype) + y_q = mx.quantized_matmul( + x_t, + w_q, + scales, + transpose=True, + mode="nvfp4", + ) + y_hat = x_t @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(y_q, y_hat, rtol=tol, atol=tol)) + def test_qmv_wide(self): # M in [2, vector_limit) routes to qmv_wide -- except K in {64, 128} # with power-of-2 bits, which stays on qmv_quad. Check both paths From 8c28c385f86d17e1da427bf8d81afe084ee17c35 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Sun, 9 Aug 2026 04:04:14 -0400 Subject: [PATCH 115/222] Fix ring hanging on peer disconnect (#4060) Co-authored-by: Cheng --- mlx/distributed/ring/ring.cpp | 45 ++++++++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 11 deletions(-) diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index ea40042844..a74122239a 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -238,6 +239,9 @@ class SocketThread { task.size -= r; delete_recv = task.size == 0; error_count = 0; + } else if (r == 0) { + error_count++; + log_info(true, "Socket", fd_, "was closed by the peer"); } else if (errno != EAGAIN) { error_count++; log_info( @@ -252,6 +256,9 @@ class SocketThread { task.size -= r; delete_send = task.size == 0; error_count = 0; + } else if (r == 0) { + error_count++; + log_info(true, "Sending to socket", fd_, "made no progress"); } else if (errno != EAGAIN) { error_count++; log_info(true, "Sending to socket", fd_, "failed with errno", errno); @@ -259,7 +266,18 @@ class SocketThread { } if (error_count >= 10) { - log_info(true, "Too many send/recv errors. Aborting..."); + log_info(true, "Too many send/recv errors. Failing pending tasks..."); + // Throw exception in invoker's thread. + auto error = std::make_exception_ptr( + std::runtime_error("[ring] connection to a peer was lost")); + for (auto& task : recvs_) { + task.promise.set_exception(error); + } + for (auto& task : sends_) { + task.promise.set_exception(error); + } + recvs_.clear(); + sends_.clear(); return; } } @@ -525,7 +543,7 @@ class RingGroup : public GroupImpl { (i % 2) ? -1 : 1))); } for (auto& f : all_gathers) { - f.wait(); + f.get(); } }); } @@ -650,7 +668,7 @@ class RingGroup : public GroupImpl { reduce_op))); } for (auto& f : all_sums) { - f.wait(); + f.get(); } }); } @@ -737,8 +755,8 @@ class RingGroup : public GroupImpl { } if (j >= 0) { - sends[b].wait(); - recvs[b].wait(); + sends[b].get(); + recvs[b].get(); if (2 * j < send_plan.size()) { reduce_op( recv_buffers[j % ALL_SUM_BUFFERS], @@ -749,8 +767,13 @@ class RingGroup : public GroupImpl { std::swap(a, b); } - sends[b].wait(); - recvs[b].wait(); + // Check valid() to avoid consuming same future twice for single packet. + if (sends[b].valid()) { + sends[b].get(); + } + if (recvs[b].valid()) { + recvs[b].get(); + } } void all_gather_impl( @@ -784,8 +807,8 @@ class RingGroup : public GroupImpl { send_segment = (send_segment + size_ + direction) % size_; recv_segment = (recv_segment + size_ + direction) % size_; - sent.wait(); - recvd.wait(); + sent.get(); + recvd.get(); } } @@ -804,7 +827,7 @@ class RingGroup : public GroupImpl { std::min(data_size, (i + 1) * segment_size) - i * segment_size)); } for (auto& f : sends) { - f.wait(); + f.get(); } } @@ -822,7 +845,7 @@ class RingGroup : public GroupImpl { std::min(data_size, (i + 1) * segment_size) - i * segment_size)); } for (auto& f : recvs) { - f.wait(); + f.get(); } } From 7eb3d85411fc902cde936b0dc1e13bd808c6b969 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:37:39 +0800 Subject: [PATCH 116/222] Propagate NaN in cummax and curmin (#4047) Co-authored-by: Cheng --- mlx/backend/cpu/scan.cpp | 18 ++++++++++++++-- mlx/backend/cpu/simd/base_simd.h | 13 ------------ mlx/backend/cpu/sort.cpp | 10 +++------ mlx/backend/metal/kernels/scan.h | 26 +++++++++++++++++++---- mlx/types/half_types.h | 25 ++++++++++++++++++---- python/tests/test_ops.py | 36 ++++++++++++++++++++++++++++++++ 6 files changed, 98 insertions(+), 30 deletions(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 4dc2f50ead..3ebbe0a3c3 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -212,7 +212,14 @@ void scan_dispatch( break; } case Scan::Min: { - auto op = [](U y, T x) { return x < y ? x : y; }; + auto op = [](U y, T x) { + if constexpr (is_floating_point_v) { + if (std::isnan(y) || std::isnan(static_cast(x))) { + return std::numeric_limits::quiet_NaN(); + } + } + return x < y ? x : y; + }; auto init = (issubdtype(in.dtype(), floating)) ? static_cast(std::numeric_limits::infinity()) : std::numeric_limits::max(); @@ -220,7 +227,14 @@ void scan_dispatch( break; } case Scan::Max: { - auto op = [](U y, T x) { return x < y ? y : x; }; + auto op = [](U y, T x) { + if constexpr (is_floating_point_v) { + if (std::isnan(y) || std::isnan(static_cast(x))) { + return std::numeric_limits::quiet_NaN(); + } + } + return x < y ? y : x; + }; auto init = (issubdtype(in.dtype(), floating)) ? static_cast(-std::numeric_limits::infinity()) : std::numeric_limits::min(); diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index be10a89ee4..6a1ee39a59 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -56,19 +56,6 @@ void store(T* dst, Simd x) { *(Simd*)dst = x; } -template -constexpr bool is_complex = false; - -template -constexpr bool is_complex().real())>> = - true; - -// std::is_signed_v is false for the custom float16_t/bfloat16_t types, so it -// skips the floored-mod sign correction for them. -template -inline constexpr bool is_signed_v = std::is_signed_v || - std::is_same_v || std::is_same_v; - template Simd rint(Simd in) { if constexpr (is_complex) { diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index 090cbb7a87..233900432a 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -15,14 +15,10 @@ namespace mlx::core { namespace { -template -inline constexpr bool is_floating_v = std::is_floating_point_v || - std::is_same_v || std::is_same_v; - // NaN-aware comparator that places NaNs at the end template bool nan_aware_less(T a, T b) { - if constexpr (is_floating_v || std::is_same_v) { + if constexpr (is_floating_point_v || std::is_same_v) { if (std::isnan(a)) return false; if (std::isnan(b)) @@ -202,7 +198,7 @@ void argsort(const array& in, array& out, int axis) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_v) { + if constexpr (is_floating_point_v) { if (std::isnan(v1)) return false; if (std::isnan(v2)) @@ -303,7 +299,7 @@ void argpartition(const array& in, array& out, int axis, int kth) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_v) { + if constexpr (is_floating_point_v) { if (std::isnan(v1)) return false; if (std::isnan(v2)) diff --git a/mlx/backend/metal/kernels/scan.h b/mlx/backend/metal/kernels/scan.h index a6bfde0018..2e5493961b 100644 --- a/mlx/backend/metal/kernels/scan.h +++ b/mlx/backend/metal/kernels/scan.h @@ -99,15 +99,24 @@ template struct CumMax { static constexpr constant U init = Limits::min; + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } + return (a >= b) ? a : b; + } + template U operator()(U a, T b) thread { - return (a >= b) ? a : b; + return combine(a, static_cast(b)); } U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x >= other) ? x : other; + x = combine(x, other); } return x; } @@ -122,15 +131,24 @@ template struct CumMin { static constexpr constant U init = Limits::max; + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } + return (a <= b) ? a : b; + } + template U operator()(U a, T b) thread { - return (a <= b) ? a : b; + return combine(a, static_cast(b)); } U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x <= other) ? x : other; + x = combine(x, other); } return x; } diff --git a/mlx/types/half_types.h b/mlx/types/half_types.h index d9d6b9bf55..560f737359 100644 --- a/mlx/types/half_types.h +++ b/mlx/types/half_types.h @@ -36,10 +36,11 @@ typedef struct _MLX_BFloat16 bfloat16_t; #endif // __ARM_FEATURE_BF16 -#ifdef ADD_HALF_BINOPS namespace mlx::core { - // clang-format off + +#ifdef ADD_HALF_BINOPS + #define fp16_bf16_binop_helper(__op__, __operator__) \ inline float __operator__(float16_t lhs, bfloat16_t rhs) { \ return static_cast(lhs) __op__ static_cast(rhs); \ @@ -52,7 +53,23 @@ fp16_bf16_binop_helper(+, operator+) fp16_bf16_binop_helper(-, operator-) fp16_bf16_binop_helper(*, operator*) fp16_bf16_binop_helper(/, operator/) -// clang-format on -} // namespace mlx::core #endif + +template +constexpr bool is_complex = false; + +template +constexpr bool is_complex().real())>> = + true; + +template +constexpr bool is_signed_v = std::is_signed_v || + std::is_same_v || std::is_same_v; + +template +constexpr bool is_floating_point_v = std::is_floating_point_v || + std::is_same_v || std::is_same_v; + +// clang-format on +} // namespace mlx::core diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index e8555b8747..b6f9ef6ebf 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2371,6 +2371,42 @@ def fn(its): mem4 = mx.get_peak_memory() self.assertEqual(mem2, mem4) + def test_cummax_cummin_nan(self): + nan = float("nan") + cases = [ + [1.0, 3.0, nan, 5.0, 4.0], + [nan, 3.0, 2.0, 5.0, 4.0], + [1.0, 2.0, 3.0, nan, 4.0], + [nan, nan, 1.0], + [5.0, 4.0, nan, 1.0, 0.0], + ] + for op, npop, init in ( + ("cummax", np.maximum, float("-inf")), + ("cummin", np.minimum, float("inf")), + ): + for arr in cases: + a_np = np.array(arr, dtype=np.float32) + a_mx = mx.array(a_np) + inc_fwd = npop.accumulate(a_np) + inc_rev = npop.accumulate(a_np[::-1])[::-1] + exc_fwd = np.concatenate([[init], inc_fwd[:-1]]) + exc_rev = np.concatenate([inc_rev[1:], [init]]) + refs = { + (False, True): inc_fwd, + (True, True): inc_rev, + (False, False): exc_fwd, + (True, False): exc_rev, + } + for (reverse, inclusive), expected in refs.items(): + got = np.array( + getattr(mx, op)(a_mx, reverse=reverse, inclusive=inclusive) + ) + self.assertTrue( + np.array_equal(got, expected, equal_nan=True), + msg=f"{op} reverse={reverse} inclusive={inclusive} " + f"arr={arr}\ngot={got}\nexp={expected}", + ) + def test_diff(self): a = mx.array([1, 2, 4, 7, 0]) self.assertEqual(mx.diff(a).tolist(), [1, 2, 3, -7]) From 313ae182615ba379c4f07c015b1e1d04d90597f0 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:27:00 -0700 Subject: [PATCH 117/222] chore: Use dispatch_inexact_types in Equal::eval_cpu (#4095) --- mlx/backend/cpu/binary.cpp | 25 +++++-------------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/mlx/backend/cpu/binary.cpp b/mlx/backend/cpu/binary.cpp index 7f61316e57..9cca16d869 100644 --- a/mlx/backend/cpu/binary.cpp +++ b/mlx/backend/cpu/binary.cpp @@ -96,26 +96,11 @@ void Equal::eval_cpu(const std::vector& inputs, array& out) { b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (a.dtype()) { - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - default: - throw std::runtime_error( - "[NanEqual::eval_cpu] Only for floating point types."); - } + dispatch_inexact_types( + a.dtype(), "[NanEqual::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } else { comparison_op_cpu(a, b, out, detail::Equal(), stream()); From 78c5b2870ec2e562902cb25fb3f03b865827b666 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:28:12 -0700 Subject: [PATCH 118/222] chore: Use dispatch_inexact_types in Matmul::eval_cpu (#4103) --- mlx/backend/cpu/matmul.cpp | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/mlx/backend/cpu/matmul.cpp b/mlx/backend/cpu/matmul.cpp index 7df331671e..12781dc1ab 100644 --- a/mlx/backend/cpu/matmul.cpp +++ b/mlx/backend/cpu/matmul.cpp @@ -7,6 +7,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/gemm.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -97,24 +98,11 @@ void matmul_general( return; } - if (out.dtype() == float32) { - matmul_dispatch( + dispatch_inexact_types(out.dtype(), "[Matmul::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + matmul_dispatch( a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == float16) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == bfloat16) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == float64) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == complex64) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else { - throw std::runtime_error("[Matmul::eval_cpu] Invalid type."); - } + }); cpu::get_command_encoder(stream).add_temporaries(std::move(temps)); } From bca2a79934481876451605251aaf5930b7d6813e Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:29:17 -0700 Subject: [PATCH 119/222] chore: Use dispatch_all_types in Gather::eval_cpu (#4105) --- mlx/backend/cpu/indexing.cpp | 48 +++--------------------------------- 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/indexing.cpp b/mlx/backend/cpu/indexing.cpp index d668b56adb..54dbfdd7af 100644 --- a/mlx/backend/cpu/indexing.cpp +++ b/mlx/backend/cpu/indexing.cpp @@ -128,50 +128,10 @@ void dispatch_gather( array& out, const std::vector& axes, const Shape& size) { - switch (out.dtype()) { - case bool_: - gather(src, inds, out, axes, size); - break; - case uint8: - gather(src, inds, out, axes, size); - break; - case uint16: - gather(src, inds, out, axes, size); - break; - case uint32: - gather(src, inds, out, axes, size); - break; - case uint64: - gather(src, inds, out, axes, size); - break; - case int8: - gather(src, inds, out, axes, size); - break; - case int16: - gather(src, inds, out, axes, size); - break; - case int32: - gather(src, inds, out, axes, size); - break; - case int64: - gather(src, inds, out, axes, size); - break; - case float16: - gather(src, inds, out, axes, size); - break; - case float32: - gather(src, inds, out, axes, size); - break; - case float64: - gather(src, inds, out, axes, size); - break; - case bfloat16: - gather(src, inds, out, axes, size); - break; - case complex64: - gather(src, inds, out, axes, size); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + gather(src, inds, out, axes, size); + }); } void Gather::eval_cpu(const std::vector& inputs, array& out) { From d64537c0820dd132bdfd97f6140922ef3544eb52 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:31:23 -0700 Subject: [PATCH 120/222] chore: Use dispatch_all_types in CPU unary (#4106) --- mlx/backend/cpu/unary.h | 48 ++++------------------------------------- 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/unary.h b/mlx/backend/cpu/unary.h index 4fab6a7544..e4c9d47b90 100644 --- a/mlx/backend/cpu/unary.h +++ b/mlx/backend/cpu/unary.h @@ -5,6 +5,7 @@ #include "mlx/backend/common/unary.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" #include "mlx/utils.h" namespace mlx::core { @@ -61,50 +62,9 @@ void unary(const array& a, array& out, Op op, Stream stream) { encoder.dispatch([a = array::unsafe_weak_copy(a), out = array::unsafe_weak_copy(out), op = op]() mutable { - switch (out.dtype()) { - case bool_: - unary_op(a, out, op); - break; - case uint8: - unary_op(a, out, op); - break; - case uint16: - unary_op(a, out, op); - break; - case uint32: - unary_op(a, out, op); - break; - case uint64: - unary_op(a, out, op); - break; - case int8: - unary_op(a, out, op); - break; - case int16: - unary_op(a, out, op); - break; - case int32: - unary_op(a, out, op); - break; - case int64: - unary_op(a, out, op); - break; - case float16: - unary_op(a, out, op); - break; - case float32: - unary_op(a, out, op); - break; - case float64: - unary_op(a, out, op); - break; - case bfloat16: - unary_op(a, out, op); - break; - case complex64: - unary_op(a, out, op); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + unary_op(a, out, op); + }); }); } From b00b34bd9c6fd80c1b90aec3cff977b2b863479a Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:32:26 -0700 Subject: [PATCH 121/222] chore: Use dispatch_all_types in binary_op_cpu (#4107) --- mlx/backend/cpu/binary.h | 49 ++++------------------------------------ 1 file changed, 5 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/binary.h b/mlx/backend/cpu/binary.h index acaca50e1b..caf54f21f0 100644 --- a/mlx/backend/cpu/binary.h +++ b/mlx/backend/cpu/binary.h @@ -6,6 +6,7 @@ #include "mlx/array.h" #include "mlx/backend/common/binary.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" @@ -309,50 +310,10 @@ void binary_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (out.dtype()) { - case bool_: - binary_op(a, b, out, bopt); - break; - case uint8: - binary_op(a, b, out, bopt); - break; - case uint16: - binary_op(a, b, out, bopt); - break; - case uint32: - binary_op(a, b, out, bopt); - break; - case uint64: - binary_op(a, b, out, bopt); - break; - case int8: - binary_op(a, b, out, bopt); - break; - case int16: - binary_op(a, b, out, bopt); - break; - case int32: - binary_op(a, b, out, bopt); - break; - case int64: - binary_op(a, b, out, bopt); - break; - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } From 5b59a9cc633bc1f5c691ac1157a81ee4fe9b6f34 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 9 Aug 2026 17:35:50 -0700 Subject: [PATCH 122/222] chore: Validate pooling kernel size, stride and padding (#4101) --- python/mlx/nn/layers/pooling.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/mlx/nn/layers/pooling.py b/python/mlx/nn/layers/pooling.py index 2031176346..eaf50228ee 100644 --- a/python/mlx/nn/layers/pooling.py +++ b/python/mlx/nn/layers/pooling.py @@ -85,6 +85,19 @@ class _Pool(Module): def __init__(self, pooling_function, kernel_size, stride, padding, padding_value): super().__init__() + class_name = type(self).__name__ + for name, values in (("kernel_size", kernel_size), ("stride", stride)): + if any(v <= 0 for v in values): + raise ValueError( + f"[{class_name}] '{name}' must be positive but got " + f"{tuple(values)}." + ) + if any(p[0] < 0 for p in padding): + raise ValueError( + f"[{class_name}] 'padding' must be non-negative but got " + f"{tuple(p[0] for p in padding)}." + ) + self._pooling_function = pooling_function self._kernel_size = kernel_size self._stride = stride From 40e9b57b0ee9c5ae3c38f7646fe2a37ac76eadec Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 9 Aug 2026 17:45:19 -0700 Subject: [PATCH 123/222] chore: Fix reducing 0-size array with non-empty axis in min/max (#4079) Co-authored-by: Cheng --- mlx/backend/cuda/reduce.cu | 2 +- mlx/backend/cuda/reduce/init_reduce.cu | 4 ++++ mlx/backend/metal/reduce.cpp | 5 ++++- mlx/ops.cpp | 23 ++++++++++++++--------- python/tests/test_reduce.py | 26 ++++++++++++++++++++++++++ 5 files changed, 49 insertions(+), 11 deletions(-) diff --git a/mlx/backend/cuda/reduce.cu b/mlx/backend/cuda/reduce.cu index 1769cdf2f8..833fdbf669 100644 --- a/mlx/backend/cuda/reduce.cu +++ b/mlx/backend/cuda/reduce.cu @@ -23,7 +23,7 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { // When all the reduced axes have size 1 at runtime, which can happen with // shapeless compilation, the reduction is the identity so just cast-copy // the input to the output. - if (out.size() == in.size()) { + if (in.size() > 0 && out.size() == in.size()) { CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General; copy_gpu(in, out, ctype, s); diff --git a/mlx/backend/cuda/reduce/init_reduce.cu b/mlx/backend/cuda/reduce/init_reduce.cu index e2d5dfa023..463d5d5a98 100644 --- a/mlx/backend/cuda/reduce/init_reduce.cu +++ b/mlx/backend/cuda/reduce/init_reduce.cu @@ -31,6 +31,10 @@ void init_reduce( out.set_data(cu::malloc_async(out.nbytes(), encoder)); } + if (out.size() == 0) { + return; + } + encoder.set_output_array(out); dispatch_all_types(in.dtype(), [&](auto type_tag) { dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) { diff --git a/mlx/backend/metal/reduce.cpp b/mlx/backend/metal/reduce.cpp index 562f966297..52d15288df 100644 --- a/mlx/backend/metal/reduce.cpp +++ b/mlx/backend/metal/reduce.cpp @@ -292,6 +292,9 @@ void init_reduce( CommandEncoder& compute_encoder, metal::Device& d, const Stream& s) { + if (out.size() == 0) { + return; + } auto [_, out_type] = remap_reduce_types(out, op_name); const std::string func_name = "init_reduce"; std::string kname = func_name; @@ -957,7 +960,7 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { // When all the reduced axes have size 1 at runtime, which can happen with // shapeless compilation, the reduction is the identity so just cast-copy // the input to the output. - if (out.size() == in.size()) { + if (in.size() > 0 && out.size() == in.size()) { CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General; copy_gpu(in, out, ctype, stream()); diff --git a/mlx/ops.cpp b/mlx/ops.cpp index f95dede289..d00130ee64 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -22,9 +22,14 @@ namespace mlx::core { namespace { +// Pass the reduction's name in `op_without_identity` when it has no identity +// element, as max and min do not. Those are undefined over an empty axis, so +// naming them here has that checked. Reductions with an identity, such as sum +// and any, reduce an empty axis fine and leave it unset. std::tuple, bool> compute_reduce_shape( const std::vector& axes, - const Shape& shape) { + const Shape& shape, + std::string_view op_without_identity = {}) { bool is_noop = true; std::set axes_set; auto ndim = shape.size(); @@ -46,6 +51,12 @@ std::tuple, bool> compute_reduce_shape( if (axes_set.count(i) == 0) { out_shape.push_back(shape[i]); } else { + if (!op_without_identity.empty() && shape[i] == 0) { + std::ostringstream msg; + msg << "[" << op_without_identity << "] Cannot " << op_without_identity + << " reduce over axis " << i << " with size 0."; + throw std::invalid_argument(msg.str()); + } out_shape.push_back(1); } is_noop &= (out_shape.back() == shape[i]); @@ -2455,11 +2466,8 @@ array max( const std::vector& axes, bool keepdims /* = false */, StreamOrDevice s /* = {}*/) { - if (a.size() == 0) { - throw std::invalid_argument("[max] Cannot max reduce zero size array."); - } auto [out_shape, sorted_axes, is_noop] = - compute_reduce_shape(axes, a.shape()); + compute_reduce_shape(axes, a.shape(), "max"); auto out = (is_noop) ? a : array( @@ -2492,14 +2500,11 @@ array min( const std::vector& axes, bool keepdims /* = false */, StreamOrDevice s /* = {}*/) { - if (a.size() == 0) { - throw std::invalid_argument("[min] Cannot min reduce zero size array."); - } if (axes.empty()) { return a; } auto [out_shape, sorted_axes, is_noop] = - compute_reduce_shape(axes, a.shape()); + compute_reduce_shape(axes, a.shape(), "min"); auto out = (is_noop) ? a : array( diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 24d5a688db..4f37ebbaf8 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -124,6 +124,32 @@ def test_edge_case(self): z = np.array(x).sum((0, 2, 3)) self.assertTrue(np.all(z == y)) + def test_zero_size(self): + # max and min have no identity, so they are only undefined when an axis + # being reduced is itself empty. An array that is empty because of some + # other axis reduces into an empty result. + for shape, axis in [ + ((0, 2), -1), + ((2, 0), 0), + ((3, 0, 2), -1), + ((0, 3, 2), (1, 2)), + ]: + a_np = np.zeros(shape, dtype=np.float32) + a_mx = mx.array(a_np) + for op in ["max", "min"]: + out = getattr(mx, op)( + a_mx, axis=list(axis) if isinstance(axis, tuple) else axis + ) + mx.eval(out) + self.assertEqual(out.shape, getattr(np, op)(a_np, axis=axis).shape) + + # Reducing an empty axis still raises, like numpy + for shape, axis in [((2, 0), -1), ((0, 2), 0), ((0, 0), -1)]: + a_mx = mx.zeros(shape) + for op in ["max", "min"]: + with self.assertRaises(ValueError): + getattr(mx, op)(a_mx, axis=axis) + def test_sum_bool(self): x = np.random.uniform(0, 1, size=(10, 10, 10)) > 0.5 y = mx.array(x) From bcfced932228c19081d90abbf7529acde7c31441 Mon Sep 17 00:00:00 2001 From: Yassine <159590674+YassineMA03@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:47:29 +0200 Subject: [PATCH 124/222] Fix SinusoidalPositionalEncoding ignoring scale=0.0 (#4098) --- python/mlx/nn/layers/positional_encoding.py | 2 +- python/tests/test_nn.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/positional_encoding.py b/python/mlx/nn/layers/positional_encoding.py index a464354255..49f300d4b1 100644 --- a/python/mlx/nn/layers/positional_encoding.py +++ b/python/mlx/nn/layers/positional_encoding.py @@ -107,7 +107,7 @@ def __init__( self._sigmas = self._sigmas * (2 * math.pi) # Save some constants that define the implementation - self.scale = scale or (2 / dims) ** 0.5 + self.scale = scale if scale is not None else (2 / dims) ** 0.5 self.cos_first = cos_first def __call__(self, x): diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 5c85133a9f..67828fd86e 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -1071,6 +1071,13 @@ def test_sin_pe(self): with self.assertRaises(ValueError): nn.SinusoidalPositionalEncoding(dims) + # An explicit scale=0.0 must be respected rather than falling back + # to the default, since 0.0 is falsy but a valid scale value. + m = nn.SinusoidalPositionalEncoding(16, scale=0.0) + self.assertEqual(m.scale, 0.0) + y = m(x) + self.assertTrue(mx.array_equal(y, mx.zeros_like(y))) + def test_sigmoid(self): x = mx.array([1.0, 0.0, -1.0]) y1 = mx.sigmoid(x) From bf561420ba8fb6270ff13e54d29d3470592c30e3 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 9 Aug 2026 22:10:16 -0700 Subject: [PATCH 125/222] Fix any and all treating -0.0 as nonzero (#4090) --- mlx/backend/cpu/reduce.cpp | 20 +++++++++++++++--- mlx/backend/metal/kernels/reduce.metal | 13 +++++++----- mlx/backend/metal/reduce.cpp | 15 +++++++++++++ python/tests/test_reduce.py | 29 ++++++++++++++++++++++++++ 4 files changed, 69 insertions(+), 8 deletions(-) diff --git a/mlx/backend/cpu/reduce.cpp b/mlx/backend/cpu/reduce.cpp index 41764f4c8d..64b3fab339 100644 --- a/mlx/backend/cpu/reduce.cpp +++ b/mlx/backend/cpu/reduce.cpp @@ -460,6 +460,9 @@ void Reduce::eval_cpu(const std::vector& inputs, array& out) { switch (reduce_type_) { case Reduce::And: case Reduce::Or: { + // Integers can be reduced as whatever type has the same width, since + // only their bits matter. Floats cannot: -0.0 compares equal to zero + // but has a bit set, so it has to be tested as a float. switch (in.dtype()) { case bool_: case uint8: @@ -468,21 +471,32 @@ void Reduce::eval_cpu(const std::vector& inputs, array& out) { break; case int16: case uint16: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case float16: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case bfloat16: - reduce_dispatch_and_or(in, out, reduce_type_, axes_); + reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; case uint32: case int32: - case float32: reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; + case float32: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case uint64: case int64: - case float64: + // complex64 stays on the integer path. Testing it as a complex + // would go through complex64_t's conversion to float and only look + // at the real part, which would miss 1j. case complex64: reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; + case float64: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; } break; } diff --git a/mlx/backend/metal/kernels/reduce.metal b/mlx/backend/metal/kernels/reduce.metal index de5dfbad7d..3f5f91be49 100644 --- a/mlx/backend/metal/kernels/reduce.metal +++ b/mlx/backend/metal/kernels/reduce.metal @@ -124,11 +124,14 @@ instantiate_init_min_max(max, Max) instantiate_row_reduce_general(name##tname, itype, otype, op) \ instantiate_col_reduce_general(name##tname, itype, otype, op) -#define instantiate_and_or(name, op) \ - instantiate_reduce_functions(name, bool_, bool, bool, op) \ - instantiate_reduce_functions(name, int16, int16_t, bool, op) \ - instantiate_reduce_functions(name, int32, int32_t, bool, op) \ - instantiate_reduce_functions(name, int64, int64_t, bool, op) +#define instantiate_and_or(name, op) \ + instantiate_reduce_functions(name, bool_, bool, bool, op) \ + instantiate_reduce_functions(name, int16, int16_t, bool, op) \ + instantiate_reduce_functions(name, int32, int32_t, bool, op) \ + instantiate_reduce_functions(name, int64, int64_t, bool, op) \ + instantiate_reduce_functions(name, float16, float16_t, bool, op) \ + instantiate_reduce_functions(name, bfloat16, bfloat16_t, bool, op) \ + instantiate_reduce_functions(name, float32, float, bool, op) instantiate_and_or(and, And) instantiate_and_or(or, Or) diff --git a/mlx/backend/metal/reduce.cpp b/mlx/backend/metal/reduce.cpp index 52d15288df..ac11ac6359 100644 --- a/mlx/backend/metal/reduce.cpp +++ b/mlx/backend/metal/reduce.cpp @@ -273,6 +273,21 @@ std::pair remap_reduce_types( } return {in.dtype(), in.dtype()}; } else if (op_name == "and" || op_name == "or") { + // Integers can be tested as whatever type has the same width, since only + // their bits matter. Floats cannot: -0.0 compares equal to zero but has a + // bit set, so it has to be tested as a float. complex64 stays on the + // integer path, since testing it as a complex would only look at the real + // part and miss 1j. + switch (in.dtype()) { + case float16: + return {float16, bool_}; + case bfloat16: + return {bfloat16, bool_}; + case float32: + return {float32, bool_}; + default: + break; + } if (in.dtype().size() == 1) { return {bool_, bool_}; } else if (in.dtype().size() == 2) { diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 4f37ebbaf8..6ac8fc1504 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -244,6 +244,35 @@ def test_long_column(self): c2 = b.sum(0) self.assertTrue(np.all(c1 == c2)) + def test_and_or_negative_zero(self): + # -0.0 equals zero but has its sign bit set, so it must not be treated + # as truthy just because its bit pattern is nonzero + for dtype in ["float32", "float16", "float64"]: + with self.subTest(dtype=dtype): + for values in [ + [0.0, -0.0], + [-0.0, -0.0], + [-0.0] * 70, + [-0.0, 0.0, 1.0], + [0.5, 0.0], + ]: + a_np = np.array(values, dtype=getattr(np, dtype)) + a_mx = mx.array(a_np) + for op in ["any", "all"]: + self.assertEqual( + getattr(mx, op)(a_mx).item(), + bool(getattr(np, op)(a_np)), + msg=f"{op} {dtype} {values}", + ) + + x_np = np.array([[0.0, -0.0], [1.0, 0.0], [-0.0, -0.0]], dtype=np.float32) + x_mx = mx.array(x_np) + for op in ["any", "all"]: + self.assertEqual( + getattr(mx, op)(x_mx, axis=1).tolist(), + getattr(np, op)(x_np, axis=1).tolist(), + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True) From a9735ebe4454b9f307d0b8c41726d7366c9d8f07 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:52:43 -0700 Subject: [PATCH 126/222] chore: Fix Eigh CPU dtype error message (#4130) --- mlx/backend/cpu/eigh.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/cpu/eigh.cpp b/mlx/backend/cpu/eigh.cpp index ce624453ba..ff5a4805a3 100644 --- a/mlx/backend/cpu/eigh.cpp +++ b/mlx/backend/cpu/eigh.cpp @@ -244,7 +244,7 @@ void Eigh::eval_cpu( break; default: throw std::runtime_error( - "[Eigh::eval_cpu] only supports float32 or float64."); + "[Eigh::eval_cpu] only supports float32, float64, or complex64."); } } From f94d4e93848c0c9224d47534eedc6f1bd93ec946 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:53:44 -0700 Subject: [PATCH 127/222] chore: Fix random distribution dtype error messages (#4131) --- mlx/random.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mlx/random.cpp b/mlx/random.cpp index 1937743793..164d5f2441 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -187,8 +187,7 @@ array normal( return complex_normal(shape, loc, scale, key, s); } else if (!issubdtype(dtype, floating)) { throw std::invalid_argument( - "[normal] Can only generate uniform numbers with " - "floating point type."); + "[normal] Only floating point and complex64 types are supported."); } auto stream = to_stream(s); @@ -459,8 +458,7 @@ array laplace( StreamOrDevice s /* = {} */) { if (!issubdtype(dtype, floating)) { throw std::invalid_argument( - "[laplace] Can only generate uniform numbers with real" - "floating point type."); + "[laplace] Only real floating point types are supported."); } auto stream = to_stream(s); From 3a67a1bdebc0e124ceec24d019c66871a0c708d1 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:54:20 -0700 Subject: [PATCH 128/222] chore: Fix safetensors dtype error formatting (#4132) --- mlx/io/safetensors.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/io/safetensors.cpp b/mlx/io/safetensors.cpp index d1b8440a06..75c8a6ddc2 100644 --- a/mlx/io/safetensors.cpp +++ b/mlx/io/safetensors.cpp @@ -102,7 +102,7 @@ Dtype dtype_from_safetensor_str(std::string_view str) { return uint8; } else { std::ostringstream msg; - msg << "[safetensor] unsupported dtype" << str; + msg << "[safetensor] unsupported dtype " << str; throw std::runtime_error(msg.str()); } } From 53087cdb314a5e4abf7382b4a2c9bde8b8dbc01b Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:54:48 -0700 Subject: [PATCH 129/222] chore: Fix no-GPU synchronize error formatting (#4133) --- mlx/backend/no_gpu/eval.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/no_gpu/eval.cpp b/mlx/backend/no_gpu/eval.cpp index 3966754679..b08f6ec6fb 100644 --- a/mlx/backend/no_gpu/eval.cpp +++ b/mlx/backend/no_gpu/eval.cpp @@ -28,7 +28,7 @@ void finalize(Stream) { } void synchronize(Stream) { - throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); + throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); } void clear_streams() {} From 5146055d94c1dd6712e63a2b0ebd4c9fb3648268 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:57:25 +0800 Subject: [PATCH 130/222] Fix mx.isinf not considering imaginary plane (#4092) --- mlx/ops.cpp | 3 +++ python/tests/test_ops.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index d00130ee64..b8973e8ad8 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1957,6 +1957,9 @@ array isinf(const array& a, StreamOrDevice s /* = {} */) { if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) { return full(a.shape(), false, bool_, s); } + if (issubdtype(a.dtype(), complexfloating)) { + return logical_or(isinf(real(a, s), s), isinf(imag(a, s), s), s); + } return logical_or(isposinf(a, s), isneginf(a, s), s); } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index b6f9ef6ebf..14dab531cf 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -468,6 +468,26 @@ def test_isinf(self): x = mx.array([0.0, float("inf")]).astype(mx.complex64) self.assertEqual(mx.isinf(x).tolist(), [False, True]) + inf = float("inf") + x = mx.array( + [ + complex(0.0, inf), + complex(inf, 0.0), + complex(inf, inf), + complex(0.0, -inf), + complex(-inf, 0.0), + complex(-inf, -inf), + complex(3.0, 4.0), + complex(0.0, 0.0), + ], + dtype=mx.complex64, + ) + self.assertEqual( + mx.isinf(x).tolist(), + [True, True, True, True, True, True, False, False], + ) + np.testing.assert_array_equal(np.isinf(np.array(x, copy=False)), mx.isinf(x)) + self.assertEqual(mx.isinf(0 * mx.array(float("inf"))).tolist(), False) x = mx.array([-2147483648, 0, 2147483647], dtype=mx.int32) @@ -488,6 +508,16 @@ def test_isfinite(self): x = x.astype(mx.bfloat16) self.assertEqual(mx.isfinite(x).tolist(), [True, False, False]) + inf = float("inf") + x = mx.array( + [complex(0.0, inf), complex(inf, 0.0), complex(3.0, 4.0)], + dtype=mx.complex64, + ) + self.assertEqual(mx.isfinite(x).tolist(), [False, False, True]) + np.testing.assert_array_equal( + np.isfinite(np.array(x, copy=False)), mx.isfinite(x) + ) + def test_tri(self): for shape in [[4], [4, 4], [2, 10]]: for diag in [-1, 0, 1, -2]: From 386a5543e2b4423e37e67f0eb307678e134c30a6 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:07:07 -0700 Subject: [PATCH 131/222] chore: Use dispatch_all_types in cpu copy (#4104) --- mlx/backend/cpu/copy.cpp | 49 ++++------------------------------------ 1 file changed, 5 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/copy.cpp b/mlx/backend/cpu/copy.cpp index d14736c15e..6b6d825252 100644 --- a/mlx/backend/cpu/copy.cpp +++ b/mlx/backend/cpu/copy.cpp @@ -7,6 +7,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" namespace mlx::core { @@ -196,50 +197,10 @@ void copy(const array& src, array& dst, CopyType ctype, Args&&... args) { template void copy(const array& src, array& dst, CopyType ctype, Args&&... args) { - switch (dst.dtype()) { - case bool_: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint8: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint16: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint32: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint64: - copy(src, dst, ctype, std::forward(args)...); - break; - case int8: - copy(src, dst, ctype, std::forward(args)...); - break; - case int16: - copy(src, dst, ctype, std::forward(args)...); - break; - case int32: - copy(src, dst, ctype, std::forward(args)...); - break; - case int64: - copy(src, dst, ctype, std::forward(args)...); - break; - case float16: - copy(src, dst, ctype, std::forward(args)...); - break; - case float32: - copy(src, dst, ctype, std::forward(args)...); - break; - case float64: - copy(src, dst, ctype, std::forward(args)...); - break; - case bfloat16: - copy(src, dst, ctype, std::forward(args)...); - break; - case complex64: - copy(src, dst, ctype, std::forward(args)...); - break; - } + dispatch_all_types(dst.dtype(), [&](auto type_tag) { + using DstT = MLX_GET_TYPE(type_tag); + copy(src, dst, ctype, std::forward(args)...); + }); } template From 07c936615dff3f007c4f07adaebdb8a4ca6d9a00 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:08:46 -0700 Subject: [PATCH 132/222] chore: Use dispatch_all_types in dispatch_gather_axis (#4109) --- mlx/backend/cpu/indexing.cpp | 48 +++--------------------------------- 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/indexing.cpp b/mlx/backend/cpu/indexing.cpp index 54dbfdd7af..fc44888a3f 100644 --- a/mlx/backend/cpu/indexing.cpp +++ b/mlx/backend/cpu/indexing.cpp @@ -241,50 +241,10 @@ void dispatch_gather_axis( const array& inds, array& out, const int axis) { - switch (out.dtype()) { - case bool_: - gather_axis(src, inds, out, axis); - break; - case uint8: - gather_axis(src, inds, out, axis); - break; - case uint16: - gather_axis(src, inds, out, axis); - break; - case uint32: - gather_axis(src, inds, out, axis); - break; - case uint64: - gather_axis(src, inds, out, axis); - break; - case int8: - gather_axis(src, inds, out, axis); - break; - case int16: - gather_axis(src, inds, out, axis); - break; - case int32: - gather_axis(src, inds, out, axis); - break; - case int64: - gather_axis(src, inds, out, axis); - break; - case float16: - gather_axis(src, inds, out, axis); - break; - case float32: - gather_axis(src, inds, out, axis); - break; - case float64: - gather_axis(src, inds, out, axis); - break; - case bfloat16: - gather_axis(src, inds, out, axis); - break; - case complex64: - gather_axis(src, inds, out, axis); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + gather_axis(src, inds, out, axis); + }); } void GatherAxis::eval_cpu(const std::vector& inputs, array& out) { From bcf1b57ce55e0dcb4464b8befc8f4944590b4dca Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:11:48 -0700 Subject: [PATCH 133/222] chore: Use dispatch_inexact_types in CPU unary_fp (#4111) --- mlx/backend/cpu/unary.h | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/mlx/backend/cpu/unary.h b/mlx/backend/cpu/unary.h index e4c9d47b90..096521df89 100644 --- a/mlx/backend/cpu/unary.h +++ b/mlx/backend/cpu/unary.h @@ -106,27 +106,9 @@ void unary_fp(const array& a, array& out, Op op, Stream stream) { encoder.dispatch([a = array::unsafe_weak_copy(a), out = array::unsafe_weak_copy(out), op = op]() mutable { - switch (out.dtype()) { - case bfloat16: - unary_op(a, out, op); - break; - case float16: - unary_op(a, out, op); - break; - case float32: - unary_op(a, out, op); - break; - case float64: - unary_op(a, out, op); - break; - case complex64: - unary_op(a, out, op); - break; - default: - std::ostringstream err; - err << "[unary_fp] Does not support " << out.dtype(); - throw std::runtime_error(err.str()); - } + dispatch_inexact_types(out.dtype(), "[unary_fp]", [&](auto type_tag) { + unary_op(a, out, op); + }); }); } From e0f1134776cda9b0b2c63d4a20dda0cd77580ca8 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:13:18 -0700 Subject: [PATCH 134/222] chore: Use dispatch_all_types in comparison_op_cpu (#4112) --- mlx/backend/cpu/binary.h | 48 ++++------------------------------------ 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/binary.h b/mlx/backend/cpu/binary.h index caf54f21f0..acbb71aae3 100644 --- a/mlx/backend/cpu/binary.h +++ b/mlx/backend/cpu/binary.h @@ -335,50 +335,10 @@ void comparison_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (a.dtype()) { - case bool_: - binary_op(a, b, out, bopt); - break; - case uint8: - binary_op(a, b, out, bopt); - break; - case uint16: - binary_op(a, b, out, bopt); - break; - case uint32: - binary_op(a, b, out, bopt); - break; - case uint64: - binary_op(a, b, out, bopt); - break; - case int8: - binary_op(a, b, out, bopt); - break; - case int16: - binary_op(a, b, out, bopt); - break; - case int32: - binary_op(a, b, out, bopt); - break; - case int64: - binary_op(a, b, out, bopt); - break; - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - } + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } From 93e8b0de3753e234dea81a113de427d5933bd3e1 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:17:56 -0700 Subject: [PATCH 135/222] chore: Use dispatch_all_types in RingGroup (#4120) --- mlx/distributed/ring/ring.cpp | 79 ++++++----------------------------- 1 file changed, 13 insertions(+), 66 deletions(-) diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index a74122239a..3e0c2a3221 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -22,72 +22,13 @@ #include "mlx/distributed/distributed_impl.h" #include "mlx/distributed/reduction_ops.h" #include "mlx/distributed/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/threadpool.h" #ifndef SOL_TCP #define SOL_TCP IPPROTO_TCP #endif -#define SWITCH_TYPE(x, ...) \ - switch ((x).dtype()) { \ - case bool_: { \ - using T = bool; \ - __VA_ARGS__; \ - } break; \ - case int8: { \ - using T = int8_t; \ - __VA_ARGS__; \ - } break; \ - case int16: { \ - using T = int16_t; \ - __VA_ARGS__; \ - } break; \ - case int32: { \ - using T = int32_t; \ - __VA_ARGS__; \ - } break; \ - case int64: { \ - using T = int64_t; \ - __VA_ARGS__; \ - } break; \ - case uint8: { \ - using T = uint8_t; \ - __VA_ARGS__; \ - } break; \ - case uint16: { \ - using T = uint16_t; \ - __VA_ARGS__; \ - } break; \ - case uint32: { \ - using T = uint32_t; \ - __VA_ARGS__; \ - } break; \ - case uint64: { \ - using T = uint64_t; \ - __VA_ARGS__; \ - } break; \ - case bfloat16: { \ - using T = bfloat16_t; \ - __VA_ARGS__; \ - } break; \ - case float16: { \ - using T = float16_t; \ - __VA_ARGS__; \ - } break; \ - case float32: { \ - using T = float; \ - __VA_ARGS__; \ - } break; \ - case float64: { \ - using T = double; \ - __VA_ARGS__; \ - } break; \ - case complex64: { \ - using T = complex64_t; \ - __VA_ARGS__; \ - } break; \ - } - namespace mlx::core::distributed::ring { constexpr const size_t ALL_SUM_SIZE = 8 * 1024 * 1024; @@ -493,18 +434,24 @@ class RingGroup : public GroupImpl { } void all_sum(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::SumOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::SumOp()); + }); } void all_max(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::MaxOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MaxOp()); + }); } void all_min(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::MinOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MinOp()); + }); } std::shared_ptr split(int color, int key = -1) override { From 94c29eedf4c082815d43e3e2b59ce2130b435b16 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:21:35 -0700 Subject: [PATCH 136/222] chore: Use dispatch_all_types in ArgSort::eval_cpu (#4123) --- mlx/backend/cpu/sort.cpp | 33 +++------------------------------ 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index 233900432a..3ba748c235 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -326,36 +326,9 @@ void ArgSort::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([in = array::unsafe_weak_copy(in), out = array::unsafe_weak_copy(out), axis_ = axis_]() mutable { - switch (in.dtype()) { - case bool_: - return argsort(in, out, axis_); - case uint8: - return argsort(in, out, axis_); - case uint16: - return argsort(in, out, axis_); - case uint32: - return argsort(in, out, axis_); - case uint64: - return argsort(in, out, axis_); - case int8: - return argsort(in, out, axis_); - case int16: - return argsort(in, out, axis_); - case int32: - return argsort(in, out, axis_); - case int64: - return argsort(in, out, axis_); - case float32: - return argsort(in, out, axis_); - case float64: - return argsort(in, out, axis_); - case float16: - return argsort(in, out, axis_); - case bfloat16: - return argsort(in, out, axis_); - case complex64: - return argsort(in, out, axis_); - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + argsort(in, out, axis_); + }); }); } From 1b312a98038d4360dc21e75a2a555d211a7e5dcc Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:31:13 -0700 Subject: [PATCH 137/222] chore: Use dispatch_all_types in Softmax::eval_cpu (#4127) --- mlx/backend/cpu/softmax.cpp | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/mlx/backend/cpu/softmax.cpp b/mlx/backend/cpu/softmax.cpp index 4c2941e965..d62e86a9c3 100644 --- a/mlx/backend/cpu/softmax.cpp +++ b/mlx/backend/cpu/softmax.cpp @@ -6,6 +6,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/types/limits.h" @@ -139,32 +140,23 @@ void Softmax::eval_cpu(const std::vector& inputs, array& out) { auto in = set_output(inputs[0]); - switch (in.dtype()) { - case float32: - softmax(in, out, stream()); - break; - case float16: + dispatch_all_types(in.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr ( + std::is_same_v || std::is_same_v) { if (precise_) { - softmax(in, out, stream()); + softmax(in, out, stream()); } else { - softmax(in, out, stream()); + softmax(in, out, stream()); } - break; - case bfloat16: - if (precise_) { - softmax(in, out, stream()); - } else { - softmax(in, out, stream()); - } - break; - case float64: - softmax(in, out, stream()); - break; - default: + } else if constexpr ( + std::is_same_v || std::is_same_v) { + softmax(in, out, stream()); + } else { throw std::runtime_error( "[softmax] Only defined for floating point types."); - break; - } + } + }); } } // namespace mlx::core From e78d894c8f718a805341e1aa1cc835d9c9b8462e Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:53:41 -0700 Subject: [PATCH 138/222] chore: Use dispatch_all_types in Arange::eval_gpu (#4128) --- mlx/backend/metal/primitives.cpp | 49 ++++++++------------------------ 1 file changed, 12 insertions(+), 37 deletions(-) diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index d5bbf797e4..45929e27dd 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include "mlx/backend/common/slicing.h" #include "mlx/backend/common/utils.h" @@ -11,6 +12,7 @@ #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" #include "mlx/utils.h" @@ -40,45 +42,18 @@ void Arange::eval_gpu(const std::vector& inputs, array& out) { auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - switch (out.dtype()) { - case bool_: // unsupported + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("[Arange::eval_gpu] Does not support bool"); - case uint8: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint64: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int8: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int64: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case float16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case float32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case bfloat16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - default: + } else if constexpr ( + std::is_integral_v || std::is_same_v || + std::is_same_v || std::is_same_v) { + arange_set_scalars(start_, start_ + step_, compute_encoder); + } else { throw std::runtime_error("[Arange::eval_gpu] Does not support type."); - } + } + }); compute_encoder.set_output_array(out, 2); compute_encoder.dispatch_threads(grid_dims, group_dims); From a076a632596eb61b1e876c0512b54b781584f756 Mon Sep 17 00:00:00 2001 From: Dwijen Patel Date: Mon, 10 Aug 2026 15:42:27 -0700 Subject: [PATCH 139/222] Pick BM from rows per expert in gather_qmm_rhs_nax (#4023) Co-authored-by: Cheng --- benchmarks/python/gather_qmm_bench.py | 21 +++++++++++++++++++ .../metal/kernels/fp_quantized_nax.metal | 4 +++- mlx/backend/metal/kernels/quantized_nax.metal | 4 +++- mlx/backend/metal/quantized.cpp | 6 ++++-- python/tests/test_quantized.py | 2 ++ 5 files changed, 33 insertions(+), 4 deletions(-) diff --git a/benchmarks/python/gather_qmm_bench.py b/benchmarks/python/gather_qmm_bench.py index 17c06d57d2..c05db91992 100644 --- a/benchmarks/python/gather_qmm_bench.py +++ b/benchmarks/python/gather_qmm_bench.py @@ -80,5 +80,26 @@ def equivalent_matmul(x, w1, w2): time_fn(equivalent_matmul, x, w1, w2) +def time_gather_qmm_short_runs(): + # Many experts and few tokens, so each expert gets N * I / E = 16 rows. + N, E, I = 512, 256, 8 + x = mx.random.normal((N, 1, 1, D)) / 1024**0.5 + w1 = mx.random.normal((E, M, D)) / 1024**0.5 + w2 = mx.random.normal((E, D, M)) / 1024**0.5 + w1 = mx.quantize(w1) + w2 = mx.quantize(w2) + indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32) + mx.eval(x, w1, w2, indices) + + def gather_mm(x, w1, w2, indices): + x, idx, inv_order = gather_sort(x, indices) + x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=True) + x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=True) + return scatter_unsort(x, inv_order, indices.shape) + + time_fn(gather_mm, x, w1, w2, indices) + + if __name__ == "__main__": time_gather_qmm() + time_gather_qmm_short_runs() diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.metal b/mlx/backend/metal/kernels/fp_quantized_nax.metal index 4d65a384d3..c736f1809e 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.metal +++ b/mlx/backend/metal/kernels/fp_quantized_nax.metal @@ -62,7 +62,9 @@ #define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits) + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 32, 64, 64, 2, 2, true, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 32, 64, 64, 2, 2, false, mode, group_size, bits) #define instantiate_quantized_modes(type, mode, group_size, bits) \ instantiate_quantized_all_aligned(type, mode, group_size, bits) \ diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 5a9c9fb874..27302ecb5f 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -78,7 +78,9 @@ #define instantiate_quantized_all_rhs(type, group_size, bits) \ instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \ - instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false) + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 32, 64, 64, 2, 2, true) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 32, 64, 64, 2, 2, false) #define instantiate_quantized_funcs(type, group_size, bits) \ instantiate_quantized_all_batched(type, group_size, bits) \ diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index cb42e817d5..b754132dc4 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1489,8 +1489,10 @@ void gather_qmm_rhs_nax( biases = ensure_row_contiguous(*biases_, d, s); } - // TODO: Tune the block sizes - int bm = 64, bn = 64, bk = 64; + // Use smaller bm for many experts and few tokens. + int E = w.size() / w.shape(-1) / w.shape(-2); + int bm = (M / E < 64) ? 32 : 64; + int bn = 64, bk = 64; int wm = 2, wn = 2; const bool align_M = (M % bm) == 0; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 806784cfc7..4ca32f2279 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1356,6 +1356,8 @@ def scatter_unsort(x, inv_order, shape=None): (32, 512, 544, 4, 2, True, "mxfp4"), (32, 512, 544, 4, 2, True, "nvfp4"), (32, 512, 544, 4, 2, True, "mxfp8"), + (39, 512, 512, 4, 2, True, "affine"), + (128, 512, 512, 4, 2, True, "affine"), (133, 512, 512, 4, 2, True, "affine"), (133, 512, 555, 4, 2, True, "affine"), (133, 512, 512, 4, 2, True, "affine"), From 98a6ec8034376a2a2bc810ed38e1b8f7b85ac0d9 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:46:00 -0700 Subject: [PATCH 140/222] Enable complex64 take_along_axis backward on Metal (#4094) --- mlx/backend/metal/indexing.cpp | 3 ++- python/tests/test_autograd.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/indexing.cpp b/mlx/backend/metal/indexing.cpp index d562fa6f22..117e402f7b 100644 --- a/mlx/backend/metal/indexing.cpp +++ b/mlx/backend/metal/indexing.cpp @@ -520,7 +520,8 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { } void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { - if (size_of(out.dtype()) == 8) { + if (size_of(out.dtype()) == 8 && + !(out.dtype() == complex64 && reduce_type_ == ScatterAxis::Sum)) { std::ostringstream msg; msg << "[ScatterAxis::eval_gpu] Does not support " << out.dtype(); throw std::invalid_argument(msg.str()); diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index 2f72fba113..2e6b33076c 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -519,6 +519,23 @@ def scatter_axis_fun(w): grad = mx.grad(scatter_axis_fun)(mx.ones((3, 3))) self.assertTrue(mx.array_equal(grad, mx.ones((3, 3)))) + def test_take_along_axis_complex_vjp(self): + x = mx.zeros((4,), dtype=mx.complex64) + indices = mx.array([3, 3], dtype=mx.int32) + cotangent = mx.array([1 + 2j, 3 + 4j], dtype=mx.complex64) + _, (gradient,) = mx.vjp( + lambda z: mx.take_along_axis(z, indices, axis=0), + [x], + [cotangent], + ) + mx.eval(gradient) + self.assertEqualArray( + gradient, + mx.array([0j, 0j, 0j, 4 + 6j], dtype=mx.complex64), + atol=0, + rtol=0, + ) + def test_scatter_add_vjp(self): def fun(src, updates): x = src.at[mx.array([1, 3])].add(updates) From 5ec30acd5aebfe92ea1653f192365cbe788c06a6 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:57:14 -0400 Subject: [PATCH 141/222] Add mx.searchsorted with CPU, Metal and CUDA kernels (#4035) Co-authored-by: Cheng --- docs/src/python/ops.rst | 1 + mlx/backend/cpu/sort.cpp | 66 ++++++++++++ mlx/backend/cuda/sort.cu | 83 ++++++++++++++- mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/jit/includes.h | 1 + mlx/backend/metal/jit_kernels.cpp | 19 ++++ mlx/backend/metal/kernels.h | 6 ++ mlx/backend/metal/kernels/CMakeLists.txt | 1 + mlx/backend/metal/kernels/searchsorted.h | 34 ++++++ mlx/backend/metal/kernels/searchsorted.metal | 28 +++++ mlx/backend/metal/nojit_kernels.cpp | 8 ++ mlx/backend/metal/sort.cpp | 47 +++++++++ mlx/backend/no_cpu/primitives.cpp | 1 + mlx/backend/no_gpu/primitives.cpp | 1 + mlx/export.cpp | 1 + mlx/ops.cpp | 31 ++++++ mlx/ops.h | 7 ++ mlx/primitives.cpp | 41 ++++++++ mlx/primitives.h | 21 ++++ python/src/ops.cpp | 34 ++++++ python/tests/test_export_import.py | 18 ++++ python/tests/test_ops.py | 104 +++++++++++++++++++ 22 files changed, 553 insertions(+), 1 deletion(-) create mode 100644 mlx/backend/metal/kernels/searchsorted.h create mode 100644 mlx/backend/metal/kernels/searchsorted.metal diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 2303ed80c6..f55ba1b4d4 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -168,6 +168,7 @@ Operations savez_compressed save_gguf save_safetensors + searchsorted sigmoid sign sin diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index 3ba748c235..d8a6f3559b 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -311,6 +311,43 @@ void argpartition(const array& in, array& out, int axis, int kth) { } } +template +void searchsorted(const array& a, const array& v, array& out) { + auto n = static_cast(a.size()); + auto a_stride = a.strides()[0]; // sequence is 1D + const T* a_ptr = a.data(); + const T* v_ptr = v.data(); + uint32_t* out_ptr = out.data(); + + auto bound = [a_ptr, a_stride, n](T x) { + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + (hi - lo) / 2; + T m = a_ptr[static_cast(mid) * a_stride]; + bool below = Right ? !nan_aware_less(x, m) : nan_aware_less(m, x); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + + if (v.flags().row_contiguous) { + for (size_t i = 0; i < v.size(); ++i) { + out_ptr[i] = bound(v_ptr[i]); + } + } else { + ContiguousIterator it(v); + for (size_t i = 0; i < v.size(); ++i) { + out_ptr[i] = bound(v_ptr[it.loc]); + it.step(); + } + } +} + } // namespace void ArgSort::eval_cpu(const std::vector& inputs, array& out) { @@ -451,4 +488,33 @@ void Partition::eval_cpu(const std::vector& inputs, array& out) { }); } +void SearchSorted::eval_cpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + auto& a = inputs[0]; + auto& v = inputs[1]; + + out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + auto& encoder = cpu::get_command_encoder(stream()); + encoder.set_input_array(a); + encoder.set_input_array(v); + encoder.set_output_array(out); + encoder.dispatch([a = array::unsafe_weak_copy(a), + v = array::unsafe_weak_copy(v), + out = array::unsafe_weak_copy(out), + right = right_]() mutable { + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if (right) { + searchsorted(a, v, out); + } else { + searchsorted(a, v, out); + } + }); + }); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu index 43756f7078..d9c77387ce 100644 --- a/mlx/backend/cuda/sort.cu +++ b/mlx/backend/cuda/sort.cu @@ -11,6 +11,7 @@ #include "mlx/dtype_utils.h" #include "mlx/primitives.h" +#include #include #include #include @@ -21,6 +22,8 @@ constexpr int N_PER_THREAD = 8; namespace cu { +namespace cg = cooperative_groups; + template __device__ __forceinline__ T nan_value(); @@ -716,6 +719,40 @@ __global__ void mb_block_merge_kernel( } } +template +__device__ __forceinline__ uint32_t +searchsorted_impl(const T* a, T v, uint32_t n, int64_t a_stride) { + LessThan lt; + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + (hi - lo) / 2; + T m = a[static_cast(mid) * a_stride]; + bool below = Right ? !lt(v, m) : lt(m, v); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +template +__global__ void searchsorted( + const T* a, + const T* v, + uint32_t* out, + int64_t size, + uint32_t n, + int64_t a_stride) { + int64_t index = cg::this_grid().thread_rank(); + if (index >= size) { + return; + } + out[index] = searchsorted_impl(a, v[index], n, a_stride); +} + } // namespace cu namespace { @@ -1074,4 +1111,48 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { gpu_sort(stream(), inputs[0], out, axis_, false); } -} // namespace mlx::core \ No newline at end of file +void SearchSorted::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("SearchSorted::eval_gpu"); + assert(inputs.size() == 2); + auto& s = stream(); + auto& a = inputs[0]; + auto v = inputs[1]; + + auto& encoder = cu::get_command_encoder(s); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + if (out.size() == 0) { + return; + } + + if (!v.flags().row_contiguous) { + v = contiguous_copy_gpu(v, s); + encoder.add_temporary(v); + } + + encoder.set_input_array(a); + encoder.set_input_array(v); + encoder.set_output_array(out); + + int64_t a_stride = a.strides()[0]; // sequence is 1D + + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using CTYPE = MLX_GET_TYPE(type_tag); + using T = cuda_type_t; + dispatch_bool(right_, [&](auto right) { + auto [num_blocks, block_dims] = + get_launch_args(out, out.size() > INT32_MAX); + encoder.add_kernel_node( + cu::searchsorted, + num_blocks, + block_dims, + gpu_ptr(a), + gpu_ptr(v), + gpu_ptr(out), + static_cast(out.size()), + static_cast(a.size()), + a_stride); + }); + }); +} + +} // namespace mlx::core diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index e7a4d9d2af..ea4a995ade 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(searchsorted kernels/sort.h) make_jit_source( reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h kernels/reduction/reduce_row.h kernels/reduction/reduce_init.h) diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index ac9fb81e26..4fb1be1110 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* searchsorted(); const char* reduce(); const char* gemm(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 1384e06c50..d40bea85b2 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -433,6 +433,25 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + bool right) { + auto lib = d.get_library(kernel_name, [&]() { + std::ostringstream kernel_source; + // The kernel compares through LessThan, which lives in sort.h. + kernel_source << metal::utils() << metal::sort() << metal::searchsorted(); + kernel_source << get_template_definition( + kernel_name, + "searchsorted", + get_type_string(in.dtype()), + right ? "true" : "false"); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_mb_sort_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 973041932a..3e4258c383 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -81,6 +81,12 @@ MTL::ComputePipelineState* get_sort_kernel( int bn, int tn); +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + bool right); + MTL::ComputePipelineState* get_mb_sort_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 6d9a0883f0..e72cb41ab4 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -141,6 +141,7 @@ if(NOT MLX_METAL_JIT) build_kernel(scan scan.h) build_kernel(softmax softmax.h) build_kernel(logsumexp logsumexp.h) + build_kernel(searchsorted searchsorted.h sort.h) build_kernel(sort sort.h) build_kernel(ternary ternary.h ternary_ops.h) build_kernel(unary unary.h unary_ops.h) diff --git a/mlx/backend/metal/kernels/searchsorted.h b/mlx/backend/metal/kernels/searchsorted.h new file mode 100644 index 0000000000..87b7c3f5bb --- /dev/null +++ b/mlx/backend/metal/kernels/searchsorted.h @@ -0,0 +1,34 @@ +// Copyright © 2026 Apple Inc. + +template +METAL_FUNC uint +searchsorted_impl(device const T* a, T v, uint n, int64_t a_stride) { + LessThan lt; + uint lo = 0; + uint hi = n; + while (lo < hi) { + uint mid = lo + (hi - lo) / 2; + T m = a[int64_t(mid) * a_stride]; + bool below = Right ? !lt(v, m) : lt(m, v); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +template +[[kernel]] void searchsorted( + device const T* a [[buffer(0)]], + device const T* v [[buffer(1)]], + device uint* out [[buffer(2)]], + constant const uint& n [[buffer(3)]], + constant const int64_t& a_stride [[buffer(4)]], + uint3 index [[thread_position_in_grid]], + uint3 grid_dim [[threads_per_grid]]) { + auto offset = + index.x + grid_dim.x * (int64_t(index.y) + int64_t(grid_dim.y) * index.z); + out[offset] = searchsorted_impl(a, v[offset], n, a_stride); +} diff --git a/mlx/backend/metal/kernels/searchsorted.metal b/mlx/backend/metal/kernels/searchsorted.metal new file mode 100644 index 0000000000..a48ee2c327 --- /dev/null +++ b/mlx/backend/metal/kernels/searchsorted.metal @@ -0,0 +1,28 @@ +// Copyright © 2026 Apple Inc. + +#include + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/sort.h" +#include "mlx/backend/metal/kernels/searchsorted.h" + +#define instantiate_searchsorted(tname, type) \ + instantiate_kernel("searchsorted_" #tname "_left", \ + searchsorted, type, false) \ + instantiate_kernel("searchsorted_" #tname "_right", \ + searchsorted, type, true) + +instantiate_searchsorted(bool_, bool) +instantiate_searchsorted(uint8, uint8_t) +instantiate_searchsorted(uint16, uint16_t) +instantiate_searchsorted(uint32, uint32_t) +instantiate_searchsorted(uint64, uint64_t) +instantiate_searchsorted(int8, int8_t) +instantiate_searchsorted(int16, int16_t) +instantiate_searchsorted(int32, int32_t) +instantiate_searchsorted(int64, int64_t) +instantiate_searchsorted(float16, half) +instantiate_searchsorted(float32, float) +instantiate_searchsorted(bfloat16, bfloat16_t) +instantiate_searchsorted(complex64, complex64_t) // clang-format on diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 9f6f8782f5..e0a3ff935d 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -100,6 +100,14 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + bool) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_mb_sort_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp index ec965e2fb9..65f144c026 100644 --- a/mlx/backend/metal/sort.cpp +++ b/mlx/backend/metal/sort.cpp @@ -365,4 +365,51 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { gpu_merge_sort(s, d, in, out, axis_, false); } +void SearchSorted::eval_gpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + auto& a = inputs[0]; + auto v = inputs[1]; + + out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + if (a.size() == 0) { + array zero = array(0, out.dtype()); + fill_gpu(zero, out, s); + compute_encoder.add_temporary(std::move(zero)); + return; + } + + if (!v.flags().row_contiguous) { + v = contiguous_copy_gpu(v, s); + compute_encoder.add_temporary(v); + } + + int64_t a_stride = a.strides()[0]; // sequence is 1D + auto n = static_cast(a.size()); + + std::string kernel_name = "searchsorted_"; + concatenate(kernel_name, type_to_name(a), right_ ? "_right" : "_left"); + auto kernel = get_searchsorted_kernel(d, kernel_name, a, right_); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(v, 1); + compute_encoder.set_output_array(out, 2); + compute_encoder.set_bytes(n, 3); + compute_encoder.set_bytes(a_stride, 4); + + size_t thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + thread_group_size = std::min(thread_group_size, out.size()); + MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1); + MTL::Size grid_dims = get_2d_grid_dims(out.shape(), out.strides()); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + } // namespace mlx::core diff --git a/mlx/backend/no_cpu/primitives.cpp b/mlx/backend/no_cpu/primitives.cpp index faaeb0c7c4..e522307156 100644 --- a/mlx/backend/no_cpu/primitives.cpp +++ b/mlx/backend/no_cpu/primitives.cpp @@ -107,6 +107,7 @@ NO_CPU(Round) NO_CPU(Scan) NO_CPU(Scatter) NO_CPU(ScatterAxis) +NO_CPU(SearchSorted) NO_CPU(Select) NO_CPU(SegmentedMM) NO_CPU(Sigmoid) diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 0e05e9d19f..b7d7a19467 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -134,6 +134,7 @@ NO_GPU(Round) NO_GPU(Scan) NO_GPU(Scatter) NO_GPU(ScatterAxis) +NO_GPU(SearchSorted) NO_GPU(Select) NO_GPU(SegmentedMM) NO_GPU(Sigmoid) diff --git a/mlx/export.cpp b/mlx/export.cpp index de5b78fd61..bcc0ad161c 100644 --- a/mlx/export.cpp +++ b/mlx/export.cpp @@ -430,6 +430,7 @@ struct PrimitiveFactory { "CumLogaddexp"), SERIALIZE_PRIMITIVE(Scatter), SERIALIZE_PRIMITIVE(ScatterAxis), + SERIALIZE_PRIMITIVE(SearchSorted), SERIALIZE_PRIMITIVE(Select), SERIALIZE_PRIMITIVE(Sigmoid), SERIALIZE_PRIMITIVE(Sign), diff --git a/mlx/ops.cpp b/mlx/ops.cpp index b8973e8ad8..cee62db438 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2803,6 +2803,37 @@ array argpartition( {a}); } +array searchsorted( + const array& sorted_sequence, + const array& values, + const std::string& side /* = "left" */, + StreamOrDevice s /* = {} */) { + if (side != "left" && side != "right") { + std::ostringstream msg; + msg << "[searchsorted] Invalid side '" << side + << "'. Expected 'left' or 'right'."; + throw std::invalid_argument(msg.str()); + } + if (sorted_sequence.ndim() != 1) { + std::ostringstream msg; + msg << "[searchsorted] The sorted sequence must be 1-D but has " + << sorted_sequence.ndim() << " dimensions."; + throw std::invalid_argument(msg.str()); + } + if (sorted_sequence.size() > UINT32_MAX) { + std::ostringstream msg; + msg << "[searchsorted] The sorted sequence has " << sorted_sequence.size() + << " elements, more than the uint32 output can index."; + throw std::invalid_argument(msg.str()); + } + auto dtype = promote_types(sorted_sequence.dtype(), values.dtype()); + return array( + values.shape(), + uint32, + std::make_shared(to_stream(s), side == "right"), + {astype(sorted_sequence, dtype, s), astype(values, dtype, s)}); +} + /** Returns topk elements of the flattened array. */ array topk(const array& a, int k, StreamOrDevice s /* = {}*/) { int size = a.size(); diff --git a/mlx/ops.h b/mlx/ops.h index 8297f77ca8..cacc4e28d6 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -857,6 +857,13 @@ MLX_API array argpartition(const array& a, int kth, StreamOrDevice s = {}); MLX_API array argpartition(const array& a, int kth, int axis, StreamOrDevice s = {}); +/** Find the indices of `values` in `sorted_sequence`. */ +MLX_API array searchsorted( + const array& sorted_sequence, + const array& values, + const std::string& side = "left", + StreamOrDevice s = {}); + /** Returns topk elements of the flattened array. */ MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index c339bc444c..343b20fc14 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -5414,6 +5414,47 @@ bool Softmax::is_equivalent(const Primitive& other) const { return precise_ == s_other.precise_; } +std::pair, std::vector> SearchSorted::vmap( + const std::vector& inputs, + const std::vector& axes) { + if (axes[0] != -1) { + throw std::invalid_argument( + "[searchsorted] Cannot vmap over the sorted sequence, only over the " + "values being searched for."); + } + auto side = right_ ? "right" : "left"; + return {{searchsorted(inputs[0], inputs[1], side, stream())}, {axes[1]}}; +} + +std::vector SearchSorted::vjp( + const std::vector& primals, + const std::vector&, + const std::vector& argnums, + const std::vector&) { + std::vector vjps; + for (auto arg : argnums) { + vjps.push_back(zeros_like(primals[arg], stream())); + } + return vjps; +} + +std::vector SearchSorted::jvp( + const std::vector& primals, + const std::vector&, + const std::vector&) { + return {zeros(primals[1].shape(), uint32, stream())}; +} + +bool SearchSorted::is_equivalent(const Primitive& other) const { + const SearchSorted& r_other = static_cast(other); + return right_ == r_other.right_; +} + +std::vector SearchSorted::output_shapes( + const std::vector& inputs) { + return {inputs[1].shape()}; +} + std::pair, std::vector> Sort::vmap( const std::vector& inputs, const std::vector& axes) { diff --git a/mlx/primitives.h b/mlx/primitives.h index 403490dbe8..3a3d0ba5e5 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -2211,6 +2211,27 @@ class Softmax : public UnaryPrimitive { bool precise_; }; +class SearchSorted : public UnaryPrimitive { + public: + explicit SearchSorted(Stream stream, bool right) + : UnaryPrimitive(stream), right_(right) {} + + void eval_cpu(const std::vector& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_VMAP() + DEFINE_GRADS() + DEFINE_NAME(SearchSorted) + bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; + auto state() const { + return right_; + } + + private: + bool right_; +}; + class Sort : public UnaryPrimitive { public: explicit Sort(Stream stream, int axis) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 8ddb3e4f96..0941793949 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3152,6 +3152,40 @@ void init_ops(nb::module_& m) { Returns: array: The ``uint32`` array containing indices that partition the input. )pbdoc"); + m.def( + "searchsorted", + &mx::searchsorted, + nb::arg(), + nb::arg(), + "side"_a = "left", + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def searchsorted(sorted_sequence: array, values: array, /, side: str = 'left', *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Find the indices that keep ``sorted_sequence`` sorted when inserting ``values``. + + Args: + sorted_sequence (array): A 1-D array sorted in ascending order. + values (array): The values to insert. May have any shape. + side (str, optional): Either ``'left'`` or ``'right'``. With + ``'left'`` the first suitable index is returned, so the result is + the number of elements strictly less than the value. With + ``'right'`` the last is returned, so the result is the number of + elements less than or equal to it. The two differ only where a + value is already present. Default: ``'left'``. + + Returns: + array: A ``uint32`` array with the same shape as ``values``, holding + indices in ``[0, sorted_sequence.size]``. + + Example: + >>> a = mx.array([1, 2, 2, 4]) + >>> mx.searchsorted(a, mx.array([0, 2, 3, 5])) + array([0, 1, 3, 4], dtype=uint32) + >>> mx.searchsorted(a, mx.array([0, 2, 3, 5]), side="right") + array([0, 3, 3, 4], dtype=uint32) + )pbdoc"); m.def( "topk", [](const mx::array& a, diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index 15d3b8c4c9..23f0645b3a 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -314,6 +314,24 @@ def fun(a, b, c): out = imported_fun(x, y, z)[0] self.assertTrue(mx.array_equal(expected, out)) + def test_export_searchsorted(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + # both sides, since the side is the primitive's only state and a lost + # state would still round trip for the default + for side in ("left", "right"): + + def fun(a, v): + return mx.searchsorted(a, v, side=side) + + x = mx.sort(mx.random.uniform(shape=(32,))) + y = mx.random.uniform(shape=(3, 5)) + mx.export_function(path, fun, (x, y)) + imported_fun = mx.import_function(path) + expected = fun(x, y) + out = imported_fun(x, y)[0] + self.assertTrue(mx.array_equal(expected, out)) + def test_export_conv(self): path = os.path.join(self.test_dir, "fn.mlxfn") diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 14dab531cf..3124b56ea8 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2651,6 +2651,110 @@ def test_argpartition(self): expected = mx.array([[0, 0], [1, 1]]) self.assertTrue(mx.array_equal(out, expected)) + def test_searchsorted(self): + def expect(out, want): + self.assertTrue(mx.array_equal(out, mx.array(want)), f"got {out}") + + a = mx.array([1, 2, 2, 4], mx.float32) + v = mx.array([0, 1, 2, 3, 5], mx.float32) + expect(mx.searchsorted(a, v), [0, 0, 1, 3, 4]) + expect(mx.searchsorted(a, v, side="right"), [0, 1, 3, 3, 4]) + self.assertEqual(mx.searchsorted(a, v).dtype, mx.uint32) + + # a local generator, so this stays deterministic without shifting the + # global numpy stream that later tests draw from + rng = np.random.RandomState(0) + for n in (1, 2, 7, 8, 9, 33, 1000): + for m in (1, 5, 64, 257): + for side in ("left", "right"): + a_np = np.sort(rng.randn(n).astype(np.float32)) + v_np = (rng.randn(m) * 2).astype(np.float32) + out = mx.searchsorted(mx.array(a_np), mx.array(v_np), side=side) + expected = np.searchsorted(a_np, v_np, side=side) + self.assertTrue(np.array_equal(np.array(out), expected)) + + # output takes the shape of the values. Compare values too, since + # checking .shape alone never forces an eval, and the 0-d case is the + # one both GPU backends special case. + a = mx.arange(16, dtype=mx.float32) + for shape in [(), (1,), (3, 4), (2, 3, 4)]: + v_np = np.asarray(rng.rand(*shape) * 20, dtype=np.float32) + out = mx.searchsorted(a, mx.array(v_np)) + self.assertEqual(out.shape, shape) + self.assertTrue( + np.array_equal(np.array(out), np.searchsorted(np.array(a), v_np)) + ) + + # non row contiguous values: transposed, sliced and broadcast views all + # have to be read in the output's order rather than the buffer's + base_np = (rng.rand(4, 6) * 20).astype(np.float32) + base = mx.array(base_np) + for v_mx, v_np in [ + (base.T, base_np.T), + (base[::2], base_np[::2]), + (base[:, ::3], base_np[:, ::3]), + (mx.broadcast_to(base[0], (3, 6)), np.broadcast_to(base_np[0], (3, 6))), + ]: + out = mx.searchsorted(a, v_mx) + expected = np.searchsorted(np.array(a), np.ascontiguousarray(v_np)) + self.assertTrue(np.array_equal(np.array(out), expected)) + + # a strided sorted sequence, including a reversed view + wide = mx.array(np.repeat(np.arange(8, dtype=np.float32) * 3, 2)) + v = mx.array([-1.0, 3.0, 7.0, 100.0]) + for a_mx in [wide[::2], wide[1::2]]: + out = mx.searchsorted(a_mx, v) + expected = np.searchsorted(np.array(a_mx), np.array(v)) + self.assertTrue(np.array_equal(np.array(out), expected)) + + desc = mx.array(np.arange(8, dtype=np.float32)[::-1].copy()) + out = mx.searchsorted(desc[::-1], v) + self.assertTrue( + np.array_equal(np.array(out), np.searchsorted(np.arange(8), np.array(v))) + ) + + # integer and mixed dtypes + ai = mx.array([1, 3, 5, 7], mx.int32) + expect(mx.searchsorted(ai, mx.array([0, 4, 8], mx.int32)), [0, 2, 4]) + # promoted, not truncated: the sequence has a 4, so 4.5 lands after it + # at 3 while a truncated 4 would land before it at 2 + expect( + mx.searchsorted( + mx.array([1, 3, 4, 7], mx.int32), mx.array([4.5], mx.float32) + ), + [3], + ) + + # empty inputs on either side + empty = mx.array([], mx.float32) + expect(mx.searchsorted(empty, mx.array([1.0, -1.0])), [0, 0]) + self.assertEqual(mx.searchsorted(mx.arange(4, dtype=mx.float32), empty).size, 0) + + # ordering follows sort, so NaN compares greater than everything + nan = mx.array([float("nan")]) + expect(mx.searchsorted(mx.array([1, 2, 3], mx.float32), nan), [3]) + a = mx.array([1, 2, float("nan")], mx.float32) + expect(mx.searchsorted(a, nan), [2]) + expect(mx.searchsorted(a, nan, side="right"), [3]) + + # vmap over the values, which is the elementwise argument + va = mx.arange(8, dtype=mx.float32) + vs = mx.random.uniform(0, 10, (3, 5)) + out = mx.vmap(lambda x: mx.searchsorted(va, x))(vs) + self.assertTrue(mx.array_equal(out, mx.searchsorted(va, vs))) + + with self.assertRaises(ValueError): + mx.vmap(lambda s: mx.searchsorted(s, mx.array([1.0])))( + mx.zeros((3, 4), mx.float32) + ) + + with self.assertRaises(ValueError): + mx.searchsorted(mx.zeros((3, 4)), mx.array([1.0])) + with self.assertRaises(ValueError): + mx.searchsorted(mx.array(1.0), mx.array([1.0])) + with self.assertRaises(ValueError): + mx.searchsorted(mx.array([1.0, 2.0]), mx.array([1.0]), side="middle") + @unittest.skipIf( os.getenv("LOW_MEMORY", None) is not None, "This test requires a lot of memory", From 47a1bd613ebb67ed03da764c6ffedb8df9939997 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 10 Aug 2026 17:46:18 -0700 Subject: [PATCH 142/222] Stop integer power from hanging on a negative exponent (#4100) Co-authored-by: Cheng --- mlx/backend/cpu/simd/base_simd.h | 5 +++++ python/tests/test_ops.py | 7 ++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index 6a1ee39a59..d69e69ecf3 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -245,6 +245,11 @@ Simd pow(Simd a, Simd b) { return std::pow(base, exp); } else { T res = 1; + if constexpr (std::is_signed_v) { + if (exp < 0) { + return 0; + } + } while (exp) { if (exp & 1) { res *= base; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 3124b56ea8..ac3bb9b49a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3643,8 +3643,13 @@ def test_integer_power(self): x = mx.power(2, mx.array([8, 8, 8, 8, 8, 8, 8, 8])) self.assertTrue(mx.all(x == 256)) - # Doesn't hang + # Doesn't hang. x = mx.power(2, -1) + self.assertEqual(x.item(), 0) + + for dtype in [mx.int8, mx.int16, mx.int32, mx.int64]: + x = mx.power(mx.array([2, -2, 1], dtype), mx.array([-1, -3, -9], dtype)) + self.assertEqual(x.tolist(), [0, 0, 0]) def test_depends(self): a = mx.array([1.0, 2.0, 3.0]) From 32364768039ece055bbd66803b5d4b051bfa448d Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:53:41 -0700 Subject: [PATCH 143/222] Fix second order gradients for sort, partition, topk and cummax/cummin (#4117) --- mlx/primitives.cpp | 26 ++++++++++++++------- python/tests/test_autograd.py | 44 +++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 8 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 343b20fc14..cec010647f 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3390,7 +3390,10 @@ std::vector Partition::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - auto sort_idx = argpartition(primals[0], kth_, axis_, stream()); + // The permutation is locally constant in the input, so cut the gradient + // there to keep higher order derivatives working. + auto sort_idx = + stop_gradient(argpartition(primals[0], kth_, axis_, stream()), stream()); return {put_along_axis( zeros_like(primals[0], stream()), sort_idx, @@ -3405,7 +3408,8 @@ std::vector Partition::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(tangents.size() == 1); - auto sort_idx = argpartition(primals[0], kth_, axis_, stream()); + auto sort_idx = + stop_gradient(argpartition(primals[0], kth_, axis_, stream()), stream()); auto out = take_along_axis(tangents[0], sort_idx, axis_, stream()); return {out}; } @@ -4390,10 +4394,14 @@ std::vector Scan::vjp( iota, array(reverse_ ? n : -1, int32), s); - auto owner = astype( - reverse_ ? cummin(masked, axis_, /* reverse = */ true, true, s) - : cummax(masked, axis_, /* reverse = */ false, true, s), - uint32, + // The owner indices are locally constant in the input, so cut the + // gradient there to keep higher order derivatives working. + auto owner = stop_gradient( + astype( + reverse_ ? cummin(masked, axis_, /* reverse = */ true, true, s) + : cummax(masked, axis_, /* reverse = */ false, true, s), + uint32, + s), s); if (!inclusive_) { @@ -5473,7 +5481,9 @@ std::vector Sort::vjp( // Sort applies a permutation to the input, so the cotangents must be // scattered back to the original positions (the transpose of the // permutation), not gathered forward as in the jvp. - auto sort_idx = argsort(primals[0], axis_, stream()); + // The permutation is locally constant in the input, so cut the gradient + // there to keep higher order derivatives working. + auto sort_idx = stop_gradient(argsort(primals[0], axis_, stream()), stream()); return {put_along_axis( zeros_like(primals[0], stream()), sort_idx, @@ -5488,7 +5498,7 @@ std::vector Sort::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(tangents.size() == 1); - auto sort_idx = argsort(primals[0], axis_, stream()); + auto sort_idx = stop_gradient(argsort(primals[0], axis_, stream()), stream()); auto out = take_along_axis(tangents[0], sort_idx, axis_, stream()); return {out}; } diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index 2e6b33076c..c85fcf09a6 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -1508,6 +1508,50 @@ def test_complex_abs_grad(self): _, (jvp,) = mx.jvp(mx.abs, [x], [t]) self.assertTrue(mx.allclose(jvp, mx.sign(x) * t)) + def test_second_order_permutation_ops(self): + # The permutation these ops apply is locally constant in the input, so + # the indices must not carry a gradient. Otherwise differentiating the + # vjp a second time fails with "Cannot calculate VJP with respect to + # indices". + def hvp(f, x, v): + return mx.grad(lambda a: mx.sum(mx.grad(f)(a) * v))(x) + + def numerical_hvp(f, x, v, eps=1e-3): + # The values below are well separated, so the permutation does not + # change over this step and the difference is exact enough. + return (mx.grad(f)(x + eps * v) - mx.grad(f)(x - eps * v)) / (2 * eps) + + x = mx.array([3.0, 1.0, 2.0, 5.0]) + v = mx.array([1.0, -2.0, 0.5, 1.5]) + + for fn in ( + lambda a: mx.sort(a), + lambda a: mx.partition(a, 2), + lambda a: mx.topk(a, 2), + lambda a: mx.cummax(a, axis=0), + lambda a: mx.cummin(a, axis=0), + lambda a: mx.cummax(a, axis=0, reverse=True), + lambda a: mx.cummax(a, axis=0, inclusive=False), + lambda a: mx.cummin(a, axis=0, reverse=True, inclusive=False), + ): + f = lambda a: mx.sum(fn(a) ** 2) + self.assertTrue( + mx.allclose(hvp(f, x, v), numerical_hvp(f, x, v), atol=1e-3) + ) + + # A non-trailing axis + y = mx.array([[3.0, 1.0], [2.0, 5.0]]) + w = mx.array([[1.0, -2.0], [0.5, 1.5]]) + for fn in ( + lambda a: mx.sort(a, axis=0), + lambda a: mx.partition(a, 1, axis=0), + lambda a: mx.cummax(a, axis=0), + ): + f = lambda a: mx.sum(fn(a) ** 2) + self.assertTrue( + mx.allclose(hvp(f, y, w), numerical_hvp(f, y, w), atol=1e-3) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 5c12b6d92309340a8285f183cc3a0b614e98ccba Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:53:53 -0700 Subject: [PATCH 144/222] Fix segfault on negative out of bounds axes in take_along_axis/put_along_axis (#4118) --- mlx/linalg.cpp | 8 ++------ mlx/ops.cpp | 22 ++++------------------ python/tests/test_ops.py | 26 ++++++++++++++++++++++++++ 3 files changed, 32 insertions(+), 24 deletions(-) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index 159d2468bc..e86c030c36 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -440,12 +440,8 @@ array cross( int axis /* = -1 */, StreamOrDevice s /* = {} */) { auto check_ax = [axis](const array& arr) { - if (axis >= static_cast(arr.ndim()) || axis + arr.ndim() < 0) { - std::ostringstream msg; - msg << "[linalg::cross] axis " << axis << " invalid for array with " - << arr.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + normalize_axis_index(axis, arr.ndim(), "[linalg::cross] "); if (arr.shape(axis) < 2 || arr.shape(axis) > 3) { throw std::invalid_argument( "[linalg::cross] The specified axis must have size 2 or 3."); diff --git a/mlx/ops.cpp b/mlx/ops.cpp index cee62db438..afc9ab489d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3655,12 +3655,8 @@ array take_along_axis( const array& indices, int axis, StreamOrDevice s /* = {} */) { - if (axis + a.ndim() < 0 || axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[take_along_axis] Received invalid axis for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + axis = normalize_axis_index(axis, a.ndim(), "[take_along_axis] "); if (indices.ndim() != a.ndim()) { std::ostringstream msg; @@ -3669,9 +3665,6 @@ array take_along_axis( throw std::invalid_argument(msg.str()); } - // Allow negative axis - axis = axis < 0 ? a.ndim() + axis : axis; - // Broadcast indices and input ignoring the take axis auto inputs = broadcast_arrays({a, indices}, std::vector{axis - int(a.ndim())}, s); @@ -3693,12 +3686,8 @@ array scatter_axis( StreamOrDevice s) { std::string prefix = (mode == ScatterAxis::None) ? "[put_along_axis]" : "[scatter_add_axis]"; - if (axis + a.ndim() < 0 || axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << prefix << " Received invalid axis for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + axis = normalize_axis_index(axis, a.ndim(), prefix + " "); if (indices.ndim() != a.ndim()) { std::ostringstream msg; @@ -3723,9 +3712,6 @@ array scatter_axis( auto inputs = broadcast_arrays({indices, upd}, s); inputs.insert(inputs.begin(), a); - // Allow negative axis - axis = axis < 0 ? a.ndim() + axis : axis; - // Broadcast src, indices, values while ignoring the take axis inputs = broadcast_arrays(inputs, {axis - int(a.ndim())}, s); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index ac3bb9b49a..a0f4fadaa3 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1429,6 +1429,32 @@ def test_take_along_axis(self): out_mlx = mx.take_along_axis(a_mlx, mx.reshape(idx_mlx, shape), axis=ax) self.assertTrue(np.array_equal(out_np, np.array(out_mlx))) + def test_along_axis_invalid_axis(self): + # Negative out of bounds axes used to slip past the bounds check + # (unsigned arithmetic made it dead code) and either raise an internal + # error or segfault. See also expand_dims in test_expand_dims. + a = mx.arange(24).reshape(2, 3, 4) + idx = mx.zeros(a.shape, dtype=mx.int32) + values = mx.ones(a.shape, dtype=a.dtype) + + for ax in [3, 4, 100, -4, -5, -100]: + with self.assertRaises(ValueError): + mx.take_along_axis(a, idx, axis=ax) + with self.assertRaises(ValueError): + mx.put_along_axis(a, idx, values, axis=ax) + + # Valid negative axes still work + for ax in [-1, -2, -3]: + self.assertEqual(mx.take_along_axis(a, idx, axis=ax).shape, a.shape) + self.assertEqual(mx.put_along_axis(a, idx, values, axis=ax).shape, a.shape) + + def test_cross_invalid_axis(self): + a = mx.array([1.0, 2.0, 3.0]) + b = mx.array([4.0, 5.0, 6.0]) + for ax in [1, 2, -2, -50]: + with self.assertRaises(ValueError): + mx.linalg.cross(a, b, axis=ax) + def test_put_along_axis(self): for ax in [None, 0, 1, 2]: a_np = np.arange(16).reshape(2, 2, 4).astype(np.int32) From fc27acdab6a64f6cec405d06b1b253d33821774b Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:54:53 -0700 Subject: [PATCH 145/222] chore: Deduplicate slow CPU convolution dtype dispatch (#4126) --- mlx/backend/cpu/conv.cpp | 117 +++++++++------------------------------ 1 file changed, 27 insertions(+), 90 deletions(-) diff --git a/mlx/backend/cpu/conv.cpp b/mlx/backend/cpu/conv.cpp index 5d4638adeb..70b5f270f0 100644 --- a/mlx/backend/cpu/conv.cpp +++ b/mlx/backend/cpu/conv.cpp @@ -670,6 +670,24 @@ void slow_conv_3D( }); } +template +void dispatch_slow_conv_type(Dtype dtype, F&& f) { + switch (dtype) { + case float32: + f.template operator()(); + break; + case float16: + f.template operator()(); + break; + case bfloat16: + f.template operator()(); + break; + default: + throw std::invalid_argument( + "[Convolution::eval] got unsupported data type."); + } +} + void dispatch_slow_conv_1D( const array& in, const array& wt, @@ -681,20 +699,8 @@ void dispatch_slow_conv_1D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_1D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_1D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_1D( in, wt, out, @@ -705,22 +711,7 @@ void dispatch_slow_conv_1D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_1D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } void dispatch_slow_conv_2D( @@ -734,20 +725,8 @@ void dispatch_slow_conv_2D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_2D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_2D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_2D( in, wt, out, @@ -758,22 +737,7 @@ void dispatch_slow_conv_2D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_2D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } void dispatch_slow_conv_3D( @@ -787,20 +751,8 @@ void dispatch_slow_conv_3D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_3D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_3D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_3D( in, wt, out, @@ -811,22 +763,7 @@ void dispatch_slow_conv_3D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_3D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } /////////////////////////////////////////////////////////////////////////////// From 63e98d8218d953c8306171cdfccbfe32b541327e Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:55:41 -0700 Subject: [PATCH 146/222] chore: chore: Use dispatch_all_types in print_constant (#4129) --- mlx/backend/common/compiled.cpp | 48 +++++++++++---------------------- 1 file changed, 16 insertions(+), 32 deletions(-) diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp index 5173be421f..9dc5a4c54c 100644 --- a/mlx/backend/common/compiled.cpp +++ b/mlx/backend/common/compiled.cpp @@ -2,46 +2,30 @@ #include "mlx/backend/common/compiled.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/utils.h" namespace mlx::core { void print_constant(std::ostream& os, const array& x) { - switch (x.dtype()) { - case float32: - return print_float_constant(os, x); - case float16: - return print_float_constant(os, x); - case bfloat16: - return print_float_constant(os, x); - case float64: - return print_float_constant(os, x); - case complex64: - return print_complex_constant(os, x); - case int8: + dispatch_all_types(x.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { + os << std::boolalpha << x.item(); + } else if constexpr (std::is_same_v) { os << static_cast(x.item()); - return; - case int16: - return print_int_constant(os, x); - case int32: - return print_int_constant(os, x); - case int64: - return print_int_constant(os, x); - case uint8: + } else if constexpr (std::is_same_v) { os << static_cast(x.item()); - return; - case uint16: - return print_int_constant(os, x); - case uint32: - return print_int_constant(os, x); - case uint64: - return print_int_constant(os, x); - case bool_: - os << std::boolalpha << x.item(); - return; - default: + } else if constexpr (is_complex) { + print_complex_constant(os, x); + } else if constexpr (is_floating_point_v) { + print_float_constant(os, x); + } else if constexpr (std::is_integral_v) { + print_int_constant(os, x); + } else { throw std::runtime_error("Unsupported constant type"); - } + } + }); } std::string get_type_string(Dtype d) { From 0eaaec995dd5dfb4d07e8896a36db5306ba9fa69 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:14:07 -0700 Subject: [PATCH 147/222] chore: Fix set_printoptions example (#4140) --- python/src/print.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/src/print.cpp b/python/src/print.cpp index 03c3fdaa79..6c08d88d3f 100644 --- a/python/src/print.cpp +++ b/python/src/print.cpp @@ -44,7 +44,7 @@ void init_print(nb::module_& m) { Example: >>> print(x) # Uses default precision - >>> mx.set_printoptions(precision=3): + >>> mx.set_printoptions(precision=3) >>> print(x) # Uses precision of 3 >>> print(x) # Uses precision of 3 (again) From 4bc4fec021b2aeca6853db1a976b0e27c3959586 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:14:56 -0700 Subject: [PATCH 148/222] chore: Fix export_function error message (#4141) --- python/src/export.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/src/export.cpp b/python/src/export.cpp index 3596e4f846..7748f87b14 100644 --- a/python/src/export.cpp +++ b/python/src/export.cpp @@ -124,8 +124,8 @@ auto wrap_export_function(nb::callable fun) { outputs_.push_back(nb::cast(outputs)); } else if (!nb::try_cast(outputs, outputs_)) { throw std::invalid_argument( - "[export_function] Outputs can be either a single array " - "a tuple or list of arrays."); + "[export_function] Outputs can be either a single array, " + "a tuple, or a list of arrays."); } return outputs_; }; From f8c45f0d03d9b5ad3799f4fc5f15bbc990a32a79 Mon Sep 17 00:00:00 2001 From: Solaris-star <67425364+Solaris-star@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:23:46 +0800 Subject: [PATCH 149/222] Fix integer keys collision in tree_unflatten (#3878) Signed-off-by: Solaris-star <820622658@qq.com> Co-authored-by: Cheng --- python/mlx/utils.py | 16 ++++++++++++---- python/tests/test_tree.py | 15 +++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/python/mlx/utils.py b/python/mlx/utils.py index bfd167c796..66f132b7ef 100644 --- a/python/mlx/utils.py +++ b/python/mlx/utils.py @@ -229,16 +229,24 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any: next_idx = "" if not next_idx else next_idx[0] children[current_idx].append((next_idx, value)) - # Assume they are a list and fail to dict if the keys are not all integers + # Assume list when all keys are integers. try: - keys = sorted((int(idx), idx) for idx in children.keys()) + keys = {} + for idx in children: + keys[int(idx)] = idx + # Guard against "01" and "1" treated as one key. + is_list = len(keys) == len(children) + except ValueError: + is_list = False + + if is_list: l = [] - for i, k in keys: + for i, k in sorted(keys.items()): # if i <= len(l), no {} will be appended. l.extend([{} for _ in range(i - len(l))]) l.append(tree_unflatten(children[k])) return l - except ValueError: + else: return {k: tree_unflatten(v) for k, v in children.items()} diff --git a/python/tests/test_tree.py b/python/tests/test_tree.py index 171c15b12b..c6f31981b1 100644 --- a/python/tests/test_tree.py +++ b/python/tests/test_tree.py @@ -121,6 +121,21 @@ class Params(NamedTuple): self.assertTrue(mx.array_equal(vector3[0], mx.array([0, 2]))) self.assertTrue(mx.array_equal(vector3[1], mx.array(4))) + def test_tree_unflatten_integer_key_collision(self): + # Non-canonical integer-like keys (e.g. "01") must not silently + # collide with "1" and shift later values. Fall back to dict tree. + tree = mlx.utils.tree_unflatten([("01", "a"), ("1", "b"), ("2", "c")]) + self.assertIsInstance(tree, dict) + self.assertEqual(tree["01"], "a") + self.assertEqual(tree["1"], "b") + self.assertEqual(tree["2"], "c") + + # Canonical list keys still unflatten as a list + self.assertEqual( + mlx.utils.tree_unflatten([("0", "a"), ("1", "b"), ("2", "c")]), + ["a", "b", "c"], + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 7e8b4ccc2a518caa430351ff8fffa625bb674e0a Mon Sep 17 00:00:00 2001 From: Michael Ellis Date: Mon, 10 Aug 2026 23:48:12 -0500 Subject: [PATCH 150/222] Fix fence tracking for donated dynamic slice offsets (#4099) Co-authored-by: Michael Ellis Co-authored-by: Cheng --- mlx/backend/metal/slicing.cpp | 2 +- tests/gpu_tests.cpp | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/slicing.cpp b/mlx/backend/metal/slicing.cpp index e92aef43db..f6ff41abc8 100644 --- a/mlx/backend/metal/slicing.cpp +++ b/mlx/backend/metal/slicing.cpp @@ -58,8 +58,8 @@ array compute_dynamic_offset( offset.copy_shared_buffer(indices); } else { offset.set_data(allocator::malloc(offset.itemsize())); + compute_encoder.add_temporary(offset); } - compute_encoder.add_temporary(offset); auto dtype = indices.dtype(); std::string lib_name = "compute_dynamic_offset_" + type_to_name(dtype); diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index d920da49d9..8bef07a616 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -480,6 +480,32 @@ TEST_CASE("test gpu validation") { eval(scatter_max(array(1), {}, array(2), std::vector{})); } +TEST_CASE("test dynamic slice update waits for its start") { + // Regression for #3880: a donated array-valued start could be read + // stale when a command-buffer boundary lands between its producer + // and the dynamic slice. The boundary occurs when the buffer's op + // or memory limits split the graph (or with MLX_MAX_OPS_PER_BUFFER + // set low); without a split the checks still assert the correct + // update position. + auto source = ones({2, 1 << 26}, int32); + auto target = zeros({4, 4}, int32); + auto update = full({1, 1}, 7, int32); + eval(source, target, update); + + { + auto recycled = zeros({2}, int32); + eval(recycled); + } + + auto out = [&] { + auto start = max(source, 1, false); + return slice_update(target, update, start, {0, 1}); + }(); + + CHECK_EQ(slice(out, {1, 1}, {2, 2}).item(), 7); + CHECK_EQ(slice(out, {0, 0}, {1, 1}).item(), 0); +} + TEST_CASE("test gpu int32 shape overflow errors") { // (2^30, 2).flatten() — product 2^31 doesn't fit in ShapeElem. // Issue #2681 reported wrapped shape (-2147483648,) and a From 199178a09ad01b518c362bbff4eb7bf78a7dbe83 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:24:54 -0700 Subject: [PATCH 151/222] chore: Deduplicate distributed all-reduce (#4144) --- mlx/distributed/ops.cpp | 48 ++++++++++++++--------------------------- 1 file changed, 16 insertions(+), 32 deletions(-) diff --git a/mlx/distributed/ops.cpp b/mlx/distributed/ops.cpp index 583a371edd..3b338dda64 100644 --- a/mlx/distributed/ops.cpp +++ b/mlx/distributed/ops.cpp @@ -20,60 +20,44 @@ Group to_group(std::optional group) { } } -} // namespace - -array all_sum( +array all_reduce( const array& x, - std::optional group_ /* = std::nullopt */, - StreamOrDevice s /* = {} */) { + const std::optional& group_, + StreamOrDevice s, + AllReduce::ReduceType reduce_type) { auto group = to_group(group_); - if (group.size() == 1) { return x; } auto stream = detail::communication_stream(group, s); - return array( x.shape(), x.dtype(), - std::make_shared(stream, group, AllReduce::Sum), + std::make_shared(stream, group, reduce_type), {x}); } -array all_max( +} // namespace + +array all_sum( const array& x, std::optional group_ /* = std::nullopt */, StreamOrDevice s /* = {} */) { - auto group = to_group(group_); - - if (group.size() == 1) { - return x; - } - auto stream = detail::communication_stream(group, s); + return all_reduce(x, group_, s, AllReduce::Sum); +} - return array( - x.shape(), - x.dtype(), - std::make_shared(stream, group, AllReduce::Max), - {x}); +array all_max( + const array& x, + std::optional group_ /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + return all_reduce(x, group_, s, AllReduce::Max); } array all_min( const array& x, std::optional group_ /* = std::nullopt */, StreamOrDevice s /* = {} */) { - auto group = to_group(group_); - - if (group.size() == 1) { - return x; - } - auto stream = detail::communication_stream(group, s); - - return array( - x.shape(), - x.dtype(), - std::make_shared(stream, group, AllReduce::Min), - {x}); + return all_reduce(x, group_, s, AllReduce::Min); } array all_gather( From c6c809b092ba1e41b3cd76c93952d08894feab4e Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:25:44 -0700 Subject: [PATCH 152/222] chore: Deduplicate TCPSocket error throw (#4145) --- mlx/distributed/utils.cpp | 43 ++++++++++++++++----------------------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/mlx/distributed/utils.cpp b/mlx/distributed/utils.cpp index 2de994cf4d..11f8df85c4 100644 --- a/mlx/distributed/utils.cpp +++ b/mlx/distributed/utils.cpp @@ -10,6 +10,16 @@ namespace mlx::core::distributed::detail { +namespace { + +[[noreturn]] void throw_socket_error(const char* tag, const char* action) { + std::ostringstream msg; + msg << tag << " " << action << " (error: " << errno << ")"; + throw std::runtime_error(msg.str()); +} + +} // namespace + /** * Parse a sockaddr from an ip and port provided as strings. */ @@ -53,9 +63,7 @@ address_t parse_address(const std::string& ip_port) { TCPSocket::TCPSocket(const char* tag) { sock_ = socket(AF_INET, SOCK_STREAM, 0); if (sock_ < 0) { - std::ostringstream msg; - msg << tag << " Couldn't create socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't create socket"); } } @@ -95,40 +103,30 @@ void TCPSocket::listen(const char* tag, const address_t& addr) { int enable = 1; success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't enable reuseaddr"); } success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int)); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't enable reuseport (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't enable reuseport"); } // Bind the socket to the address and port success = bind(sock_, addr.get(), addr.len); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't bind socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't bind socket"); } // Prepare waiting for connections success = ::listen(sock_, 0); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't listen (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't listen"); } } TCPSocket TCPSocket::accept(const char* tag) { int peer = ::accept(sock_, nullptr, nullptr); if (peer < 0) { - std::ostringstream msg; - msg << tag << " Accept failed (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Accept failed"); } return TCPSocket(peer); @@ -173,10 +171,7 @@ TCPSocket TCPSocket::connect( // Create the socket sock = socket(AF_INET, SOCK_STREAM, 0); if (sock < 0) { - std::ostringstream msg; - msg << tag << " Couldn't create socket to connect (error: " << errno - << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't create socket to connect"); } success = ::connect(sock, addr.get(), addr.len); @@ -195,9 +190,7 @@ TCPSocket TCPSocket::connect( } if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't connect (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't connect"); } return TCPSocket(sock); From 191d72464657adac72482dfbeb97d4f060ffff91 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:26:26 -0700 Subject: [PATCH 153/222] chore: Simplify BitwiseBinary::eval_gpu (#4147) --- mlx/backend/metal/binary.cpp | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/mlx/backend/metal/binary.cpp b/mlx/backend/metal/binary.cpp index 51771e9b1b..7d6b2a7131 100644 --- a/mlx/backend/metal/binary.cpp +++ b/mlx/backend/metal/binary.cpp @@ -235,23 +235,7 @@ BINARY_GPU(Power) BINARY_GPU(Subtract) void BitwiseBinary::eval_gpu(const std::vector& inputs, array& out) { - switch (op_) { - case BitwiseBinary::And: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::Or: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::Xor: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::LeftShift: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::RightShift: - binary_op_gpu(inputs, out, name()); - break; - } + binary_op_gpu(inputs, out, name()); } } // namespace mlx::core From 4be2e0e82ef36c478fa16926bc61f60d0c1e1418 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:27:30 -0700 Subject: [PATCH 154/222] chore: Deduplicate metal_kernel validation (#4149) --- mlx/backend/common/metal_kernel.cpp | 29 ++++++++++------------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index 43ed8c2047..7a2eef51fc 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -283,27 +283,18 @@ CustomKernelFunction metal_kernel( std::optional init_value = std::nullopt, bool verbose = false, StreamOrDevice s_ = {}) { - if (inputs.size() != input_names.size()) { - std::ostringstream msg; - msg << "[metal_kernel] Expected `inputs` to have size " - << input_names.size() << " but got size " << inputs.size() << "." - << std::endl; - throw std::invalid_argument(msg.str()); - } - if (output_shapes.size() != output_names.size()) { - std::ostringstream msg; - msg << "[metal_kernel] Expected `output_shapes` to have size " - << output_names.size() << " but got size " << output_shapes.size() - << "." << std::endl; - throw std::invalid_argument(msg.str()); - } - if (output_dtypes.size() != output_names.size()) { + auto check_size = [](size_t actual, size_t expected, const char* name) { + if (actual == expected) { + return; + } std::ostringstream msg; - msg << "[metal_kernel] Expected `output_dtypes` to have size " - << output_names.size() << " but got size " << output_dtypes.size() - << "." << std::endl; + msg << "[metal_kernel] Expected `" << name << "` to have size " + << expected << " but got size " << actual << "." << std::endl; throw std::invalid_argument(msg.str()); - } + }; + check_size(inputs.size(), input_names.size(), "inputs"); + check_size(output_shapes.size(), output_names.size(), "output_shapes"); + check_size(output_dtypes.size(), output_names.size(), "output_dtypes"); auto s = resolve_metal_kernel_stream(s_); From c489fe282959f33742c2584f5c85a9db3ed2fa02 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:28:16 -0700 Subject: [PATCH 155/222] chore: Simplify dtype_to_compute_type (#4151) --- mlx/backend/cuda/gemms/cublas_gemm.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/mlx/backend/cuda/gemms/cublas_gemm.cpp b/mlx/backend/cuda/gemms/cublas_gemm.cpp index a790d946a5..e05595b8f3 100644 --- a/mlx/backend/cuda/gemms/cublas_gemm.cpp +++ b/mlx/backend/cuda/gemms/cublas_gemm.cpp @@ -15,17 +15,14 @@ namespace { cublasComputeType_t dtype_to_compute_type(Dtype dtype) { switch (dtype) { case float16: - return CUBLAS_COMPUTE_32F; case bfloat16: return CUBLAS_COMPUTE_32F; case float32: + case complex64: return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; case float64: return CUBLAS_COMPUTE_64F; - case complex64: - return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32 - : CUBLAS_COMPUTE_32F; default: throw std::runtime_error( fmt::format( From d4936ab7dd6fa56c3761b1ada098612270bbf27b Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:29:47 -0700 Subject: [PATCH 156/222] chore: Use dispatch_all_types in Scatter::eval_cpu (#4155) --- mlx/backend/cpu/indexing.cpp | 48 +++--------------------------------- 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/indexing.cpp b/mlx/backend/cpu/indexing.cpp index fc44888a3f..c7e24fd75e 100644 --- a/mlx/backend/cpu/indexing.cpp +++ b/mlx/backend/cpu/indexing.cpp @@ -430,50 +430,10 @@ void Scatter::eval_cpu(const std::vector& inputs, array& out) { updates = array::unsafe_weak_copy(updates), inds = std::move(inds), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint8: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int8: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case bfloat16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case complex64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + dispatch_scatter(out, inds, updates, axes_, reduce_type_); + }); }); } From f66bcc9c99b7354f7a5db147621f5a366c1c9937 Mon Sep 17 00:00:00 2001 From: Ishaan Samantray Date: Tue, 11 Aug 2026 02:31:26 -0400 Subject: [PATCH 157/222] chore: Improve LUF/SVD error messages (#4167) --- mlx/backend/cpu/luf.cpp | 4 ++-- mlx/backend/cpu/svd.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cpu/luf.cpp b/mlx/backend/cpu/luf.cpp index e7e34c4d57..2e396f4b70 100644 --- a/mlx/backend/cpu/luf.cpp +++ b/mlx/backend/cpu/luf.cpp @@ -69,8 +69,8 @@ void luf_impl( if (info < 0) { std::stringstream ss; - ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info - << " because argument had an illegal value"; + ss << "[LUF::eval_cpu] LU factorization failed with error code " + << info << " because argument had an illegal value"; throw std::runtime_error(ss.str()); } diff --git a/mlx/backend/cpu/svd.cpp b/mlx/backend/cpu/svd.cpp index 932ededdb4..80101e9072 100644 --- a/mlx/backend/cpu/svd.cpp +++ b/mlx/backend/cpu/svd.cpp @@ -91,7 +91,7 @@ struct SVDWork< if (info != 0) { std::stringstream ss; - ss << "svd_impl: sgesvdx_ failed with code " << info; + ss << "[SVD::eval_cpu] SVD failed with error code " << info; throw std::runtime_error(ss.str()); } } @@ -181,7 +181,7 @@ struct SVDWork> { if (info != 0) { std::stringstream ss; - ss << "svd_impl: sgesvdx_ failed with code " << info; + ss << "[SVD::eval_cpu] SVD failed with error code " << info; throw std::runtime_error(ss.str()); } } From 067462b81a52df252e745b9faae92a7f96afb8dc Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:32:03 -0700 Subject: [PATCH 158/222] chore: Use dispatch_all_types in ArgPartition::eval_cpu (#4156) --- mlx/backend/cpu/sort.cpp | 33 +++------------------------------ 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index d8a6f3559b..81020ac1de 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -407,36 +407,9 @@ void ArgPartition::eval_cpu(const std::vector& inputs, array& out) { out = array::unsafe_weak_copy(out), axis_ = axis_, kth_ = kth_]() mutable { - switch (in.dtype()) { - case bool_: - return argpartition(in, out, axis_, kth_); - case uint8: - return argpartition(in, out, axis_, kth_); - case uint16: - return argpartition(in, out, axis_, kth_); - case uint32: - return argpartition(in, out, axis_, kth_); - case uint64: - return argpartition(in, out, axis_, kth_); - case int8: - return argpartition(in, out, axis_, kth_); - case int16: - return argpartition(in, out, axis_, kth_); - case int32: - return argpartition(in, out, axis_, kth_); - case int64: - return argpartition(in, out, axis_, kth_); - case float32: - return argpartition(in, out, axis_, kth_); - case float64: - return argpartition(in, out, axis_, kth_); - case float16: - return argpartition(in, out, axis_, kth_); - case bfloat16: - return argpartition(in, out, axis_, kth_); - case complex64: - return argpartition(in, out, axis_, kth_); - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + argpartition(in, out, axis_, kth_); + }); }); } From a9ab0f695a0af07b2460b96d2ca368ad8e5c9333 Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Tue, 11 Aug 2026 12:02:52 +0530 Subject: [PATCH 159/222] chore: Add complex to python scalar type (#4168) --- python/mlx/_stub_patterns.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt index 2a3fbce133..974ce0c7a5 100644 --- a/python/mlx/_stub_patterns.txt +++ b/python/mlx/_stub_patterns.txt @@ -7,7 +7,7 @@ mlx.core.__prefix__: __dlpack_device__: Callable[..., Any] mlx.core.__suffix__: - scalar: TypeAlias = int | float | bool + scalar: TypeAlias = int | float | bool | complex list_or_scalar: TypeAlias = scalar | list["list_or_scalar"] StreamOrDevice: TypeAlias = Stream | ThreadLocalStream | Device | DeviceType | None bool_: Dtype = ... From 249ea075eb05279f521e89224ace7e297b2e0764 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:33:44 -0700 Subject: [PATCH 160/222] chore: Use dispatch_all_types in copy_inplace_dispatch (#4170) --- mlx/backend/cpu/copy.cpp | 48 ++++------------------------------------ 1 file changed, 4 insertions(+), 44 deletions(-) diff --git a/mlx/backend/cpu/copy.cpp b/mlx/backend/cpu/copy.cpp index 6b6d825252..381719eb0c 100644 --- a/mlx/backend/cpu/copy.cpp +++ b/mlx/backend/cpu/copy.cpp @@ -209,50 +209,10 @@ inline void copy_inplace_dispatch( array& dst, CopyType ctype, Args&&... args) { - switch (src.dtype()) { - case bool_: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint8: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint16: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint32: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint64: - copy(src, dst, ctype, std::forward(args)...); - break; - case int8: - copy(src, dst, ctype, std::forward(args)...); - break; - case int16: - copy(src, dst, ctype, std::forward(args)...); - break; - case int32: - copy(src, dst, ctype, std::forward(args)...); - break; - case int64: - copy(src, dst, ctype, std::forward(args)...); - break; - case float16: - copy(src, dst, ctype, std::forward(args)...); - break; - case float32: - copy(src, dst, ctype, std::forward(args)...); - break; - case float64: - copy(src, dst, ctype, std::forward(args)...); - break; - case bfloat16: - copy(src, dst, ctype, std::forward(args)...); - break; - case complex64: - copy(src, dst, ctype, std::forward(args)...); - break; - } + dispatch_all_types(src.dtype(), [&](auto type_tag) { + using SrcT = MLX_GET_TYPE(type_tag); + copy(src, dst, ctype, std::forward(args)...); + }); } } // namespace From 9d47fa63f85a4dc70692caf711509ef1773dc4ea Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:35:37 -0700 Subject: [PATCH 161/222] chore: Simplify scan_gpu_inplace (#4154) --- mlx/backend/metal/scan.cpp | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/mlx/backend/metal/scan.cpp b/mlx/backend/metal/scan.cpp index cd5184aa3e..18935191ff 100644 --- a/mlx/backend/metal/scan.cpp +++ b/mlx/backend/metal/scan.cpp @@ -59,14 +59,14 @@ void scan_gpu_inplace( auto kernel = get_scan_kernel(d, kname, reverse, inclusive, reduce_type_str, in, out); - if (contiguous) { - 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); - size_t size = in.shape(axis); - compute_encoder.set_bytes(size, 2); + 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); + size_t size = in.shape(axis); + compute_encoder.set_bytes(size, 2); + if (contiguous) { // Compute the thread grid int n_reads = (in.itemsize() <= 4) ? 4 : 2; constexpr int simd_size = 32; @@ -89,15 +89,9 @@ void scan_gpu_inplace( MTL::Size group_dims(thread_group_size, 1, 1); compute_encoder.dispatch_threads(grid_dims, group_dims); } else { - 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); - size_t size = in.shape(axis); size_t stride = in.strides()[axis]; int bn = 32; size_t stride_blocks = (stride + bn - 1) / bn; - compute_encoder.set_bytes(size, 2); compute_encoder.set_bytes(stride, 3); compute_encoder.set_bytes(stride_blocks, 4); From 66a040789e779a0c8e7585987a4fd3f832794fd8 Mon Sep 17 00:00:00 2001 From: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:37:45 +0800 Subject: [PATCH 162/222] Skip empty NAX GEMM output groups (#3941) --- .../metal/kernels/steel/gemm/gemm_nax.h | 10 ++++++++++ .../steel/gemm/kernels/steel_gemm_fused_nax.h | 19 +++++++++++-------- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h b/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h index 40066be7a4..c42aa1306f 100644 --- a/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h @@ -45,11 +45,17 @@ auto gemm_loop( NAXTile Dtile; Dtile.clear(); + const bool has_output = sgp_sm > 0 && sgp_sn > 0; + int gemm_k_iterations_ = gemm_k_iterations_aligned; STEEL_PRAGMA_NO_UNROLL for (int kk0 = 0; kk0 < gemm_k_iterations_; kk0++) { threadgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + continue; + } STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { @@ -94,6 +100,10 @@ auto gemm_loop( if constexpr (!kAlignedK) { simdgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + return Dtile; + } const short rem_bk = K - gemm_k_iterations_ * BK; diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h index 80d90333d7..8267707ade 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h @@ -200,14 +200,17 @@ template < params->gemm_k_iterations_aligned, sgp_sm, sgp_sn); - if (use_out_source) { - gemm_epilogue( - Dtile, C, params, addmm_params, sgp_sm, sgp_sn); - } - if constexpr (kAlignedM && kAlignedN) { - Dtile.store(D, int(params->ldd)); - } else { - Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + if ((kAlignedM.value || sgp_sm > 0) && + (kAlignedN.value || sgp_sn > 0)) { + if (use_out_source) { + gemm_epilogue( + Dtile, C, params, addmm_params, sgp_sm, sgp_sn); + } + if constexpr (kAlignedM && kAlignedN) { + Dtile.store(D, int(params->ldd)); + } else { + Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + } } }); }); From cd5b92cb4f4bffabe0e61b7bdcbf83268db384f9 Mon Sep 17 00:00:00 2001 From: "Brian C." <94733710+deBrian07@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:28:11 -0700 Subject: [PATCH 163/222] docs: fix stale CLI invocations in LLM inference example (#4172) --- docs/src/examples/llama-inference.rst | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/src/examples/llama-inference.rst b/docs/src/examples/llama-inference.rst index 7e06895e35..cd6e0b561b 100644 --- a/docs/src/examples/llama-inference.rst +++ b/docs/src/examples/llama-inference.rst @@ -322,14 +322,14 @@ several unnecessary copies from disk to numpy and then from numpy to MLX. It will be replaced in the future with direct loading to MLX. You can download the full example code in `mlx-examples`_. Assuming, the -existence of ``weights.pth`` and ``tokenizer.model`` in the current working -directory we can play around with our inference script as follows (the timings -are representative of an M1 Ultra and the 7B parameter Llama model): +existence of the PyTorch Llama weights in ``llama-7B/`` we can play around +with our inference script as follows (the timings are representative of an M1 +Ultra and the 7B parameter Llama model): .. code-block:: bash - $ python convert.py weights.pth llama-7B.mlx.npz - $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely' + $ python convert.py --torch-path llama-7B/ + $ python llama.py --prompt 'Call me Ishmael. Some years ago never mind how long precisely' [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ @@ -347,7 +347,7 @@ time as well as the prompt processing time remains almost constant. .. code-block:: bash - $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' + $ python llama.py --prompt 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ @@ -355,7 +355,7 @@ time as well as the prompt processing time remains almost constant. ------ [INFO] Prompt processing: 0.579 s [INFO] Full generation: 4.690 s - $ python llama.py --num-tokens 500 llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' + $ python llama.py --max-tokens 500 --prompt 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' [INFO] Loading model from disk: 5.628 s Press enter to start generation ------ From 596dc79f43429d2de91f867fc377581dcc47f862 Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 11 Aug 2026 19:34:11 +0900 Subject: [PATCH 164/222] Increase ccache size in CI (#3999) --- .github/actions/setup/action.yml | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index f0344a427e..4a1a27cf9c 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -92,15 +92,16 @@ runs: key: v7-${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.ccache-toolkit || inputs.toolkit }} max-size: |- ${{ case(inputs.ccache-key == 'release', - case(startsWith(inputs.toolkit, 'cuda'), '150MB', + case(startsWith(inputs.toolkit, 'cuda'), + case(runner.os == 'Linux', '300MB', + '160MB'), '60MB'), case(startsWith(inputs.toolkit, 'cuda'), - case(runner.os == 'Linux' && runner.arch == 'x64', '500MB', - runner.os == 'Linux', '240MB', - '120MB'), - runner.os == 'macOS' && inputs.toolkit == 'metal', '200MB', - runner.os == 'macOS', '150MB', - '60MB')) + case(runner.os == 'Linux' && runner.arch == 'x64', '600MB', + runner.os == 'Linux', '400MB', + '320MB'), + runner.os == 'macOS' && inputs.toolkit == 'metal', '300MB', + '200MB')) }} save: ${{ !startsWith(github.ref, 'refs/pull/') && (inputs.ccache-save != 'false') }} # ccache-action bug: running "apt-get update" fails on large arm runner. From 9562ea7f9ce75b1764485d6f6ead90aa37d7b105 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Isaac=20Hern=C3=A1ndez?= Date: Tue, 11 Aug 2026 20:15:32 -0300 Subject: [PATCH 165/222] Fix metal hadamard_transform for n = m with no power-of-2 factor (#4054) Co-authored-by: Cheng --- mlx/backend/metal/hadamard.cpp | 41 ++++++++++++++++++---------------- python/tests/test_ops.py | 27 ++++++++++++++++++++++ 2 files changed, 49 insertions(+), 19 deletions(-) diff --git a/mlx/backend/metal/hadamard.cpp b/mlx/backend/metal/hadamard.cpp index 223c6f3f5f..c696e4d0ff 100644 --- a/mlx/backend/metal/hadamard.cpp +++ b/mlx/backend/metal/hadamard.cpp @@ -69,7 +69,8 @@ void hadamard_mn_contiguous( int n = n1 * n2; int read_width_n1 = n1 == 2 ? 2 : 4; int read_width_n2 = n2 == 2 ? 2 : 4; - int read_width_m = (n == 2 || m == 28) ? 2 : 4; + // The m stage strides by n / read_width_m, so cap the read width at n. + int read_width_m = (n == 1) ? 1 : ((n == 2 || m == 28) ? 2 : 4); int max_radix_1 = std::min(n1, 16); int max_radix_2 = std::min(n2, 16); float scale_n1 = 1.0; @@ -97,17 +98,16 @@ void hadamard_mn_contiguous( auto lib = d.get_library(kname, [&]() { std::string kernel; concatenate( - kernel, - metal::utils(), - gen_hadamard_codelet(m), - metal::hadamard(), - get_template_definition( - "n2" + kname, - "hadamard_n", - get_type_string(x.dtype()), - n2, - max_radix_2, - read_width_n2)); + kernel, metal::utils(), gen_hadamard_codelet(m), metal::hadamard()); + if (n2 > 1) { + kernel += get_template_definition( + "n2" + kname, + "hadamard_n", + get_type_string(x.dtype()), + n2, + max_radix_2, + read_width_n2); + } if (n1 > 1) { kernel += get_template_definition( "n1" + kname, @@ -143,18 +143,21 @@ void hadamard_mn_contiguous( // Launch the transform for n2 auto& compute_encoder = metal::get_command_encoder(s); - auto kernel = d.get_kernel("n2" + kname, lib); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(n1 > 1 ? y : x, 0); - compute_encoder.set_output_array(y, 1); - compute_encoder.set_bytes(scale_n2, 2); - compute_encoder.dispatch_threads(grid_dims_n2, group_dims_n2); + if (n2 > 1) { + auto kernel = d.get_kernel("n2" + kname, lib); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(n1 > 1 ? y : x, 0); + compute_encoder.set_output_array(y, 1); + compute_encoder.set_bytes(scale_n2, 2); + compute_encoder.dispatch_threads(grid_dims_n2, group_dims_n2); + } // Launch the strided transform for m if (m > 1) { auto kernel = d.get_kernel("m" + kname, lib); compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(y, 0); + // With n == 1 no earlier stage copied x into y, so read from x. + compute_encoder.set_input_array(n > 1 ? y : x, 0); compute_encoder.set_output_array(y, 1); compute_encoder.set_bytes(scale_m, 2); compute_encoder.dispatch_threads(grid_dims_m, group_dims_m); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a0f4fadaa3..2571c58557 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3481,6 +3481,33 @@ def parse_h_string(h_str): y_bf16.astype(mx.float16), y, atol=atol * 2 ) + @unittest.skipIf(not mx.metal.is_available(), "Metal only") + def test_hadamard_m_only(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # n = m * 2^0, so only the m stage runs. test_hadamard sweeps k from 1. + tests = product( + (12, 20, 28), # m + (None, 0.25), # scale + ) + for m, scale in tests: + for shape in ((m,), (4, m), (3, 5, m)): + with self.subTest(m=m, shape=shape, scale=scale): + x = mx.array( + np.random.RandomState(3).normal(size=shape).astype(np.float32) + ) + kwargs = {} if scale is None else {"scale": scale} + y_cpu = mx.hadamard_transform(x, stream=mx.cpu, **kwargs) + y_gpu = mx.hadamard_transform(x, stream=mx.gpu, **kwargs) + mx.eval(y_cpu, y_gpu) + self.assertEqual(y_gpu.shape, x.shape) + self.assertLess(mx.abs(y_cpu - y_gpu).max().item(), 1e-5) + # non-donatable input: the malloc'd-output path + y_nd = mx.hadamard_transform(x + 0.0, stream=mx.gpu, **kwargs) + mx.eval(y_nd) + self.assertLess(mx.abs(y_cpu - y_nd).max().item(), 1e-5) + def test_hadamard_grad_vmap(self): np.random.seed(4) From ec3ad74d7862223c375f1d218c10295371d386f6 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:24:21 -0700 Subject: [PATCH 166/222] chore: Use dispatch_all_types in Partition::eval_cpu (#4175) --- mlx/backend/cpu/sort.cpp | 33 +++------------------------------ 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index 81020ac1de..767e7ef1d3 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -428,36 +428,9 @@ void Partition::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([out = array::unsafe_weak_copy(out), axis_ = axis_, kth_ = kth_]() mutable { - switch (out.dtype()) { - case bool_: - return partition(out, axis_, kth_); - case uint8: - return partition(out, axis_, kth_); - case uint16: - return partition(out, axis_, kth_); - case uint32: - return partition(out, axis_, kth_); - case uint64: - return partition(out, axis_, kth_); - case int8: - return partition(out, axis_, kth_); - case int16: - return partition(out, axis_, kth_); - case int32: - return partition(out, axis_, kth_); - case int64: - return partition(out, axis_, kth_); - case float32: - return partition(out, axis_, kth_); - case float64: - return partition(out, axis_, kth_); - case float16: - return partition(out, axis_, kth_); - case bfloat16: - return partition(out, axis_, kth_); - case complex64: - return partition(out, axis_, kth_); - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + partition(out, axis_, kth_); + }); }); } From 9e37e5b48ef254872e9c1d3864ef3e3a1590bb97 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:25:38 -0700 Subject: [PATCH 167/222] chore: Use dispatch_all_types in ScatterAxis::eval_cpu (#4176) --- mlx/backend/cpu/indexing.cpp | 99 +++--------------------------------- 1 file changed, 8 insertions(+), 91 deletions(-) diff --git a/mlx/backend/cpu/indexing.cpp b/mlx/backend/cpu/indexing.cpp index c7e24fd75e..15a2af3cb9 100644 --- a/mlx/backend/cpu/indexing.cpp +++ b/mlx/backend/cpu/indexing.cpp @@ -555,53 +555,10 @@ void ScatterAxis::eval_cpu(const std::vector& inputs, array& out) { idx = array::unsafe_weak_copy(idx), updates = array::unsafe_weak_copy(updates), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint8: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint16: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int8: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int16: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case float16: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - case float32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case float64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case bfloat16: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - case complex64: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); + }); }); } @@ -662,50 +619,10 @@ void MaskedScatter::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([mask = array::unsafe_weak_copy(mask), src = array::unsafe_weak_copy(src), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - masked_scatter_impl(mask, src, out); - break; - case uint8: - masked_scatter_impl(mask, src, out); - break; - case uint16: - masked_scatter_impl(mask, src, out); - break; - case uint32: - masked_scatter_impl(mask, src, out); - break; - case uint64: - masked_scatter_impl(mask, src, out); - break; - case int8: - masked_scatter_impl(mask, src, out); - break; - case int16: - masked_scatter_impl(mask, src, out); - break; - case int32: - masked_scatter_impl(mask, src, out); - break; - case int64: - masked_scatter_impl(mask, src, out); - break; - case float16: - masked_scatter_impl(mask, src, out); - break; - case float32: - masked_scatter_impl(mask, src, out); - break; - case float64: - masked_scatter_impl(mask, src, out); - break; - case bfloat16: - masked_scatter_impl(mask, src, out); - break; - case complex64: - masked_scatter_impl(mask, src, out); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + masked_scatter_impl(mask, src, out); + }); }); } From fb0818fc3b38c7ddf6d59e9b7175e5fe2570143c Mon Sep 17 00:00:00 2001 From: Kolja Wawrowsky <3075215+apocryphx@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:27:40 -0700 Subject: [PATCH 168/222] docs: document softmax's `precise` argument (#4178) Co-authored-by: Claude Opus 5 --- python/src/ops.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 0941793949..05c3735953 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3278,7 +3278,7 @@ void init_ops(nb::module_& m) { "precise"_a = false, "stream"_a = nb::none(), nb::sig( - "def softmax(a: array, /, axis: None | int | Sequence[int] = None, *, stream: StreamOrDevice = None) -> array"), + "def softmax(a: array, /, axis: None | int | Sequence[int] = None, *, precise: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the softmax along the given axis. @@ -3293,6 +3293,10 @@ void init_ops(nb::module_& m) { axis (int or list(int), optional): Optional axis or axes to compute the softmax over. If unspecified this performs the softmax over the full array. + precise (bool, optional): Accumulate in ``float32`` for inputs of + lower precision. Otherwise the accumulation type matches the + input, which can lose precision over long reduction axes. + Default: ``False``. Returns: array: The output of the softmax. From 64d392d61b3bd6aa8deec847af84ba545bb62f4a Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 11 Aug 2026 16:28:31 -0700 Subject: [PATCH 169/222] Send out of range trig arguments to libm (#4157) --- mlx/backend/cpu/simd/math.h | 22 ++++++++++++++++++++-- python/tests/test_ops.py | 7 +++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cpu/simd/math.h b/mlx/backend/cpu/simd/math.h index f9fc8317a5..88676323a5 100644 --- a/mlx/backend/cpu/simd/math.h +++ b/mlx/backend/cpu/simd/math.h @@ -115,12 +115,30 @@ Simd sincos(Simd in) { } } +// sincos reduces the argument by rounding x * 4/pi into a uint32, which is +// only exact while that product fits in a float's mantissa. Past that the +// result degrades and then leaves [-1, 1] altogether, so send those arguments +// to libm instead. 2^23 keeps x * 4/pi under 2^24 with room to spare. +template +Simd sincos_checked(Simd x) { + Simd xf = x; + if (any(abs(xf) > Simd(8388608.0f))) { + Simd out; + for (int i = 0; i < N; ++i) { + float v = xf[i]; + out[i] = static_cast(Sine ? std::sin(v) : std::cos(v)); + } + return out; + } + return sincos(x); +} + template Simd sin(Simd x) { if constexpr (is_complex) { return std::sin(x.value); } else { - return sincos(x); + return sincos_checked(x); } } @@ -129,7 +147,7 @@ Simd cos(Simd x) { if constexpr (is_complex) { return std::cos(x.value); } else { - return sincos(x); + return sincos_checked(x); } } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 2571c58557..e2da815bd8 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1192,6 +1192,13 @@ def test_sin(self): self.assertTrue(np.allclose(result, expected)) + # Large arguments still have to land in [-1, 1] + big = np.array([1e8, 1e9, 1e10, 1e20, 1e30], dtype=np.float32) + for op, npop in [(mx.sin, np.sin), (mx.cos, np.cos)]: + out = np.array(op(mx.array(big))) + self.assertTrue(np.all(np.abs(out) <= 1.0)) + self.assertTrue(np.allclose(out, npop(big), atol=1e-6)) + def test_cos(self): a = mx.array( [0, math.pi / 4, math.pi / 2, math.pi, 3 * math.pi / 4, 2 * math.pi] From 9f35f77427a681a99125d886925f1cd9ca745181 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Tue, 11 Aug 2026 19:48:19 -0400 Subject: [PATCH 170/222] Add reflect and symmetric padding modes to mx.pad (#3608) Co-authored-by: Cheng Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com> --- ACKNOWLEDGMENTS.md | 1 + mlx/ops.cpp | 85 ++++++++++++++++++++++++++++++++++++++++ python/src/ops.cpp | 4 +- python/tests/test_ops.py | 33 ++++++++++++++++ tests/ops_tests.cpp | 43 ++++++++++++++++++++ 5 files changed, 165 insertions(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 186908f09c..d83832be8f 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -20,6 +20,7 @@ MLX was developed with contributions from the following individuals: - Paul Paczuski: Improved stability of BCE loss calculation - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops. - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. +- katlun-lgtm: Added `reflect` and `symmetric` padding modes. diff --git a/mlx/ops.cpp b/mlx/ops.cpp index afc9ab489d..7cc283e399 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1460,6 +1460,85 @@ array tile( return reshape(x, std::move(final_shape), s); } +array reflect_pad( + const array& a, + const std::vector& axes, + const Shape& low_pad_size, + const Shape& high_pad_size, + const Shape& out_shape, + bool include_edge, + StreamOrDevice s /* = {} */) { + array out = zeros(out_shape, a.dtype(), s); + Shape starts(a.ndim(), 0); + auto stops = a.shape(); + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + starts[ax] = low_pad_size[i]; + stops[ax] += low_pad_size[i]; + } + // Copy over values from the unpadded array + array padded = slice_update(out, a, starts, stops, s); + + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + int n = a.shape(ax); + int L = low_pad_size[i]; + int H = high_pad_size[i]; + if (L == 0 && H == 0) { + continue; + } + // reflect skips the edge value (period 2(n-1)); symmetric repeats it + // (period 2n). + int offset = (!include_edge && n > 1) ? 1 : 0; + int tile = n - offset; + + if (L > 0) { + int filled_start = low_pad_size[i]; + int remaining = L; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_start + offset; + src_stops[ax] = filled_start + offset + chunk; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_start - chunk; + dst_stops[ax] = filled_start; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_start -= chunk; + remaining -= chunk; + } + } + + if (H > 0) { + int filled_end = low_pad_size[i] + n; + int remaining = H; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_end - offset - chunk; + src_stops[ax] = filled_end - offset; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_end; + dst_stops[ax] = filled_end + chunk; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_end += chunk; + remaining -= chunk; + } + } + } + return padded; +} + array edge_pad( const array& a, const std::vector& axes, @@ -1552,6 +1631,12 @@ array pad( {a, astype(pad_value, a.dtype(), s)}); } else if (mode == "edge") { return edge_pad(a, axes, low_pad_size, high_pad_size, out_shape, s); + } else if (mode == "reflect") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, false, s); + } else if (mode == "symmetric") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, true, s); } else { std::ostringstream msg; msg << "Invalid padding mode (" << mode << ") passed to pad"; diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 05c3735953..a666588f51 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3513,7 +3513,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), + "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge', 'reflect', 'symmetric'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Pad an array with a constant value @@ -3528,6 +3528,8 @@ void init_ops(nb::module_& m) { mode: Padding mode. One of the following strings: "constant" (default): Pads with a constant value. "edge": Pads with the edge values of array. + "reflect": Pads with the reflection of the array, without repeating the edge values. + "symmetric": Pads with the reflection of the array, repeating the edge values. constant_values (array or scalar, optional): Optional constant value to pad the edges of the array with. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index e2da815bd8..c6a8245ebe 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2272,6 +2272,39 @@ def test_nan_to_num(self): out_mx = mx.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000) self.assertTrue(np.allclose(out_mx, out_np)) + def test_pad_reflect_symmetric(self): + # mx.pad reflect/symmetric must match numpy.pad exactly. Covers + # in-bounds, multi-reflect (pad larger than the axis, exercising the + # tiling loop), asymmetric + # per-axis widths, zero-width sides, and degenerate axes (n == 1, n == 2). + cases = [ + ((8,), [(2, 3)]), + ((8,), [(0, 4)]), + ((8,), [(3, 0)]), + ((8,), [(7, 8)]), + ((4,), [(10, 7)]), # multi-reflect + ((4,), [(20, 20)]), # multi-reflect, both sides + ((3,), [(9, 1)]), # multi-reflect + ((1,), [(3, 2)]), # degenerate axis + ((2,), [(5, 6)]), # smallest non-trivial, multi-reflect + ((5, 6), [(2, 3), (1, 2)]), + ((5, 6), [(9, 9), (11, 0)]), # both axes multi-reflect + ((3, 4, 5), [(1, 1), (0, 0), (2, 2)]), + ((3, 4, 5), [(4, 4), (0, 0), (7, 3)]), + ] + for mode in ("reflect", "symmetric"): + for shape, pw in cases: + a_npy = np.random.randn(*shape).astype(np.float32) + a_mlx = mx.array(a_npy) + b_npy = np.pad(a_npy, pw, mode=mode) + b_mlx = mx.pad(a_mlx, pw, mode=mode) + self.assertEqual(b_mlx.shape, tuple(b_npy.shape)) + self.assertTrue( + np.array_equal(np.array(b_mlx), b_npy), + msg=f"mismatch mode={mode} shape={shape} pad={pw}", + ) + self.assertEqual(b_mlx.dtype, mx.float32) + def test_as_strided(self): x_npy = np.random.randn(128).astype(np.float32) x_mlx = mx.array(x_npy) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 0dd9385e14..f7a2b8ab92 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -2978,6 +2978,49 @@ TEST_CASE("test pad") { 0.0f}, {4, 4}); CHECK(array_equal(padded_x, expected).item()); + + // reflect padding (mirror without repeating the edge value) + x = array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {5}); + CHECK(array_equal( + pad(x, {{2, 2}}, array(0.0f), "reflect"), + array({3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 4.0f, 3.0f}, {9})) + .item()); + CHECK(array_equal( + pad(x, {{0, 3}}, array(0.0f), "reflect"), + array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 4.0f, 3.0f, 2.0f}, {8})) + .item()); + + // symmetric padding (mirror repeating the edge value) + CHECK(array_equal( + pad(x, {{2, 2}}, array(0.0f), "symmetric"), + array({2.0f, 1.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 5.0f, 4.0f}, {9})) + .item()); + CHECK(array_equal( + pad(x, {{3, 0}}, array(0.0f), "symmetric"), + array({3.0f, 2.0f, 1.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {8})) + .item()); + + // multi-reflect: pad larger than the axis repeats the reflection (numpy + // parity) + x = array({1.0f, 2.0f, 3.0f}, {3}); + CHECK(array_equal( + pad(x, {{5, 5}}, array(0.0f), "reflect"), + array( + {2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f}, + {13})) + .item()); } TEST_CASE("test power") { From 2ddcb6f7232d051797928a3000ed76d1d3f1d717 Mon Sep 17 00:00:00 2001 From: Kolja Wawrowsky <3075215+apocryphx@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:45:17 -0700 Subject: [PATCH 171/222] docs: Writing a Fast KV Cache (#4019) Co-authored-by: Claude Opus 5 Co-authored-by: Cheng --- docs/src/index.rst | 1 + docs/src/usage/kv_cache.rst | 91 +++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 docs/src/usage/kv_cache.rst diff --git a/docs/src/index.rst b/docs/src/index.rst index e29d396a89..8b3028be70 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -39,6 +39,7 @@ are the CPU and GPU. usage/lazy_evaluation usage/unified_memory usage/indexing + usage/kv_cache usage/saving_and_loading usage/function_transforms usage/compile diff --git a/docs/src/usage/kv_cache.rst b/docs/src/usage/kv_cache.rst new file mode 100644 index 0000000000..f24666cf46 --- /dev/null +++ b/docs/src/usage/kv_cache.rst @@ -0,0 +1,91 @@ +.. _kv_cache: + +Writing a Fast KV Cache +======================= + +Autoregressive generation appends one position to the key and value arrays for +each generated token. The append strategy can dominate KV-cache performance. + +Avoid appending naively with :func:`concatenate`: + +.. code-block:: python + + # Avoid this + cache = mx.zeros((1, 0, d)) + for x in steps: + cache = mx.concatenate([cache, x], axis=1) + mx.eval(cache) + +Instead, preallocate fixed-size chunks and update the cache in place: + +.. code-block:: python + + chunk = 256 + cache = mx.zeros((1, chunk, d)) + offset = 0 + for x in steps: + if offset == cache.shape[1]: + cache = mx.concatenate([cache, mx.zeros((1, chunk, d))], axis=1) + cache = mx.slice_update(cache, x, mx.array(offset), (1,)) + offset += 1 + mx.eval(cache) + + keys = cache[:, :offset] + +You can also use indexed assignment, which may read more naturally: +``cache[:, offset : offset + 1, :] = x``. + +The following measurements append one position per step to 20 ``bfloat16`` +caches of shape ``[1, 4, N, 512]`` on an M4 Max: + +.. list-table:: + :widths: 20 30 30 + :header-rows: 1 + + * - Context + - Concatenate + - Preallocate + update + * - 512 + - 0.90 ms / step + - 0.24 ms / step + * - 1024 + - 1.11 ms / step + - 0.21 ms / step + * - 4096 + - 3.73 ms / step + - 0.22 ms / step + +With preallocation, step time remains nearly constant as context length grows. +Concatenation becomes progressively slower. + +Why Concatenating Is Slow +------------------------- + +Concatenation has two costs: copying data and preventing buffer reuse. + +:func:`concatenate` creates a new array and copies the existing cache at every +step. Appending ``n`` positions therefore copies on the order of ``n^2`` +elements. + +Growing the cache also prevents buffer reuse. MLX pools freed device buffers, +but reuses a buffer only for a similarly sized request. It does not split a +larger buffer for a smaller request or combine smaller buffers for a larger +one. Because a growing cache has a new size at every step, each allocation +typically comes from the driver while freed buffers remain unused. + +This allocation work occurs on the CPU. In profiles, it appears as GPU idle +time between kernels rather than as a slow kernel, which can make the model +appear to be the bottleneck. + +Preallocation avoids both costs. The cache shape remains fixed, eliminating +both repeated copies and per-step allocations. + +Choosing a Chunk Size +--------------------- + +Use a chunk size that is a multiple of 256. + +This both amortizes growth and enables the fused cuDNN attention kernel on +CUDA. For single-token attention, the key and value arrays must be slices of a +contiguous cache with a capacity that is a multiple of 256, and at least 256 +positions must be in use. Other chunk sizes silently use a slower path. From d04e5dbff89eb0f1522de71810d9cd7b7ddc05b2 Mon Sep 17 00:00:00 2001 From: Andrew Geyko <83419363+gordofreemo@users.noreply.github.com> Date: Wed, 12 Aug 2026 02:58:32 +0200 Subject: [PATCH 172/222] Fix and enable non-transposed NAX qmm (#4051) Co-authored-by: Cheng --- mlx/backend/metal/kernels/quantized_nax.h | 33 +++++++++------- mlx/backend/metal/quantized.cpp | 5 ++- python/tests/test_quantized.py | 47 +++++++++++++++++++++++ 3 files changed, 70 insertions(+), 15 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index e67be9a06d..81f6a7a252 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1094,7 +1094,6 @@ METAL_FUNC void qmm_n_nax_tgp_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { (void)lid; - (void)M; static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -1115,23 +1114,20 @@ METAL_FUNC void qmm_n_nax_tgp_impl( bits>; // Set the block - const int K_w = K * bytes_per_pack / pack_factor; - const int K_g = K / group_size; const int y_row = tid.y * BM; const int y_col = tid.x * BN; auto wl = (const device uint8_t*)w; + // Here w is [K, N]: packed and group-quantized along N, with row stride N. x += y_row * static_cast(K); - wl += y_col * K_w; - scales += y_col * K_g; - biases += y_col * K_g; + wl += y_col * bytes_per_pack / pack_factor; + scales += y_col / group_size; + biases += y_col / group_size; y += y_row * static_cast(N) + y_col; - // Make the x loader and mma operation - // const short num_els = min(BM, M - y_row); - // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + // Make the weight loader + loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -1144,6 +1140,8 @@ METAL_FUNC void qmm_n_nax_tgp_impl( const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); + const short sgp_sm = min(int(SM), M - (y_row + tm)); + const short ldb_tgp = BN_padded; constexpr bool transpose_a = false; @@ -1168,7 +1166,12 @@ METAL_FUNC void qmm_n_nax_tgp_impl( volatile int compiler_barrier; - Atile.load(x + kk1, K); + if (sgp_sm == SM) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + Btile.template load(Ws + tn + kk1 * ldb_tgp); tile_matmad_nax( @@ -1188,7 +1191,11 @@ METAL_FUNC void qmm_n_nax_tgp_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - Dtile.store(y + tm * N + tn, N); + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm)); + } } template < @@ -1678,4 +1685,4 @@ template < }); }); } -} \ No newline at end of file +} diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index b754132dc4..ab3e1deacc 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1035,8 +1035,9 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { - if (metal::is_nax_available() && transpose && (K % 64 == 0) && - (env::enable_tf32() || x.dtype() != float32)) { + // The non-transposed kernel requires N % 64 == 0. + if (metal::is_nax_available() && (transpose || (N % 64 == 0)) && + (K % 64 == 0) && (env::enable_tf32() || x.dtype() != float32)) { return qmm_nax( /* const array& x = */ x, /* const array& w = */ w, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 4ca32f2279..56ba75db32 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -353,6 +353,53 @@ def test_qmm_large_dims(self): tol = 1e-3 if dtype == mx.float32 else 1.5e-3 self.assertLess((y_q - y_hat).abs().max(), tol) + def test_qmm_non_transposed(self): + # The non-transposed matmul (w is [K, N]) is reachable mainly from the + # vjp of a quantized linear layer, so it gets much less coverage than + # the transposed one. Sweep it over transformer-sized K/N and over M + # values that leave a partial M-tile. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 + + def check(M, K, N, group_size, bits, batch=()): + x = mx.random.normal(shape=(*batch, M, K), key=k1) / K**0.5 + w = mx.random.normal(shape=(K, N), key=k2) / K**0.5 + x = x.astype(dtype) + w = w.astype(dtype) + w_q, scales, biases = mx.quantize(w, group_size, bits) + w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) + y_q = mx.quantized_matmul(x, w_q, scales, biases, False, group_size, bits) + y_hat = x @ w_hat + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), tol) + + # M sweep. 33..63 is the interesting range: a whole simdgroup of the + # threadgroup's M-tile falls past the end of the matrix. + for M in [1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129]: + for group_size, bits in [(64, 4), (128, 4), (64, 8)]: + with self.subTest(M=M, group_size=group_size, bits=bits): + check(M, 512, 1024, group_size, bits) + + # Transformer-sized K/N, aligned and unaligned M. + for K, N in [(2048, 2048), (512, 2048), (2048, 512), (11008, 2048)]: + for M in [100, 256]: + with self.subTest(shape=(M, K, N)): + check(M, K, N, 64, 4) + + # Batched x, unaligned M. + for batch in [(2,), (2, 3)]: + for M in [33, 250]: + with self.subTest(batch=batch, M=M): + check(M, 512, 1024, 64, 4, batch=batch) + + # M > 2**15 with a partial M-tile, so the per-simdgroup row count is a + # distance that does not fit in an int16. Same failure mode as the one + # test_qmm_large_dims covers for the transposed kernel. + with self.subTest(shape=(33000, 128, 64)): + check(33000, 128, 64, 64, 4) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From 31b5cbb0933e6d3adc49f148946cd2c3ccb9c292 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 11 Aug 2026 18:00:48 -0700 Subject: [PATCH 173/222] Add scatter reduce for JACCL (#3901) --- mlx/backend/cpu/distributed.cpp | 19 +- mlx/distributed/jaccl/jaccl.cpp | 11 +- mlx/distributed/jaccl/lib/jaccl/group.h | 10 + mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 81 +++++--- mlx/distributed/jaccl/lib/jaccl/mesh.h | 15 +- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 211 +++++++++++++++++++- mlx/distributed/jaccl/lib/jaccl/ring.cpp | 40 ++++ mlx/distributed/jaccl/lib/jaccl/ring.h | 10 + mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 193 ++++++++++++++++++ 9 files changed, 550 insertions(+), 40 deletions(-) diff --git a/mlx/backend/cpu/distributed.cpp b/mlx/backend/cpu/distributed.cpp index 22dc4b4cc8..6b788cde25 100644 --- a/mlx/backend/cpu/distributed.cpp +++ b/mlx/backend/cpu/distributed.cpp @@ -98,6 +98,23 @@ void Recv::eval_cpu( void ReduceScatter::eval_cpu( const std::vector& inputs, std::vector& outputs) { - throw std::runtime_error("[ReduceScatter] Not implemented yet."); + assert(inputs.size() == 0); + assert(outputs.size() == 1); + + auto [in, copied] = ensure_row_contiguous(inputs[0], stream()); + outputs[0].set_data(allocator::malloc(outputs[0].nbytes())); + + switch (reduce_type_) { + case Sum: + distributed::detail::sum_scatter(group(), in, outputs[0], stream()); + break; + default: + throw std::runtime_error("Only scatter sum is supported for now"); + } + if (copied) { + auto& enc = cpu::get_command_encoder(stream()); + enc.add_temporary(in); + } } + } // namespace mlx::core::distributed diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 3d4b8e070e..99a4ac3be3 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -148,7 +148,16 @@ class JACCLGroup : public GroupImpl { } void sum_scatter(const array& input, array& output, Stream stream) override { - throw std::runtime_error("[jaccl] sum_scatter not supported."); + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = output.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->sum_scatter(in_ptr, out_ptr, n_bytes, dtype); + }); } std::shared_ptr split(int color, int key = -1) override { diff --git a/mlx/distributed/jaccl/lib/jaccl/group.h b/mlx/distributed/jaccl/lib/jaccl/group.h index 758109c256..603765a511 100644 --- a/mlx/distributed/jaccl/lib/jaccl/group.h +++ b/mlx/distributed/jaccl/lib/jaccl/group.h @@ -28,6 +28,16 @@ class Group { virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0; + /** + * Reduce scatter with a sum reduction. + * + * The input holds size() contiguous chunks of n_bytes each (total + * size() * n_bytes bytes). After the call, output (n_bytes bytes) on rank r + * contains the elementwise sum over all ranks of the r-th input chunk. + */ + virtual void + sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) = 0; + virtual void send(const void* input, size_t n_bytes, int dst) = 0; virtual void recv(void* output, size_t n_bytes, int src) = 0; virtual void barrier() = 0; diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 2409e12754..5b749f2bd9 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -28,15 +28,7 @@ MeshGroup::MeshGroup( side_channel_.barrier(); // Create the mesh implementation object - mesh_ = MeshImpl(rank_, size_, connections_, buffers_); - ring_ = RingImpl( - rank_, - size_, - &connections_[(rank_ + size_ - 1) % size_], - &connections_[(rank_ + 1) % size_], - 1, - ring_send_buffers_, - ring_recv_buffers_); + mesh_ = MeshImpl(rank_, size_, connections_, buffers_, scatter_buffers_); } void MeshGroup::initialize() { @@ -83,8 +75,7 @@ void MeshGroup::initialize() { void MeshGroup::allocate_buffers() { // Deregister any buffers and free the memory buffers_.clear(); - ring_send_buffers_.clear(); - ring_recv_buffers_.clear(); + scatter_buffers_.clear(); // Allocate the memory for (int k = 0; k < BUFFER_SIZES; k++) { @@ -93,10 +84,9 @@ void MeshGroup::allocate_buffers() { for (int j = 0; j < size_; j++) { buffers_.emplace_back(FRAME_SIZE * (1 << k)); } - // Ring buffers (1 for each direction) - for (int j = 0; j < 2; j++) { - ring_send_buffers_.emplace_back(FRAME_SIZE * (1 << k)); - ring_recv_buffers_.emplace_back(FRAME_SIZE * (1 << k)); + // Scatter buffers (size_ send slots followed by size_ recv slots) + for (int j = 0; j < 2 * size_; j++) { + scatter_buffers_.emplace_back(FRAME_SIZE * (1 << k)); } } } @@ -122,18 +112,19 @@ void MeshGroup::allocate_buffers() { } } - // Ring buffers (see ring group for the logic below) - int left = (rank_ + size_ - 1) % size_; - int right = (rank_ + 1) % size_; - // We register send buffers to both the right and the left. - ring_send_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0] - .register_to_protection_domain(connections_[right].protection_domain); - ring_recv_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0] - .register_to_protection_domain(connections_[left].protection_domain); - ring_send_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 1] - .register_to_protection_domain(connections_[left].protection_domain); - ring_recv_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 1] - .register_to_protection_domain(connections_[right].protection_domain); + // Scatter buffers. Slot p (send to peer p) and slot size_ + p (recv from + // peer p) are both registered to peer p's protection domain. The slots + // for our own rank are unused but kept for uniform indexing. + int scatter_base = k * NUM_BUFFERS * 2 * size_ + i * 2 * size_; + for (int j = 0; j < size_; j++) { + if (j == rank_) { + continue; + } + scatter_buffers_[scatter_base + j].register_to_protection_domain( + connections_[j].protection_domain); + scatter_buffers_[scatter_base + size_ + j] + .register_to_protection_domain(connections_[j].protection_domain); + } } } } @@ -176,6 +167,17 @@ void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) { static_cast(input), static_cast(output), n_bytes); } +void MeshGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + void MeshGroup::send(const void* input, size_t n_bytes, int dst) { mesh_.send(static_cast(input), n_bytes, dst); } @@ -198,13 +200,30 @@ void MeshGroup::all_reduce( auto in_ptr = static_cast(input); auto out_ptr = static_cast(output); int64_t count = n_bytes / sizeof(T); - if (size_ > 2 && - ((std::is_same_v && count > 256 * 1024) || - count >= 8 * 1024 * 1024 / static_cast(sizeof(T)))) { - ring_.all_reduce<2>(in_ptr, out_ptr, count, 1, reduce_op); + if (size_ > 2 && n_bytes > 32 * 1024) { + // Large messages are bandwidth bound so use the reduce scatter + all gather + // path which moves size_x less data per link than the fully connected + // all_reduce. + mesh_.all_reduce_scatter_gather(in_ptr, out_ptr, count, reduce_op); } else { + // Small messages are latency bound so use the single phase fully + // connected all_reduce cause it is a bit better. mesh_.all_reduce(in_ptr, out_ptr, count, reduce_op); } } +template +void MeshGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + // n_bytes is the size of the output (one chunk). The input holds size_ such + // chunks laid out contiguously. + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + mesh_.sum_scatter(in_ptr, out_ptr, count, reduce_op); +} + } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.h b/mlx/distributed/jaccl/lib/jaccl/mesh.h index 14823db873..8e5164aee2 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.h @@ -5,7 +5,6 @@ #include "jaccl/group.h" #include "jaccl/mesh_impl.h" #include "jaccl/rdma.h" -#include "jaccl/ring_impl.h" namespace jaccl { @@ -44,6 +43,9 @@ class MeshGroup : public Group { void all_gather(const void* input, void* output, size_t n_bytes) override; + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + void send(const void* input, size_t n_bytes, int dst) override; void recv(void* output, size_t n_bytes, int src) override; @@ -57,6 +59,13 @@ class MeshGroup : public Group { size_t n_bytes, ReduceOp reduce_op); + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + /** * Performs the connection initialization. Namely, after this call all * Connection objects should have a queue pair in RTS state and all buffers @@ -74,11 +83,9 @@ class MeshGroup : public Group { SideChannel side_channel_; std::vector connections_; std::vector buffers_; - std::vector ring_send_buffers_; - std::vector ring_recv_buffers_; + std::vector scatter_buffers_; MeshImpl mesh_; - RingImpl ring_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 6327f41d35..75c540e917 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -19,11 +19,13 @@ class MeshImpl { int rank, int size, std::vector& conns, - std::vector& buffers) + std::vector& buffers, + std::vector& scatter_buffers) : rank_(rank), size_(size), connections_(conns), buffers_(buffers), + scatter_buffers_(scatter_buffers), staging_mem_( std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)) {} @@ -200,8 +202,12 @@ class MeshImpl { } void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); + // Copy our data to the appropriate place. Skip when in place (the scatter + // gather all reduce passes our own reduced shard which already lives at + // out_ptr + rank_ * n_bytes). + if (in_ptr != out_ptr + rank_ * n_bytes) { + std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); + } // Fully connected all gather char* data = out_ptr; @@ -279,6 +285,163 @@ class MeshImpl { } } + template + void sum_scatter(const T* in, T* out, int64_t count, ReduceOp reduce_op) { + // Fully connected reduce scatter. + // + // The input holds size_ contiguous chunks of `count` elements each. Every + // rank keeps chunk rank_ and needs the elementwise reduction of that chunk + // across all ranks. To that end each rank p sends its chunk j to rank j and + // receives every peer's chunk rank_. The output is seeded with this rank's + // own chunk rank_ and every received chunk is reduced into it. + // + // Unlike all_reduce/all_gather each peer receives a *different* chunk, so + // we use the dedicated scatter buffers: per (sz, buff) tile there are size_ + // send slots (slot p -> peer p) and size_ recv slots (slot p <- peer p). + + const T* our_chunk = in + static_cast(rank_) * count; + + auto [sz, buffer_size] = buffer_size_from_message(count * sizeof(T)); + int64_t N = buffer_size / sizeof(T); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; + int64_t total = static_cast(count); + int num_peers = size_ - 1; + + // Seed the output with our own chunk. + std::copy_n(our_chunk, total, out); + + // Counters to maintain the state of transfers + int in_flight = 0; + int64_t read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int64_t write_offset[MESH_MAX_PEERS] = {0}; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + int64_t elems = std::min(N, total - read_offset); + scatter_post_recv_all(sz, buff); + // Stage the chunk destined for each peer p (our input chunk p) into that + // peer's send buffer. + for (int p = 0; p < size_; p++) { + if (p == rank_) { + continue; + } + const T* src = in + static_cast(p) * count + read_offset; + std::copy( + src, src + elems, scatter_send_buffer(sz, buff, p).begin()); + } + scatter_post_send_all(sz, buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + // Send completed. Once every peer received this buffer, refill it with + // the next slice of each peer's chunk and post the next sends. + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + int64_t elems = std::min(N, total - read_offset); + for (int p = 0; p < size_; p++) { + if (p == rank_) { + continue; + } + const T* src = in + static_cast(p) * count + read_offset; + std::copy( + src, + src + elems, + scatter_send_buffer(sz, buff, p).begin()); + } + scatter_post_send_all(sz, buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + // Recv completed. Reduce the peer's contribution into our chunk and, if + // there is more data to fetch from that peer, post another recv. + else if (work_type == RECV_WR) { + int64_t elems = std::min(N, total - write_offset[rank]); + reduce_op( + scatter_recv_buffer(sz, buff, rank).begin(), + out + write_offset[rank], + elems); + write_offset[rank] += N; + if (write_offset[rank] + N * (PIPELINE - 1) < total) { + scatter_recv_from(sz, rank, buff); + in_flight++; + } + } + } + } + } + + template + void all_reduce_scatter_gather( + const T* in, + T* out, + int64_t count, + ReduceOp reduce_op) { + // Bandwidth optimal all reduce for large messages: a reduce scatter + // followed by an all gather. Compared to the fully connected all_reduce + // (which sends every rank's whole input to every peer) this moves size_x + // less data per link at the cost of an extra communication phase, so it is + // preferred for large messages where bandwidth dominates latency. + // + // The input is split into size_ equal chunks of `chunk` elements. The + // reduce scatter reduces chunk rank_ across all ranks and leaves it at + // out + rank_ * chunk. The all gather then reads that reduced shard in + // place and distributes every rank's reduced chunk into the output. No + // intermediate buffer is needed: the shard already lives where the all + // gather expects our contribution. + // + // Any trailing elements that do not divide evenly (fewer than size_) are + // handled with the fully connected all_reduce on the tail. + + int64_t chunk = count / size_; + int64_t base = chunk * size_; + + if (chunk > 0) { + // Reduce scatter our chunk directly into its final location in the + // output. sum_scatter only reads `in` and writes out + rank_ * chunk so + // this is safe even when in aliases out. + T* shard = out + static_cast(rank_) * chunk; + sum_scatter(in, shard, chunk, reduce_op); + + // All gather every rank's reduced chunk into the output. Our own shard is + // already in place so all_gather skips the self copy and only fills the + // other ranks' slices, which held the (already sent) input chunks. + all_gather( + reinterpret_cast(shard), + reinterpret_cast(out), + chunk * sizeof(T)); + } + + // Reduce the trailing elements (fewer than size_) with the fully connected + // all reduce so every rank ends up with the same tail. + if (base < count) { + all_reduce(in + base, out + base, count - base, reduce_op); + } + } + void send(const char* in_ptr, int64_t n_bytes, int dst) { constexpr int PIPELINE = 2; constexpr int WC_NUM = PIPELINE; @@ -417,10 +580,52 @@ class MeshImpl { } } + // Scatter buffer helpers. Per (sz, buff) tile there are 2 * size_ slots: the + // first size_ are send buffers (slot p is sent to peer p) and the next size_ + // are recv buffers (slot p receives from peer p). + SharedBuffer& scatter_send_buffer(int sz, int buff, int peer) { + return scatter_buffers_ + [sz * NUM_BUFFERS * 2 * size_ + buff * 2 * size_ + peer]; + } + + SharedBuffer& scatter_recv_buffer(int sz, int buff, int peer) { + return scatter_buffers_ + [sz * NUM_BUFFERS * 2 * size_ + buff * 2 * size_ + size_ + peer]; + } + + void scatter_send_to(int sz, int rank, int buff) { + connections_[rank].post_send( + scatter_send_buffer(sz, buff, rank), SEND_WR << 16 | buff << 8 | rank); + } + + void scatter_recv_from(int sz, int rank, int buff) { + connections_[rank].post_recv( + scatter_recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank); + } + + void scatter_post_send_all(int sz, int buff) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + scatter_send_to(sz, i, buff); + } + } + + void scatter_post_recv_all(int sz, int buff) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + scatter_recv_from(sz, i, buff); + } + } + int rank_; int size_; std::span connections_; std::span buffers_; + std::span scatter_buffers_; std::unique_ptr staging_mem_; }; diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index 7fbfac04ca..a08035f3f3 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -168,6 +168,17 @@ void RingGroup::all_gather(const void* input, void* output, size_t n_bytes) { n_conns_); } +void RingGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + void RingGroup::send(const void* input, size_t n_bytes, int dst) { int right = (rank_ + 1) % size_; int left = (rank_ + size_ - 1) % size_; @@ -215,4 +226,33 @@ void RingGroup::all_reduce( ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, n_conns_, reduce_op); } +template +void RingGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + // n_bytes is the size of the output (one chunk). The input holds size_ such + // chunks laid out contiguously so the full element count is size_ * count. + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + if (size_ == 1) { + std::copy_n(in_ptr, count, out_ptr); + return; + } + + int64_t total = static_cast(size_) * count; + + // Mirror the all_reduce heuristics: use a single wire for small messages and + // scale up to n_conns_ wires for large ones where bandwidth dominates. + if (total < size_ * 2 * n_conns_ || n_bytes <= 65536) { + ring_.reduce_scatter(in_ptr, out_ptr, total, 1, reduce_op); + return; + } + + ring_.reduce_scatter( + in_ptr, out_ptr, total, n_conns_, reduce_op); +} + } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h index 94799dd267..c901c15a85 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring.h @@ -45,6 +45,9 @@ class RingGroup : public Group { void all_gather(const void* input, void* output, size_t n_bytes) override; + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + void send(const void* input, size_t n_bytes, int dst) override; void recv(void* output, size_t n_bytes, int src) override; @@ -58,6 +61,13 @@ class RingGroup : public Group { size_t n_bytes, ReduceOp reduce_op); + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + /** * Performs the connection initialization. Namely, after this call all * Connection objects should have a queue pair in RTS state and all buffers diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 412b178dc3..ad05f7b208 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -240,6 +240,199 @@ class RingImpl { CopyOp{}); } + // Standalone ring reduce scatter. + // + // The input holds size_ contiguous chunks of `chunk = size / size_` elements + // (size is guaranteed divisible by size_ for a sum_scatter). Rank r produces + // the full reduction of chunk r in out_ptr (chunk elements). Compared to the + // reduce scatter phase embedded in the ring all_reduce this converges every + // rank on its *own* chunk index so it can be exposed directly, and it writes + // straight into the chunk sized output using it as the rolling accumulator. + template + void reduce_scatter( + const T* in_ptr, + T* out_ptr, + int64_t size, + int n_wires, + ReduceOp reduce_op) { + int64_t chunk = size / size_; + // Two directional regions, each split across the wires. + int64_t size_per_wire = (chunk + (2 * n_wires) - 1) / (2 * n_wires); + + dispatch_wires(n_wires, [&](int lw) { + reduce_scatter_wire( + in_ptr, out_ptr, chunk, size_per_wire, n_wires, lw, reduce_op); + }); + } + + // Perform the dual direction ring reduce scatter for a single wire lw. + // + // The output chunk is divided into 2 directional regions and each region is + // split into n_wires contiguous slices of size_per_wire elements. Direction 0 + // flows left (send right / recv left), direction 1 flows right. Both + // directions converge on chunk rank_, so together they fill the whole output + // chunk. This wire owns slice lw of every region and only touches + // left_[lw] / right_[lw]. + // + // The rolling partial for each direction lives in the chunk sized output + // itself: every step sends the partial produced by the previous step and the + // recv writes the next partial into the same location. To avoid clobbering a + // partial that is still being staged for sending, a recv slice is only + // reduced into the output once the matching send slice has been staged (i.e. + // recv_count[lr] < send_count[lr]). + template + void reduce_scatter_wire( + const T* in_ptr, + T* out_ptr, + int64_t chunk, + int64_t size_per_wire, + int n_wires, + int lw, + ReduceOp reduce_op) { + constexpr int MAX_DIR = 2; + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * 2 * MAX_DIR; + + auto [sz, buffer_bytes] = + buffer_size_from_message(size_per_wire * sizeof(T)); + int64_t N = buffer_bytes / sizeof(T); + int64_t total = static_cast(size_) * chunk; + + // This wire's element offset within the output chunk in each direction and + // the end of each direction's region. + int64_t wire_offset[MAX_DIR]; + int64_t region_end[MAX_DIR]; + for (int lr = 0; lr < MAX_DIR; lr++) { + wire_offset[lr] = lr * n_wires * size_per_wire + + static_cast(lw) * size_per_wire; + region_end[lr] = std::min(chunk, (lr + 1) * n_wires * size_per_wire); + } + + // Input windows (count space, chunk aligned). Both directions converge on + // chunk rank_ (verified: rank r finishes chunk r). Direction 0 flows left + // and rotates backward; direction 1 flows right and rotates forward. + int64_t in_send_offset[MAX_DIR]; + int64_t in_recv_offset[MAX_DIR]; + in_send_offset[0] = ((rank_ - 1 + size_) % size_) * chunk; + in_recv_offset[0] = ((rank_ - 2 + 2 * size_) % size_) * chunk; + in_send_offset[1] = ((rank_ + 1) % size_) * chunk; + in_recv_offset[1] = ((rank_ + 2) % size_) * chunk; + + int64_t n_steps = (size_per_wire + N - 1) / N; + + int in_flight = 0; + int send_count[MAX_DIR] = {0}; + int recv_count[MAX_DIR] = {0}; + // Recv completions that arrived before their matching send slice was staged + // are deferred until the send catches up so we never overwrite a partial + // that is still being read for sending. + int deferred_recv[MAX_DIR] = {0}; + + // Apply as many deferred recv slices for direction lr as are now unblocked + // by staged sends. Recv completions arrive in FIFO order per queue pair, so + // recv slice recv_count[lr] lives in pipeline buffer recv_count[lr] % + // PIPELINE. + auto drain_deferred = [&](int lr) { + while (deferred_recv[lr] > 0 && recv_count[lr] < send_count[lr]) { + int slice = recv_count[lr]; + int b = slice % PIPELINE; + int64_t offset = wire_offset[lr] + static_cast(slice) * N; + int64_t n = std::min(N, region_end[lr] - offset); + reduce_op( + recv_buffer(sz, b, lr, lw).template begin(), + in_ptr + in_recv_offset[lr] + offset, + out_ptr + offset, + std::max(0, n)); + recv_count[lr]++; + deferred_recv[lr]--; + if (recv_count[lr] + (PIPELINE - 1) < n_steps) { + recv_from(sz, b, lr, lw); + in_flight++; + } + } + }; + + for (int k = 0; k < size_ - 1; k++) { + // Step 0 forwards this rank's own input; later steps forward the partial + // accumulated in the output by the previous step. + const T* send_base = (k == 0) ? in_ptr : out_ptr; + // For step 0 the send window is in the input (count space); for later + // steps the partial lives at wire_offset in the chunk sized output. + int64_t send_base_offset[MAX_DIR]; + for (int lr = 0; lr < MAX_DIR; lr++) { + send_base_offset[lr] = (k == 0) ? in_send_offset[lr] : 0; + } + + // Prefill the pipeline + int buff = 0; + while (buff < n_steps && buff < PIPELINE) { + for (int lr = 0; lr < MAX_DIR; lr++) { + recv_from(sz, buff, lr, lw); + } + for (int lr = 0; lr < MAX_DIR; lr++) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_base_offset[lr] + offset, + send_base + send_base_offset[lr] + + std::max(offset, std::min(offset + N, region_end[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); + } + + buff++; + in_flight += 2 * MAX_DIR; + } + + // Main loop + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll_wire(lw, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int lr = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR) { + if (send_count[lr] < n_steps) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_base_offset[lr] + offset, + send_base + send_base_offset[lr] + + std::max(offset, std::min(offset + N, region_end[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); + in_flight++; + } + // A newly staged send may unblock deferred recvs. + drain_deferred(lr); + } + + else if (work_type == RECV_WR) { + // Only reduce into the output once the matching send slice has been + // staged; otherwise defer until send_count catches up. + deferred_recv[lr]++; + drain_deferred(lr); + } + } + } + + // Advance the input windows around the ring for the next step and reset + // the per step counters. + in_send_offset[0] = (in_send_offset[0] + total - chunk) % total; + in_recv_offset[0] = (in_recv_offset[0] + total - chunk) % total; + in_send_offset[1] = (in_send_offset[1] + chunk) % total; + in_recv_offset[1] = (in_recv_offset[1] + chunk) % total; + for (int lr = 0; lr < MAX_DIR; lr++) { + send_count[lr] = recv_count[lr] = 0; + deferred_recv[lr] = 0; + } + } + } + // Run size_ - 1 pipelined ring steps for a single wire in every direction. // // At every step each direction sends its current slice to a neighbor and From 56c26e83e2479398118c65053e55545f81d71851 Mon Sep 17 00:00:00 2001 From: Ishaan Samantray Date: Tue, 11 Aug 2026 22:47:08 -0400 Subject: [PATCH 174/222] Fix diag for zero-size input (#4165) --- mlx/ops.cpp | 5 +++++ python/tests/test_ops.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 7cc283e399..ce29cb4e3d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -6312,6 +6312,11 @@ array diag(const array& a, int k /* = 0 */, StreamOrDevice s /* = {} */) { int a_size = a.size(); int n = a_size + std::abs(k); auto res = zeros({n, n}, a.dtype(), s); + if (a_size == 0) { + // Nothing to place on the diagonal, and scattering into the 0x0 output + // produced when k is zero is an error. + return res; + } std::vector indices; auto s1 = std::max(0, -k); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index c6a8245ebe..d569dbd3c7 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3169,6 +3169,27 @@ def test_diag(self): expected = mx.array(np.diag(x, k=-1)) self.assertTrue(mx.array_equal(result, expected)) + def test_diag_zero_size(self): + # A zero-size 1-D input builds a |k| x |k| matrix of zeros. k = 0 makes + # that 0 x 0, which used to fail while every other k worked. + for k in (-2, -1, 0, 1, 2): + for dtype, nptype in ( + (mx.float32, np.float32), + (mx.int32, np.int32), + (mx.complex64, np.complex64), + ): + result = mx.diag(mx.zeros((0,), dtype=dtype), k=k) + expected = np.diag(np.zeros((0,), dtype=nptype), k=k) + self.assertEqual(result.shape, expected.shape, msg=f"k={k} {dtype}") + self.assertEqual(result.dtype, dtype) + self.assertTrue(np.array_equal(np.array(result), expected)) + + # A zero-size 2-D input already worked; keep it covered. + for shape in ((0, 0), (0, 3), (3, 0)): + result = mx.diag(mx.zeros(shape)) + expected = np.diag(np.zeros(shape, dtype=np.float32)) + self.assertEqual(result.shape, expected.shape, msg=f"{shape}") + def test_trace(self): a_mx = mx.arange(9, dtype=mx.int64).reshape((3, 3)) a_np = np.arange(9, dtype=np.int64).reshape((3, 3)) From 13e79684d71fd46084883ddc20381d7c68623255 Mon Sep 17 00:00:00 2001 From: Franco De E Date: Tue, 11 Aug 2026 21:37:25 -0600 Subject: [PATCH 175/222] Add an inverse-CDF path to categorical sampling (#4177) Co-authored-by: Cheng --- mlx/random.cpp | 30 ++++++++++++++++++++++++++++++ python/tests/test_random.py | 15 +++++++++++++++ tests/random_tests.cpp | 13 +++++++++++++ 3 files changed, 58 insertions(+) diff --git a/mlx/random.cpp b/mlx/random.cpp index 164d5f2441..2932b3eaf8 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -388,12 +388,42 @@ int get_valid_axis(int axis, int ndim) { return ax; } +// O(N + M) in memory, where the Gumbel-max trick needs O(N * M) +array categorical_inverse_cdf( + const array& logits, + const Shape& shape, + const std::optional& key, + StreamOrDevice s) { + auto dtype = promote_types(logits.dtype(), float32); + auto x = astype(logits, dtype, s); + auto m = max(x, s); + + // exp(x - m) is NaN when m is infinite, where the distribution is uniform + // over the maximal entries + auto w = where( + isinf(m, s), + astype(equal(x, m, s), dtype, s), + exp(subtract(x, m, s), s), + s); + + auto cdf = cumsum(w, 0, /* reverse = */ false, /* inclusive = */ false, s); + auto u = multiply(uniform(shape, float32, key, s), sum(w, s), s); + + return subtract(searchsorted(cdf, u, "right", s), array(1u, uint32), s); +} + array categorical_impl( const array& logits, int axis, const Shape& shape, const std::optional& key /*= nullopt */, StreamOrDevice s) { + // searchsorted only takes 1D sequence. + auto n = logits.shape(axis); + if (n > 0 && logits.size() == static_cast(n) && !shape.empty()) { + return categorical_inverse_cdf(reshape(logits, {n}, s), shape, key, s); + } + auto gumbel_shape = shape; auto offset = axis + shape.size() - logits.ndim() + 1; gumbel_shape.insert(gumbel_shape.begin() + offset, logits.shape(axis)); diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 238db5971b..96a789b608 100644 --- a/python/tests/test_random.py +++ b/python/tests/test_random.py @@ -348,6 +348,21 @@ def test_categorical(self): with self.assertRaises(ValueError): mx.random.categorical(logits, shape=[10, 5], num_samples=5) + # Single distribution. + logits = mx.zeros((20,)) + + out = mx.random.categorical(logits, num_samples=7) + self.assertEqual(out.shape, (7,)) + self.assertEqual(out.dtype, mx.uint32) + self.assertTrue(mx.max(out).item() < 20) + + out = mx.random.categorical(logits, 0, [5, 3]) + self.assertEqual(out.shape, (5, 3)) + self.assertTrue(mx.max(out).item() < 20) + + self.assertEqual(mx.random.categorical(logits, num_samples=1).shape, (1,)) + self.assertEqual(mx.random.categorical(logits, num_samples=0).shape, (0,)) + def test_permutation(self): x = sorted(mx.random.permutation(4).tolist()) self.assertEqual([0, 1, 2, 3], x) diff --git a/tests/random_tests.cpp b/tests/random_tests.cpp index 04a86352f1..9dfe98a5d8 100644 --- a/tests/random_tests.cpp +++ b/tests/random_tests.cpp @@ -642,6 +642,19 @@ TEST_CASE("test categorical") { CHECK_EQ(categorical(logits, -1, 7).shape(), Shape{5, 4, 7}); CHECK_EQ(categorical(logits, -2, 7).shape(), Shape{5, 3, 7}); CHECK_EQ(categorical(logits, -3, 7).shape(), Shape{4, 3, 7}); + + // Infinities mean the same thing when several samples are drawn + logits = array({1.0f, -2.0f, inf, 4.0f, 3.0f}); + CHECK(all(equal(categorical(logits, 0, 5), array(2u))).item()); + + logits = array({-inf, -2.0f, -inf, -inf}); + CHECK(all(equal(categorical(logits, 0, 5), array(1u))).item()); + + // A -inf category carries no mass and is never drawn + logits = array({-inf, 0.0f, -inf, 0.0f}); + out = categorical(logits, 0, 100); + CHECK(all(logical_or(equal(out, array(1u)), equal(out, array(3u)))) + .item()); } TEST_CASE("test laplace") { From 307692cb0ec40f4302b9753f3518f3f633d603f7 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:53:20 -0400 Subject: [PATCH 176/222] Fail when a requested RDMA device is not found (#4180) --- mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index a0e7153c0d..d65b1a536c 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -276,6 +276,7 @@ std::vector create_connections( } // Search for the name and try to open the device + bool found = false; for (int i = 0; i < num_devices; i++) { if (name == ibv().get_device_name(devices[i])) { auto ctx = ibv().open_device(devices[i]); @@ -285,9 +286,19 @@ std::vector create_connections( throw std::runtime_error(msg.str()); } connections.emplace_back(ctx); + found = true; break; } } + + // Returning a shorter vector than device_names would leave the callers, + // which size themselves from device_names, indexing past its end. + if (!found) { + std::ostringstream msg; + msg << "[jaccl] Could not find device " << name << " (" << num_devices + << " available)"; + throw std::runtime_error(msg.str()); + } } ibv().free_device_list(devices); From ae8e9a732c6beb10bb120e2611d40a8b44636f89 Mon Sep 17 00:00:00 2001 From: robertomeroni <150194833+robertomeroni@users.noreply.github.com> Date: Wed, 12 Aug 2026 05:53:39 +0200 Subject: [PATCH 177/222] Bound GGUF tensor data offsets against the file mapping (#4179) --- mlx/io/gguf.cpp | 23 ++++++++++++ tests/load_tests.cpp | 89 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 758488e4a8..40cca573e5 100644 --- a/mlx/io/gguf.cpp +++ b/mlx/io/gguf.cpp @@ -211,6 +211,28 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { return metadata; } +// gguflib computes weights_data as ctx->data + ctx->data_off + the tensor's +// offset field in unsigned arithmetic, without comparing the result against the +// mapping, so a crafted offset can point outside the file or -- if the addition +// wraps -- back inside it at the wrong bytes. +void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) { + auto fail = [&tensor](const std::string& what) { + std::ostringstream msg; + msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) + << "' " << what << ". Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + }; + if (tensor.offset < ctx->data_off) { + fail("has a data offset that overflows the data section"); + } + if (tensor.offset > ctx->size) { + fail("has a data offset past the end of the file"); + } + if (tensor.bsize > ctx->size - tensor.offset) { + fail("extends past the end of the file"); + } +} + std::unordered_map load_arrays(gguf_ctx* ctx) { std::unordered_map array_map; gguf_tensor tensor; @@ -225,6 +247,7 @@ std::unordered_map load_arrays(gguf_ctx* ctx) { }; while (gguf_get_tensor(ctx, &tensor)) { + check_tensor_in_file(ctx, tensor); if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1 || tensor.type == GGUF_TYPE_Q8_0) { gguf_load_quantized(array_map, tensor); diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 33ad34bc15..8974919476 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -168,6 +168,95 @@ TEST_CASE("test gguf") { } } +// Writes a one-tensor GGUF (name "t", ndim 1, dim 4, type F32) whose tensor +// data offset field is set verbatim to `tensor_data_offset`. Writes +// `data_bytes` bytes of tensor data, defaulting to the full four floats. +void write_raw_gguf( + const std::string& path, + uint64_t tensor_data_offset, + size_t data_bytes = 4 * sizeof(float)) { + std::ofstream out(path, std::ios::binary); + auto u32 = [&out](uint32_t v) { + out.write(reinterpret_cast(&v), 4); + }; + auto u64 = [&out](uint64_t v) { + out.write(reinterpret_cast(&v), 8); + }; + out.write("GGUF", 4); + u32(3); // version + u64(1); // tensor_count + u64(0); // metadata_kv_count + u64(1); // tensor name length + out.write("t", 1); + u32(1); // ndim + u64(4); // dim[0] + u32(0); // GGUF_TYPE_F32 + u64(tensor_data_offset); + while (out.tellp() % 32 != 0) { // default GGUF alignment + out.put(0); + } + std::vector data(data_bytes, 0); + out.write(data.data(), data.size()); +} + +TEST_CASE("test gguf tensor data offset validation") { + // A crafted tensor data offset must be rejected rather than turned into a + // pointer outside the mapping. See ml-explore/mlx#4136. + SUBCASE("valid offset loads") { + std::string file_path = get_temp_file("test_gguf_offset_ok.gguf"); + write_raw_gguf(file_path, 0); + auto [weights, metadata] = load_gguf(file_path); + CHECK_EQ(weights.size(), 1); + CHECK(array_equal(weights.at("t"), zeros({4}, float32)).item()); + } + + SUBCASE("offset past the end of the file") { + std::string file_path = get_temp_file("test_gguf_offset_past_end.gguf"); + write_raw_gguf(file_path, 1ull << 20); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset far past the end of the file") { + std::string file_path = get_temp_file("test_gguf_offset_far_past.gguf"); + write_raw_gguf(file_path, 1ull << 40); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset that overflows the data section base") { + // Wraps back to an in-mapping address, so an end-pointer-only check would + // silently read the wrong bytes instead of reading out of bounds. + std::string file_path = get_temp_file("test_gguf_offset_wrap.gguf"); + write_raw_gguf(file_path, ~0ull - 8); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor extends past the end of the file") { + // In-range offset, but the data is truncated: only the extent check + // catches this. + std::string file_path = get_temp_file("test_gguf_truncated.gguf"); + write_raw_gguf(file_path, 0, 4 * sizeof(float) - 1); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor starts inside the file but ends past it") { + // A small offset, so the start is in range and only the extent decides. + // Pins the extent check to the tensor's own start rather than to the + // start of the data section. + std::string file_path = get_temp_file("test_gguf_partial_overrun.gguf"); + write_raw_gguf(file_path, 8); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset just past the end of the file") { + // Only a few bytes past the end rather than far outside it, so the + // resulting pointer is still in the mapped page and reads succeed + // silently. Pins the offset bound to the file size exactly. + std::string file_path = get_temp_file("test_gguf_offset_just_past.gguf"); + write_raw_gguf(file_path, 20); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } +} + TEST_CASE("test gguf metadata") { std::string file_path = get_temp_file("test_arr.gguf"); using dict = std::unordered_map; From d935276decaa152cafc91e6f732093650a08999d Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 11 Aug 2026 20:53:54 -0700 Subject: [PATCH 178/222] Shift by the max in log_softmax (#4169) --- python/mlx/nn/layers/activations.py | 4 ++++ python/tests/test_nn.py | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/python/mlx/nn/layers/activations.py b/python/mlx/nn/layers/activations.py index 58fef50687..305b5ec7cc 100644 --- a/python/mlx/nn/layers/activations.py +++ b/python/mlx/nn/layers/activations.py @@ -68,6 +68,10 @@ def log_softmax(x, axis=-1): Applies :math:`x - \log \sum_i e^{x_i}` element wise. """ + # Shift by the max first. Subtracting the logsumexp of x directly loses the + # normalizer, since adding it to a large max rounds away before the + # subtraction happens. + x = x - mx.stop_gradient(mx.max(x, axis=axis, keepdims=True)) return x - mx.logsumexp(x, axis=axis, keepdims=True) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 67828fd86e..c5e6db94a7 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -1216,6 +1216,11 @@ def test_log_softmax(self): self.assertEqual(y.shape, (3,)) self.assertEqual(y.dtype, mx.float32) + # Equal logits normalize to log(1/n) whatever their magnitude + for v in [1e0, 1e4, 1e8, 1e20, 1e36]: + y = nn.log_softmax(mx.array([[v, v]])) + self.assertTrue(mx.allclose(y, mx.full((1, 2), -0.6931472), atol=1e-5)) + def test_log_sigmoid(self): x = mx.array([1.0, -1.0, 0.0]) y = nn.log_sigmoid(x) From 74c9acd12b21d9cccf4dff2880b3a853d48217a6 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 11 Aug 2026 20:54:04 -0700 Subject: [PATCH 179/222] Shift by the max in cross_entropy (#4188) --- python/mlx/nn/losses.py | 5 +++++ python/tests/test_losses.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 9226c10eee..184df2a2e0 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -83,6 +83,11 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) + # Shift by the max first. The loss only depends on differences between + # logits, but subtracting the logsumexp of large logits loses the gap to + # rounding before the subtraction happens. + logits = logits - mx.stop_gradient(mx.max(logits, axis=axis, keepdims=True)) + if targets_as_probs: score = mx.sum(logits * targets, axis=axis) else: diff --git a/python/tests/test_losses.py b/python/tests/test_losses.py index e282fb4053..96174f3c89 100644 --- a/python/tests/test_losses.py +++ b/python/tests/test_losses.py @@ -17,6 +17,22 @@ def test_cross_entropy(self): loss = nn.losses.cross_entropy(logits, indices, reduction="none") self.assertTrue(mx.allclose(loss, expected)) + # The loss only depends on the gaps between logits, so a large shared + # offset must not change it + indices = mx.array([0]) + base = mx.array([[2.0, -1.0]]) + expected = nn.losses.cross_entropy(base, indices, reduction="none") + for offset in [1e4, 1e6]: + loss = nn.losses.cross_entropy(base + offset, indices, reduction="none") + self.assertTrue(mx.allclose(loss, expected, atol=1e-5)) + + # Equal logits give log(n) whatever their magnitude + for v in [1e0, 1e4, 1e8, 1e20]: + loss = nn.losses.cross_entropy( + mx.array([[v, v]]), indices, reduction="none" + ) + self.assertTrue(mx.allclose(loss, mx.array([0.6931472]), atol=1e-5)) + probs = mx.array([[1.0, 0.0], [0.0, 1.0]]) loss = nn.losses.cross_entropy(logits, probs, reduction="none") self.assertTrue(mx.isnan(loss).all()) # produce NaNs, like PyTorch From 152baf48c89c60ada9a83813e771f47b01520680 Mon Sep 17 00:00:00 2001 From: Ishaan Samantray Date: Tue, 11 Aug 2026 23:54:41 -0400 Subject: [PATCH 180/222] Fix linalg.norm(x, keepdims=True) not keeping dims (#4166) --- mlx/linalg.cpp | 9 ++++++++- python/tests/test_linalg.py | 3 +++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index e86c030c36..0d06b56469 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -166,7 +166,14 @@ array norm( bool keepdims /* = false */, StreamOrDevice s /* = {} */) { if (!axis) { - return norm(flatten(a, s), std::vector{0}, keepdims, s); + auto out = norm(flatten(a, s), std::vector{0}, keepdims, s); + if (keepdims) { + // The flatten above collapses the input to one dimension, so keepdims + // has to restore the rank of the original array rather than the + // flattened one. + out = reshape(out, Shape(a.ndim(), 1), s); + } + return out; } if (axis.value().size() > 2) { diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index 9b3859976f..52e02e4d9a 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -38,6 +38,7 @@ def test_norm(self): with self.subTest( shape=shape, ord=o, axis=axis, keepdims=keepdims ): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -51,6 +52,7 @@ def test_norm(self): out_np = np.linalg.norm(x_np, ord=o, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, ord=o, keepdims=keepdims) with self.subTest(shape=shape, ord=o, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -63,6 +65,7 @@ def test_norm(self): out_np = np.linalg.norm(x_np, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, keepdims=keepdims) with self.subTest(shape=shape, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)) # tests for negative indexing: -1/1/inf/-inf/ From a4356745781db7730c135f92d19f94dd0b953340 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:55:13 -0700 Subject: [PATCH 181/222] chore: Use dispatch_inexact_types in binary_float_op_cpu (#4194) --- mlx/backend/cpu/binary.h | 23 +++-------------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/mlx/backend/cpu/binary.h b/mlx/backend/cpu/binary.h index acbb71aae3..5333aff9a1 100644 --- a/mlx/backend/cpu/binary.h +++ b/mlx/backend/cpu/binary.h @@ -360,26 +360,9 @@ void binary_float_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (out.dtype()) { - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - default: - throw std::runtime_error( - "[binary_float] Only supports floating point types."); - } + dispatch_inexact_types(out.dtype(), "[binary_float]", [&](auto type_tag) { + binary_op(a, b, out, bopt); + }); }); } From 36fde27577e4649a22a8f1f2a34e548d19dff943 Mon Sep 17 00:00:00 2001 From: Vedant Singhal <88609765+Ved235@users.noreply.github.com> Date: Wed, 12 Aug 2026 09:28:48 +0530 Subject: [PATCH 182/222] Implement batched matmul for large 1D dot products (#3580) --- mlx/backend/metal/kernels/CMakeLists.txt | 1 + mlx/backend/metal/kernels/dot.h | 68 ++++++++++++++++++++++++ mlx/backend/metal/kernels/dot.metal | 18 +++++++ mlx/backend/metal/matmul.cpp | 60 +++++++++++++++++++++ python/tests/test_blas.py | 31 +++++++++++ 5 files changed, 178 insertions(+) create mode 100644 mlx/backend/metal/kernels/dot.h create mode 100644 mlx/backend/metal/kernels/dot.metal diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index e72cb41ab4..ecabbda550 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -49,6 +49,7 @@ endfunction(build_kernel) build_kernel(arg_reduce) build_kernel(conv steel/conv/params.h) +build_kernel(dot) build_kernel(layer_norm) build_kernel(random) build_kernel(rms_norm) diff --git a/mlx/backend/metal/kernels/dot.h b/mlx/backend/metal/kernels/dot.h new file mode 100644 index 0000000000..6917b99f0e --- /dev/null +++ b/mlx/backend/metal/kernels/dot.h @@ -0,0 +1,68 @@ +#pragma once + +#include + +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +template < + typename T, + const int ITEMS_PER_THREAD, + const int TG_SIZE, + const uint SIMD_GROUPS> +[[kernel]] void dot_product( + const device T* a [[buffer(0)]], + const device T* b [[buffer(1)]], + device float* output [[buffer(2)]], + const constant int& n [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_id [[simdgroup_index_in_threadgroup]], + uint tg_id [[threadgroup_position_in_grid]]) { + constexpr int VEC = 16 / sizeof(T); + int start = (tg_id * TG_SIZE + simd_id * 32) * ITEMS_PER_THREAD + lane * VEC; + + float4 c = 0.0f; + + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < ITEMS_PER_THREAD; i += VEC) { + int idx = start + i * ITEMS_PER_THREAD; + if (idx + VEC <= n) { + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < VEC; j += 4) { + c += float4(*reinterpret_cast*>( + a + idx + j)) * + float4(*reinterpret_cast*>( + b + idx + j)); + } + } else { + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < VEC; ++j) { + int nidx = idx + j; + if (nidx < n) { + c[j & 3] += float(a[nidx]) * float(b[nidx]); + } + } + } + } + + threadgroup float smem[SIMD_GROUPS]; + + float sum = c[0] + c[1] + c[2] + c[3]; + sum = simd_sum(sum); + + if (lane == 0) { + smem[simd_id] = sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tid < SIMD_GROUPS) { + sum = smem[tid]; + sum = simd_sum(sum); + if (tid == 0) { + output[tg_id] = sum; + } + } +} diff --git a/mlx/backend/metal/kernels/dot.metal b/mlx/backend/metal/kernels/dot.metal new file mode 100644 index 0000000000..5c9ea3be64 --- /dev/null +++ b/mlx/backend/metal/kernels/dot.metal @@ -0,0 +1,18 @@ +#include + +#include "mlx/backend/metal/kernels/dot.h" + +#define instantiate_dot_product_kernel( \ + name, itype, items_per_thread, tg_size, simd_groups) \ + instantiate_kernel( \ + "dot_product_" #name "_it" #items_per_thread "_tg" #tg_size \ + "_sg" #simd_groups, \ + dot_product, \ + itype, \ + items_per_thread, \ + tg_size, \ + simd_groups) + +instantiate_dot_product_kernel(float32, float, 32, 512, 16); +instantiate_dot_product_kernel(float16, half, 32, 512, 16); +instantiate_dot_product_kernel(bfloat16, bfloat16_t, 32, 512, 16); diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index b61a58adb3..8c0b46d6de 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -15,6 +15,7 @@ #include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/kernels/steel/gemm/params.h" #include "mlx/backend/metal/matmul.h" +#include "mlx/backend/metal/reduce.h" #include "mlx/backend/metal/utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -1040,6 +1041,54 @@ void steel_matmul_axpby( // GEMV dispatch /////////////////////////////////////////////////////////////////////////////// +void dot_product( + const Stream& s, + metal::Device& d, + const array& a, + const array& b, + array& out, + int K, + std::vector& copies) { + constexpr int thread_group_size = 512; + constexpr int items_per_thread = 32; + constexpr int simd_groups = thread_group_size / 32; + auto& compute_encoder = metal::get_command_encoder(s); + std::string kname = "dot_product_" + type_to_name(a); + concatenate( + kname, + "_it", + items_per_thread, + "_tg", + thread_group_size, + "_sg", + simd_groups); + auto kernel = d.get_kernel(kname); + + int n = K; + int threads = (n + items_per_thread - 1) / items_per_thread; + int blocks = (threads + thread_group_size - 1) / thread_group_size; + + array partials({blocks}, float32, nullptr, {}); + partials.set_data(allocator::malloc(partials.nbytes())); + copies.push_back(partials); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(b, 1); + compute_encoder.set_output_array(partials, 2); + compute_encoder.set_bytes(n, 3); + compute_encoder.dispatch_threads( + MTL::Size(size_t(blocks) * thread_group_size, 1, 1), + MTL::Size(thread_group_size, 1, 1)); + + array tempResult(out.shape(), float32, nullptr, {}); + tempResult.set_data(allocator::malloc(tempResult.nbytes())); + copies.push_back(tempResult); + all_reduce_dispatch(partials, tempResult, "sum", compute_encoder, d, s); + copy_gpu(tempResult, out, CopyType::Scalar, s); + + compute_encoder.add_temporaries(std::move(copies)); +} + template void gemv_axbpy( const Stream& s, @@ -1455,6 +1504,17 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { ///////////////////////////////////////////////////////////////////////////// // Gemv specialization + if (M == 1 && N == 1 && batch_size_out == 1 && a.flags().row_contiguous && + b.flags().row_contiguous && a.dtype() != complex64) { + return dot_product( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ a, + /* const array& b = */ b, + /* array& out = */ out, + /* int K = */ K, + /* std::vector& copies = */ copies); + } // The wide gemv route streams the weight matrix once per <= 5 input // vectors instead of running a row-padded GEMM tile. if (!a_transposed && b_transposed && diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 957c46aec9..7dbbaea049 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -538,6 +538,37 @@ def test_matrix_vector_edgecases(self): ) self.assertTrue(np.array_equal(c_mlx, c_npy)) + def test_dot_product(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + def run_test(dtype, size, offset, atol): + with self.subTest(dtype=str(dtype), size=size, offset=offset): + np.random.seed(42) + scale = size**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, size + offset).astype(np.float32) + ).astype(dtype)[offset:] + b_mx = mx.array( + np.random.normal(0.0, scale, size + offset).astype(np.float32) + ).astype(dtype)[offset:] + + expected = np.inner( + np.array(a_mx.astype(mx.float32)), + np.array(b_mx.astype(mx.float32)), + ) + actual = np.array(mx.inner(a_mx, b_mx).astype(mx.float32)) + self.assertTrue(np.allclose(actual, expected, atol=atol)) + + for dtype, atol in ( + (mx.float32, 1e-5), + (mx.float16, 2e-3), + (mx.bfloat16, 2e-3), + ): + for size in (1023, 1024, 1025, 16385, 131072, 1000000): + for offset in (0, 1): + run_test(dtype, size, offset, atol) + def test_wide_matmul(self): if mx.default_device() == mx.cpu: self.skipTest("requires GPU") From 01d4e1238d8d47d821c1199ed5aedad60ff8ba30 Mon Sep 17 00:00:00 2001 From: Philip John Basile Date: Wed, 12 Aug 2026 04:11:48 -0400 Subject: [PATCH 183/222] Re-disable qmm_n_nax and fix group_size < 64 (#4202) Co-authored-by: Cheng --- mlx/backend/metal/kernels/fp_quantized_nax.h | 2 +- mlx/backend/metal/kernels/quantized_nax.h | 4 +- mlx/backend/metal/quantized.cpp | 7 +- python/tests/test_quantized.py | 102 +++++++++++++------ 4 files changed, 78 insertions(+), 37 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 7c452e6a64..cf64ff7f46 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -178,7 +178,7 @@ struct QuantizedBlockLoader { if (reduction_dim == 1) { scales += n_groups; } else { - scales += n_groups * group_stride; + scales += group_stride; } } }; diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 81f6a7a252..31e51a5b7e 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -826,8 +826,8 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; + scales += group_stride; + biases += group_stride; } } }; diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index ab3e1deacc..f659e16c93 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1035,9 +1035,10 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { - // The non-transposed kernel requires N % 64 == 0. - if (metal::is_nax_available() && (transpose || (N % 64 == 0)) && - (K % 64 == 0) && (env::enable_tf32() || x.dtype() != float32)) { + bool has_nax_kernel = + metal::is_nax_available() && (transpose || mode == "affine"); + if (has_nax_kernel && transpose && (K % 64 == 0) && + (env::enable_tf32() || x.dtype() != float32)) { return qmm_nax( /* const array& x = */ x, /* const array& w = */ w, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 56ba75db32..15bc892bd8 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,5 +1,6 @@ # Copyright © 2023-2026 Apple Inc. +import os import platform import subprocess import unittest @@ -353,6 +354,7 @@ def test_qmm_large_dims(self): tol = 1e-3 if dtype == mx.float32 else 1.5e-3 self.assertLess((y_q - y_hat).abs().max(), tol) + @unittest.skipIf("CI" in os.environ, "too slow in CI") def test_qmm_non_transposed(self): # The non-transposed matmul (w is [K, N]) is reachable mainly from the # vjp of a quantized linear layer, so it gets much less coverage than @@ -360,45 +362,83 @@ def test_qmm_non_transposed(self): # values that leave a partial M-tile. key = mx.random.key(0) k1, k2 = mx.random.split(key) - dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 - tol = 1e-3 if dtype == mx.float32 else 1.5e-3 - def check(M, K, N, group_size, bits, batch=()): - x = mx.random.normal(shape=(*batch, M, K), key=k1) / K**0.5 - w = mx.random.normal(shape=(K, N), key=k2) / K**0.5 - x = x.astype(dtype) - w = w.astype(dtype) + modes = ["mxfp4", "nvfp4", "mxfp8"] + if mx.default_device() == mx.gpu: + dtypes = [mx.float16, mx.bfloat16] + else: + dtypes = [mx.float32] + + def check_affine(M, K, N, group_size, bits, dtype, batch=()): + x = mx.random.normal(shape=(*batch, M, K), key=k1, dtype=dtype) / K**0.5 + w = mx.random.normal(shape=(K, N), key=k2, dtype=dtype) / K**0.5 w_q, scales, biases = mx.quantize(w, group_size, bits) w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) y_q = mx.quantized_matmul(x, w_q, scales, biases, False, group_size, bits) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 + self.assertLess((y_q - y_hat).abs().max(), tol) + + def check_fp(M, K, N, mode, dtype, batch=()): + x = mx.random.normal(shape=(*batch, M, K), key=k1, dtype=dtype) / K**0.5 + w = mx.random.normal(shape=(K, N), key=k2, dtype=dtype) / K**0.5 + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + y_q = mx.quantized_matmul(x, w_q, scales, None, False, mode=mode) + y_hat = x @ w_hat + self.assertEqual(y_q.shape, y_hat.shape) + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 self.assertLess((y_q - y_hat).abs().max(), tol) - # M sweep. 33..63 is the interesting range: a whole simdgroup of the - # threadgroup's M-tile falls past the end of the matrix. - for M in [1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129]: - for group_size, bits in [(64, 4), (128, 4), (64, 8)]: - with self.subTest(M=M, group_size=group_size, bits=bits): - check(M, 512, 1024, group_size, bits) - - # Transformer-sized K/N, aligned and unaligned M. - for K, N in [(2048, 2048), (512, 2048), (2048, 512), (11008, 2048)]: - for M in [100, 256]: - with self.subTest(shape=(M, K, N)): - check(M, K, N, 64, 4) - - # Batched x, unaligned M. - for batch in [(2,), (2, 3)]: - for M in [33, 250]: - with self.subTest(batch=batch, M=M): - check(M, 512, 1024, 64, 4, batch=batch) - - # M > 2**15 with a partial M-tile, so the per-simdgroup row count is a - # distance that does not fit in an int16. Same failure mode as the one - # test_qmm_large_dims covers for the transposed kernel. - with self.subTest(shape=(33000, 128, 64)): - check(33000, 128, 64, 64, 4) + for dtype in dtypes: + # M sweep. 33..63 is the interesting range: a whole simdgroup of the + # threadgroup's M-tile falls past the end of the matrix. + for M in [1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129]: + for group_size, bits in [(64, 4), (128, 4), (64, 8)]: + with self.subTest( + M=M, group_size=group_size, bits=bits, dtype=dtype + ): + check_affine(M, 512, 1024, group_size, bits, dtype) + for mode in modes: + with self.subTest(M=M, mode=mode, dtype=dtype): + check_fp(M, 512, 1024, mode, dtype) + + # Transformer-sized K/N, aligned and unaligned M. + for K, N in [(2048, 2048), (512, 2048), (2048, 512), (11008, 2048)]: + for M in [100, 256]: + with self.subTest(shape=(M, K, N), dtype=dtype): + check_affine(M, K, N, 64, 4, dtype) + for mode in modes: + with self.subTest(shape=(M, K, N), mode=mode, dtype=dtype): + check_fp(M, 512, 1024, mode, dtype) + + # Batched x, unaligned M. + for batch in [(2,), (2, 3)]: + for M in [33, 250]: + with self.subTest(batch=batch, M=M, dtype=dtype): + check_affine(M, 512, 1024, 64, 4, dtype, batch=batch) + for mode in modes: + with self.subTest(batch=batch, mode=mode, dtype=dtype): + check_fp(M, 512, 1024, mode, dtype, batch=batch) + + # M > 2**15 with a partial M-tile, so the per-simdgroup row count is a + # distance that does not fit in an int16. Same failure mode as the one + # test_qmm_large_dims covers for the transposed kernel. + with self.subTest(shape=(33000, 128, 64), dtype=dtype): + check_affine(33000, 128, 64, 64, 4, dtype) + check_fp(33000, 128, 64, mode, dtype) + + # K=64 is the single reduction-tile control; K > 64 spans two or more + # tiles, which exposed the over-advanced scale pointer. + for M in [8, 33, 65]: + for K in [64, 128, 256]: + for bits in [2, 4, 8]: + with self.subTest(M=M, K=K, bits=bits, dtype=dtype): + check_affine(M, K, 128, 32, bits, dtype) + for mode in modes: + with self.subTest(M=M, K=K, mode=mode, dtype=dtype): + check_fp(M, K, 128, mode, dtype) def test_qmm_vjp(self): key = mx.random.key(0) From d34f630d983e309f8815a5961b2f4dab26adae11 Mon Sep 17 00:00:00 2001 From: Bunlong Heng Date: Wed, 12 Aug 2026 04:12:55 -0400 Subject: [PATCH 184/222] chore: shell-quote the working directory in the distributed launch script (#4189) --- python/mlx/_distributed_utils/launch.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 464bff7e70..4771c1fb5b 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -121,11 +121,12 @@ def make_launch_script(rank, python, cwd, files, env, command, is_local): # Change the working directory if one was requested. Otherwise attempt to # change to the current one but don't fail if it wasn't possible. d = cwd or os.getcwd() - script += f"if [[ -d {repr(d)} ]]; then " - script += f" cd {repr(d)}; " + qd = shlex.quote(d) + script += f"if [[ -d {qd} ]]; then " + script += f" cd {qd}; " if cwd is not None: script += "else " - script += f" echo 'Failed to change directory to' {repr(d)} >2; " + script += f" echo 'Failed to change directory to' {qd} >&2; " script += "fi; " # Add the environment variables that were requested From 21d897d3b842d63352f0f60424063cb93909b3f4 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Wed, 12 Aug 2026 04:13:46 -0400 Subject: [PATCH 185/222] Report when no usable GID is found (#4191) --- mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index d65b1a536c..be51364787 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -181,18 +181,30 @@ const Destination& Connection::info() { ibv_port_attr port_attr; ibv().query_port(ctx, 1, &port_attr); - ibv_gid gid; + ibv_gid gid = {}; + bool found_gid = false; for (int i = 0; i < port_attr.gid_tbl_len; i++) { ibv_gid tmp; if (ibv().query_gid(ctx, 1, i, &tmp) == 0) { if (*(uint64_t*)&tmp.raw[0] == 0 && *(uint16_t*)&tmp.raw[8] == 0 && *(uint16_t*)&tmp.raw[10] == 0xffff) { gid = tmp; + found_gid = true; break; } } } + // Fail here rather than hand an unset GID to the queue pair. + if (!found_gid) { + std::ostringstream msg; + msg << "[jaccl] No IPv4-mapped GID for this device. Thunderbolt RDMA ports " + << "only publish one once the interface has an IPv4 address; assign a " + << "link-local address to it, for example: ifconfig inet " + << "169.254.0.1 netmask 255.255.0.0 alias"; + throw std::runtime_error(msg.str()); + } + src.local_id = port_attr.lid; src.queue_pair_number = queue_pair->qp_num; src.packet_sequence_number = 7; From 3abd0fd6b3eb9d9d3a34cb65e8a2189c57260399 Mon Sep 17 00:00:00 2001 From: Nilesh Patil <128893479+nileshpatil6@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:22:31 +0530 Subject: [PATCH 186/222] Keep double precision for python floats in float64 operations (#4173) --- python/src/ops.cpp | 15 ++++++++++----- python/src/utils.cpp | 12 +++++++++--- python/tests/test_double.py | 38 +++++++++++++++++++++++++++++++++++++ 3 files changed, 57 insertions(+), 8 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index a666588f51..d8892a45d2 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3491,18 +3491,23 @@ void init_ops(nb::module_& m) { const ScalarOrArray& constant_value, mx::StreamOrDevice s) { if (auto pv = std::get_if(&pad_width); pv) { - return mx::pad(a, *pv, to_array(constant_value), mode, s); + return mx::pad(a, *pv, to_array(constant_value, a.dtype()), mode, s); } else if (auto pv = std::get_if>(&pad_width); pv) { return mx::pad( - a, std::get<0>(*pv), to_array(constant_value), mode, s); + a, + std::get<0>(*pv), + to_array(constant_value, a.dtype()), + mode, + s); } else if (auto pv = std::get_if>(&pad_width); pv) { - return mx::pad(a, *pv, to_array(constant_value), mode, s); + return mx::pad(a, *pv, to_array(constant_value, a.dtype()), mode, s); } else { auto v = std::get>>(pad_width); if (v.size() == 1) { - return mx::pad(a, v[0], to_array(constant_value), mode, s); + return mx::pad( + a, v[0], to_array(constant_value, a.dtype()), mode, s); } else { - return mx::pad(a, v, to_array(constant_value), mode, s); + return mx::pad(a, v, to_array(constant_value, a.dtype()), mode, s); } } }, diff --git a/python/src/utils.cpp b/python/src/utils.cpp index 4416531dc8..84214cd47b 100644 --- a/python/src/utils.cpp +++ b/python/src/utils.cpp @@ -31,9 +31,15 @@ mx::array to_array( return mx::array(val, (out_t == mx::bool_) ? mx::int32 : out_t); } else if (auto pv = std::get_if(&v); pv) { auto out_t = dtype.value_or(mx::float32); - return mx::array( - nb::cast(*pv), - mx::issubdtype(out_t, mx::floating) ? out_t : mx::float32); + if (!mx::issubdtype(out_t, mx::floating)) { + out_t = mx::float32; + } + if (out_t == mx::float64) { + // Cast straight to double: going through float would round the value + // to float32 precision while the result still advertises float64. + return mx::array(nb::cast(*pv), out_t); + } + return mx::array(nb::cast(*pv), out_t); } else if (auto pv = std::get_if>(&v); pv) { return mx::array(static_cast(*pv), mx::complex64); } else if (auto pv = std::get_if(&v); pv) { diff --git a/python/tests/test_double.py b/python/tests/test_double.py index f957e8fc96..65603cd937 100644 --- a/python/tests/test_double.py +++ b/python/tests/test_double.py @@ -296,6 +296,44 @@ def test_conversion(self): b = a.tolist() self.assertEqual(b, [1.0, 2.0]) + def test_python_float_keeps_double_precision(self): + with mx.stream(mx.cpu): + # https://github.com/ml-explore/mlx/issues/4160 + for v in (0.37, 0.1, math.pi): + self.assertEqual(mx.full((3,), v, dtype=mx.float64).tolist(), [v] * 3) + + # https://github.com/ml-explore/mlx/issues/4159 + a = mx.array([1.0], dtype=mx.float64) + for v in (0.1, 1e-4, math.pi): + out = a * v + self.assertEqual(out.dtype, mx.float64) + self.assertEqual(out.tolist(), [v]) + + # values outside the float32 range used to saturate to inf or zero + self.assertEqual((a * 1e300).tolist(), [1e300]) + self.assertEqual((a * 1e-300).tolist(), [1e-300]) + + # every op that pairs a scalar with an array goes through the same + # conversion + zero = mx.array([0.0], dtype=mx.float64) + self.assertEqual(mx.maximum(zero, 0.1).tolist(), [0.1]) + self.assertEqual(mx.minimum(a, 0.1).tolist(), [0.1]) + self.assertEqual(mx.clip(a, None, 0.1).tolist(), [0.1]) + self.assertEqual(mx.where(mx.array([False]), zero, 0.1).tolist(), [0.1]) + self.assertEqual( + mx.pad(a, 1, constant_values=0.1).tolist(), [0.1, 1.0, 0.1] + ) + + # the python float is still weak: it does not widen the array + for dtype in (mx.float16, mx.bfloat16, mx.float32): + self.assertEqual((mx.array([1.0], dtype=dtype) * 0.1).dtype, dtype) + self.assertEqual((mx.array([1], dtype=mx.int32) * 0.1).dtype, mx.float32) + + # pad() keeps returning the input's dtype for every input type + for dtype in (mx.int32, mx.float16, mx.bfloat16, mx.float32): + padded = mx.pad(mx.array([1.0], dtype=dtype), 1, constant_values=0.5) + self.assertEqual(padded.dtype, dtype) + def test_linspace(self): with mx.stream(mx.cpu): vals = mx.linspace(0, math.pi, 2, mx.float64) From e3d389eed07d052264796c9d80b3aeef5eb10073 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Wed, 12 Aug 2026 10:54:00 -0700 Subject: [PATCH 187/222] Jaccl ring multi wire fix (#4193) --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 57 +++++++++++---------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index ad05f7b208..15b81a2177 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -121,19 +121,19 @@ class RingImpl { int n_wires, int lw, ReduceOp reduce_op) { - // The element offset (within a chunk) of this wire's slice in each - // direction and the end of each direction's region. Wire slices are - // contiguous rather than interleaved. Direction lr owns the chunk region - // [lr * n_wires * size_per_wire, (lr + 1) * n_wires * size_per_wire) (the - // last region is clamped to chunk_size). + // This wire's slice within the chunk: [wire_offset, wire_end). Clamp to the + // per wire end (not the whole region) so the last frame can't spill into + // the next wire's slice when size_per_wire is not a multiple of N. int64_t wire_offset[MAX_DIR]; - int64_t region_end[MAX_DIR]; + int64_t wire_end[MAX_DIR]; int64_t send_offset[MAX_DIR]; int64_t recv_offset[MAX_DIR]; for (int lr = 0; lr < MAX_DIR; lr++) { wire_offset[lr] = lr * n_wires * size_per_wire + static_cast(lw) * size_per_wire; - region_end[lr] = std::min(chunk_size, (lr + 1) * n_wires * size_per_wire); + int64_t region_end = + std::min(chunk_size, (lr + 1) * n_wires * size_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + size_per_wire); send_offset[lr] = rank_ * chunk_size; } recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; @@ -152,7 +152,7 @@ class RingImpl { size_ * chunk_size, size_per_wire, wire_offset, - region_end, + wire_end, send_offset, recv_offset, in_ptr, @@ -166,7 +166,7 @@ class RingImpl { size_ * chunk_size, size_per_wire, wire_offset, - region_end, + wire_end, send_offset, recv_offset, in_ptr, @@ -180,7 +180,7 @@ class RingImpl { size_ * chunk_size, size_per_wire, wire_offset, - region_end, + wire_end, send_offset, recv_offset, out_ptr, @@ -219,7 +219,10 @@ class RingImpl { // Both directions send the same contiguous slice of each rank's region. int64_t slice = static_cast(lw) * n_bytes_per_wire; int64_t wire_offset[2] = {slice, slice}; - int64_t region_end[2] = {n_bytes, n_bytes}; + // Clamp to this wire's slice end so the last frame can't spill into the + // next wire's slice. + int64_t wire_end_bytes = std::min(n_bytes, slice + n_bytes_per_wire); + int64_t wire_end[2] = {wire_end_bytes, wire_end_bytes}; int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; int64_t recv_offset[2] = { ((rank_ + size_ - 1) % size_) * n_bytes, @@ -232,7 +235,7 @@ class RingImpl { n_bytes * size_, n_bytes_per_wire, wire_offset, - region_end, + wire_end, send_offset, recv_offset, out_ptr, @@ -298,14 +301,16 @@ class RingImpl { int64_t N = buffer_bytes / sizeof(T); int64_t total = static_cast(size_) * chunk; - // This wire's element offset within the output chunk in each direction and - // the end of each direction's region. + // This wire's slice within the chunk: [wire_offset, wire_end). Clamp to the + // per wire end (not the whole region) so the last frame can't spill into + // the next wire's slice when size_per_wire is not a multiple of N. int64_t wire_offset[MAX_DIR]; - int64_t region_end[MAX_DIR]; + int64_t wire_end[MAX_DIR]; for (int lr = 0; lr < MAX_DIR; lr++) { wire_offset[lr] = lr * n_wires * size_per_wire + static_cast(lw) * size_per_wire; - region_end[lr] = std::min(chunk, (lr + 1) * n_wires * size_per_wire); + int64_t region_end = std::min(chunk, (lr + 1) * n_wires * size_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + size_per_wire); } // Input windows (count space, chunk aligned). Both directions converge on @@ -337,7 +342,7 @@ class RingImpl { int slice = recv_count[lr]; int b = slice % PIPELINE; int64_t offset = wire_offset[lr] + static_cast(slice) * N; - int64_t n = std::min(N, region_end[lr] - offset); + int64_t n = std::min(N, wire_end[lr] - offset); reduce_op( recv_buffer(sz, b, lr, lw).template begin(), in_ptr + in_recv_offset[lr] + offset, @@ -374,7 +379,7 @@ class RingImpl { std::copy( send_base + send_base_offset[lr] + offset, send_base + send_base_offset[lr] + - std::max(offset, std::min(offset + N, region_end[lr])), + std::max(offset, std::min(offset + N, wire_end[lr])), send_buffer(sz, buff, lr, lw).template begin()); send_count[lr]++; send_to(sz, buff, lr, lw); @@ -401,7 +406,7 @@ class RingImpl { std::copy( send_base + send_base_offset[lr] + offset, send_base + send_base_offset[lr] + - std::max(offset, std::min(offset + N, region_end[lr])), + std::max(offset, std::min(offset + N, wire_end[lr])), send_buffer(sz, buff, lr, lw).template begin()); send_count[lr]++; send_to(sz, buff, lr, lw); @@ -444,10 +449,10 @@ class RingImpl { // The first step sends from in_ptr, later steps from out_ptr; this lets the // all reduce seed its output without an up front copy. Callers that stage // into out_ptr (both all gather paths) pass out_ptr for in_ptr. wire_offset - // is this wire's element offset within a chunk, region_end the end of each - // direction's region. send_offset and recv_offset are the starting windows, - // updated in place so callers can chain passes. stride is the elements a - // window moves per step and total is the wrap around modulus. + // is this wire's element offset within a chunk, wire_end the end of its + // slice. send_offset and recv_offset are the starting windows, updated in + // place so callers can chain passes. stride is the elements a window moves + // per step and total is the wrap around modulus. template inline void ring_pass( int lw, @@ -456,7 +461,7 @@ class RingImpl { int64_t total, int64_t size_per_wire, const int64_t (&wire_offset)[MAX_DIR], - const int64_t (®ion_end)[MAX_DIR], + const int64_t (&wire_end)[MAX_DIR], int64_t (&send_offset)[MAX_DIR], int64_t (&recv_offset)[MAX_DIR], const T* in_ptr, @@ -477,9 +482,9 @@ class RingImpl { auto set_limits = [&]() { for (int lr = 0; lr < MAX_DIR; lr++) { send_limits[lr] = std::min( - region_end[lr], std::max(0, size - send_offset[lr])); + wire_end[lr], std::max(0, size - send_offset[lr])); recv_limits[lr] = std::min( - region_end[lr], std::max(0, size - recv_offset[lr])); + wire_end[lr], std::max(0, size - recv_offset[lr])); } }; set_limits(); From 52960f80f89b3720ec4cef2a18a3a390103abc5f Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:44:47 -0400 Subject: [PATCH 188/222] Report why creating a queue pair failed (#4209) --- ACKNOWLEDGMENTS.md | 1 + mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index d83832be8f..7005396e71 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -21,6 +21,7 @@ MLX was developed with contributions from the following individuals: - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops. - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. - katlun-lgtm: Added `reflect` and `symmetric` padding modes. +- Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting. diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index be51364787..f449d8c814 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -2,8 +2,10 @@ #include #include +#include #include #include +#include #include "jaccl/rdma.h" @@ -170,7 +172,12 @@ void Connection::create_queue_pair() { queue_pair = ibv().create_qp(protection_domain, &init_attr); if (queue_pair == nullptr) { - throw std::runtime_error("[jaccl] Couldn't create queue pair"); + int err = errno; + std::string error_message = std::generic_category().message(err); + std::ostringstream msg; + msg << "[jaccl] Creating the queue pair failed with '" << error_message + << " (" << errno << ")'."; + throw std::runtime_error(msg.str()); } } From 210a1e2c413839059aad1016ab7452efe6f36403 Mon Sep 17 00:00:00 2001 From: Noah Lyons Date: Wed, 12 Aug 2026 20:21:29 -0400 Subject: [PATCH 189/222] Add antialias support to nn.Upsample for linear and cubic modes (#3677) Co-authored-by: Cheng --- python/mlx/nn/layers/upsample.py | 222 +++++++++++++++++++++++++++++-- python/tests/test_upsample.py | 217 ++++++++++++++++++++++++++++++ 2 files changed, 430 insertions(+), 9 deletions(-) diff --git a/python/mlx/nn/layers/upsample.py b/python/mlx/nn/layers/upsample.py index f4499e50d6..0be9f87bfb 100644 --- a/python/mlx/nn/layers/upsample.py +++ b/python/mlx/nn/layers/upsample.py @@ -3,6 +3,7 @@ import operator from functools import partial, reduce from itertools import product +from math import ceil from typing import Callable, Literal, Tuple, Union import mlx.core as mx @@ -51,6 +52,118 @@ def _linear_indices(N, scale, align_corners, dim, ndims): ) +def _aa_indices(N, scale, align_corners, dim, ndims, kernel_fn, kernel_radius): + """Compute antialiased interpolation indices for a given kernel. + + When downscaling (scale < 1), the kernel support widens by 1/scale to + act as a low-pass filter, preventing aliasing. Out-of-bounds taps are + zeroed and weights are renormalized per output pixel. + + For upscaling (scale >= 1), the kernel is applied at its native width + without widening. This matches PyTorch's F.interpolate(antialias=True) + behavior where the kernel coefficient (e.g. a=-0.5 for cubic) is used + for both up and downsampling when antialias=True. + + Args: + N: input size for this dimension + scale: scale factor for this dimension + align_corners: align_corners flag + dim: which spatial dimension + ndims: number of spatial dimensions + kernel_fn: callable(distance) -> weight, operating in normalized + filter coordinates where |distance| < kernel_radius has support + kernel_radius: support radius of the kernel in normalized coords + (1.0 for triangle/linear, 2.0 for cubic) + """ + indices = _scaled_indices(N, scale, align_corners, dim, ndims) + + # For downscale, widen the filter by 1/scale + if scale < 1: + inv_scale = 1.0 / scale + else: + inv_scale = 1.0 + + support = kernel_radius * inv_scale + num_taps = ceil(support) + 1 + + # Compute per-tap weights, zero out-of-bounds, then normalize + all_idx = [] + all_w = [] + for k in range(-num_taps + 1, num_taps): + idx = mx.floor(indices) + k + # Map distance to normalized filter coordinates + dist = mx.abs(indices - idx) / inv_scale + w = kernel_fn(dist) + # Zero out-of-bounds taps + w = mx.where((idx >= 0) & (idx < N), w, 0.0) + all_idx.append(idx) + all_w.append(w) + + # Normalize so weights sum to 1 per output pixel + w_sum = sum(all_w) + w_sum = mx.where(w_sum > 0, w_sum, 1.0) + + result = [] + for idx, w in zip(all_idx, all_w): + w = mx.expand_dims(w / w_sum, -1) + idx = mx.clip(idx, a_min=0, a_max=N - 1).astype(mx.uint32) + result.append((idx, w)) + + return tuple(result) + + +def _triangle_kernel(x): + """Triangle (linear) filter kernel. Support radius = 1.""" + return mx.maximum(1.0 - x, 0.0) + + +def _cubic_kernel(x): + """Keys cubic kernel with a=-0.5 (PIL/Pillow convention). + + This coefficient is used by PyTorch when antialias=True for both + bilinear and bicubic modes. The non-antialiased cubic path uses + a=-0.75 (OpenCV convention) -- see ``_cubic_indices``. + + Support radius = 2. + """ + a = -0.5 + w_inner = ((a + 2.0) * x - (a + 3.0)) * x * x + 1 + w_outer = (((x - 5) * x + 8) * x - 4) * a + return mx.where(x <= 1.0, w_inner, mx.where(x <= 2.0, w_outer, 0.0)) + + +def _linear_aa_indices(N, scale, align_corners, dim, ndims): + """Linear interpolation with antialiasing (triangle kernel).""" + return _aa_indices( + N, + scale, + align_corners, + dim, + ndims, + kernel_fn=_triangle_kernel, + kernel_radius=1.0, + ) + + +def _cubic_aa_indices(N, scale, align_corners, dim, ndims): + """Cubic interpolation with antialiasing (Keys cubic, a=-0.5). + + Note: the non-antialiased cubic path (``_cubic_indices``) uses a=-0.75 + (OpenCV convention). When ``antialias=True``, PyTorch switches to a=-0.5 + (PIL convention). This coefficient change affects the interpolant shape, + not just the filter width. See ``_cubic_kernel`` for details. + """ + return _aa_indices( + N, + scale, + align_corners, + dim, + ndims, + kernel_fn=_cubic_kernel, + kernel_radius=2.0, + ) + + def _cubic_indices(N, scale, align_corners, dim, ndims): indices = _scaled_indices(N, scale, align_corners, dim, ndims) indices_l1 = mx.floor(indices) @@ -60,8 +173,9 @@ def _cubic_indices(N, scale, align_corners, dim, ndims): @partial(mx.compile, shapeless=True) def _get_weight(ind, grid, dist): - # PyTorch uses -0.5 for antialiasing=true (compatibility with PIL) - # and uses -0.75 for antialiasing=false (compatibility with OpenCV) + # a=-0.75 (OpenCV convention) for non-antialiased cubic. + # When antialias=True, _cubic_aa_indices uses a=-0.5 (PIL convention) + # via _cubic_kernel instead. a = -0.75 x = mx.abs(ind - grid) if dist == 1: @@ -89,6 +203,14 @@ def _get_weight(ind, grid, dist): ) +def _validate_antialias_align_corners(align_corners, antialias): + if antialias and align_corners: + raise ValueError( + "[Upsample] antialias=True with align_corners=True is not " + "supported. Use align_corners=False for antialiased interpolation." + ) + + def upsample_nearest(x: mx.array, scale_factor: Tuple): dims = x.ndim - 2 if dims != len(scale_factor): @@ -145,7 +267,41 @@ def _interpolate( return sum(wi * xi for wi, xi in zip(weights, samples)) -def upsample_linear(x: mx.array, scale_factor: Tuple, align_corners: bool = False): +def _interpolate_separable( + x: mx.array, scale_factor: Tuple, indices_fn: Callable, align_corners: bool = False +): + dims = x.ndim - 2 + if dims != len(scale_factor): + raise ValueError("A scale needs to be provided for each spatial dimension") + + _, *N, _ = x.shape + out = x + + for i, (n, s) in enumerate(zip(N, scale_factor)): + axis = i + 1 + samples = [] + for idx, weight in indices_fn(n, s, align_corners, i, dims): + sample = mx.take(out, idx.reshape(-1), axis=axis) + samples.append(sample * weight) + out = sum(samples) + + return out + + +def upsample_linear( + x: mx.array, + scale_factor: Tuple, + align_corners: bool = False, + antialias: bool = False, +): + _validate_antialias_align_corners(align_corners, antialias) + if antialias: + return _interpolate_separable( + x=x, + scale_factor=scale_factor, + indices_fn=_linear_aa_indices, + align_corners=align_corners, + ) return _interpolate( x=x, scale_factor=scale_factor, @@ -154,7 +310,20 @@ def upsample_linear(x: mx.array, scale_factor: Tuple, align_corners: bool = Fals ) -def upsample_cubic(x: mx.array, scale_factor: Tuple, align_corners: bool = False): +def upsample_cubic( + x: mx.array, + scale_factor: Tuple, + align_corners: bool = False, + antialias: bool = False, +): + _validate_antialias_align_corners(align_corners, antialias) + if antialias: + return _interpolate_separable( + x=x, + scale_factor=scale_factor, + indices_fn=_cubic_aa_indices, + align_corners=align_corners, + ) return _interpolate( x=x, scale_factor=scale_factor, @@ -185,6 +354,23 @@ class Upsample(Module): ``align_corners=True`` then the top and left edge of the input and output will be matching as will the bottom right edge. + .. note:: + When ``antialias=True`` is used with ``"linear"`` or ``"cubic"`` mode, + an antialiased filter is applied during downsampling (scale factor < 1), + producing smoother results by avoiding aliasing artifacts. For 2D + integer-ratio downscales with ``align_corners=False``, this matches the + behavior of PyTorch's ``F.interpolate(antialias=True)``. Non-integer + scale factors are supported but may differ from PyTorch because of + existing index-selection differences. + + For ``"cubic"`` mode, enabling ``antialias`` also changes the cubic + kernel coefficient from ``a=-0.75`` (OpenCV convention) to ``a=-0.5`` + (PIL/Pillow convention), matching PyTorch's behavior. This affects the + interpolant shape, not just the filter width. + + ``antialias=True`` with ``align_corners=True`` is not supported and + will raise a ``ValueError``. + Parameters: scale_factor (float or tuple): The multiplier for the spatial size. If a ``float`` is provided, it is the multiplier for all spatial dimensions. @@ -195,6 +381,11 @@ class Upsample(Module): align_corners (bool, optional): Changes the way the corners are treated during ``"linear"`` and ``"cubic"`` upsampling. See the note above and the examples below for more details. Default: ``False``. + antialias (bool, optional): If ``True``, apply an antialiasing filter + when downsampling with ``"linear"`` or ``"cubic"`` mode. For + ``"cubic"`` mode this also switches the kernel coefficient to + ``a=-0.5``. Not supported with ``"nearest"`` mode or with + ``align_corners=True``. Default: ``False``. Examples: >>> import mlx.core as mx @@ -230,22 +421,35 @@ def __init__( scale_factor: Union[float, Tuple], mode: Literal["nearest", "linear", "cubic"] = "nearest", align_corners: bool = False, + antialias: bool = False, ): super().__init__() if mode not in ["nearest", "linear", "cubic"]: raise ValueError(f"[Upsample] Got unsupported upsampling algorithm: {mode}") + if antialias and mode == "nearest": + raise ValueError( + "[Upsample] Antialiasing is not supported for nearest neighbor upsampling" + ) if isinstance(scale_factor, (list, tuple)): - self.scale_factor = tuple(map(float, scale_factor)) + scale_factor = tuple(map(float, scale_factor)) else: - self.scale_factor = float(scale_factor) + scale_factor = float(scale_factor) + + _validate_antialias_align_corners(align_corners, antialias) + + self.scale_factor = scale_factor self.mode = mode self.align_corners = align_corners + self.antialias = antialias def _extra_repr(self) -> str: - return ( + repr_str = ( f"scale_factor={self.scale_factor}, mode={self.mode!r}, " f"align_corners={self.align_corners}" ) + if self.antialias: + repr_str += ", antialias=True" + return repr_str def __call__(self, x: mx.array) -> mx.array: dims = x.ndim - 2 @@ -270,8 +474,8 @@ def __call__(self, x: mx.array) -> mx.array: if self.mode == "nearest": return upsample_nearest(x, scale_factor) elif self.mode == "linear": - return upsample_linear(x, scale_factor, self.align_corners) + return upsample_linear(x, scale_factor, self.align_corners, self.antialias) elif self.mode == "cubic": - return upsample_cubic(x, scale_factor, self.align_corners) + return upsample_cubic(x, scale_factor, self.align_corners, self.antialias) else: raise Exception(f"Unknown interpolation mode: {self.mode}") diff --git a/python/tests/test_upsample.py b/python/tests/test_upsample.py index 631853cce0..9a8ab9184c 100644 --- a/python/tests/test_upsample.py +++ b/python/tests/test_upsample.py @@ -95,6 +95,223 @@ def run_upsample( dtype=dtype, ) + @unittest.skipIf(not has_torch, "requires Torch") + def test_torch_upsample_antialias(self): + """Test antialiased downsampling matches PyTorch F.interpolate(antialias=True).""" + + def run_antialias( + N, + C, + idim, + scale_factor, + mode, + align_corners=False, + dtype="float32", + atol=1e-5, + ): + with self.subTest( + N=N, + C=C, + idim=idim, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + ): + np_dtype = getattr(np, dtype) + np.random.seed(0) + iH, iW = idim + in_np = np.random.normal(-1.0, 1.0, (N, iH, iW, C)).astype(np_dtype) + + in_mx = mx.array(in_np) + in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu") + + out_mx = nn.Upsample( + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + antialias=True, + )(in_mx) + mode_pt = { + "linear": "bilinear", + "cubic": "bicubic", + }[mode] + out_pt = F.interpolate( + in_pt, + scale_factor=scale_factor, + mode=mode_pt, + align_corners=align_corners, + antialias=True, + ) + out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) + self.assertEqual(out_pt.shape, out_mx.shape) + self.assertTrue( + np.allclose(out_pt, out_mx, atol=atol), + f"antialias {mode} ac={align_corners} scale={scale_factor} max_diff=" + f"{np.abs(out_pt - np.array(out_mx)).max():.2e}", + ) + + for dtype in ("float32",): + for N, C in ((1, 1), (2, 3)): + # Test downscale with antialias — use integer-ratio scales + # to avoid the pre-existing _scaled_indices step divergence + # for non-integer ratios (see issue #2186). + for idim, scale_factor in ( + ((4, 4), (0.5, 0.5)), + ((8, 8), (0.5, 0.5)), + ((8, 8), (0.25, 0.25)), + ((16, 16), (0.5, 0.5)), + ((16, 16), (0.25, 0.25)), + ((32, 32), (0.5, 0.5)), + ((32, 32), (0.25, 0.25)), + ((64, 64), (0.5, 0.5)), + ((10, 10), (0.5, 0.5)), + ((12, 12), (0.25, 0.25)), + ((8, 16), (0.5, 0.5)), # non-square + ((16, 8), (0.5, 0.25)), # different scales per dim + ): + for mode in ("linear", "cubic"): + # align_corners=True + antialias has a known + # interaction with _scaled_indices that requires + # further work. Test align_corners=False for now. + run_antialias( + N, + C, + idim, + scale_factor, + mode, + align_corners=False, + dtype=dtype, + ) + + @unittest.skipIf(not has_torch, "requires Torch") + def test_antialias_upscale_linear_is_noop(self): + """For linear mode, antialias has no effect on upscaling.""" + np.random.seed(0) + in_np = np.random.normal(-1.0, 1.0, (1, 4, 4, 3)).astype(np.float32) + in_mx = mx.array(in_np) + + for scale in (2.0, 3.0): + with self.subTest(scale=scale): + out_aa = nn.Upsample( + scale_factor=scale, + mode="linear", + align_corners=False, + antialias=True, + )(in_mx) + out_no = nn.Upsample( + scale_factor=scale, + mode="linear", + align_corners=False, + antialias=False, + )(in_mx) + self.assertTrue( + np.allclose(np.array(out_aa), np.array(out_no), atol=1e-7), + "linear antialias should be no-op for upscaling", + ) + + @unittest.skipIf(not has_torch, "requires Torch") + def test_antialias_upscale_cubic_matches_pytorch(self): + """For cubic mode, antialias changes a from -0.75 to -0.5 even on upscale.""" + np.random.seed(0) + in_np = np.random.normal(-1.0, 1.0, (1, 8, 8, 3)).astype(np.float32) + in_mx = mx.array(in_np) + in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)) + + for scale in (2.0, 3.0): + with self.subTest(scale=scale): + out_mx = nn.Upsample( + scale_factor=scale, + mode="cubic", + align_corners=False, + antialias=True, + )(in_mx) + out_pt = F.interpolate( + in_pt, + scale_factor=scale, + mode="bicubic", + align_corners=False, + antialias=True, + ) + out_pt_np = out_pt.permute(0, 2, 3, 1).numpy(force=True) + self.assertTrue( + np.allclose(np.array(out_mx), out_pt_np, atol=1e-5), + f"cubic antialias upscale {scale}x max_diff=" + f"{np.abs(np.array(out_mx) - out_pt_np).max():.2e}", + ) + + def test_antialias_non_integer_scale_smoke(self): + """Smoke test for non-integer scale factors (no PyTorch comparison).""" + np.random.seed(42) + in_np = np.random.normal(0, 1, (1, 32, 32, 3)).astype(np.float32) + in_mx = mx.array(in_np) + + for scale in (0.3, 0.7, 0.6): + for mode in ("linear", "cubic"): + with self.subTest(scale=scale, mode=mode): + out = nn.Upsample( + scale_factor=scale, + mode=mode, + align_corners=False, + antialias=True, + )(in_mx) + mx.eval(out) + out_np = np.array(out) + + # Correct shape + expected = int(32 * scale) + self.assertEqual(out_np.shape, (1, expected, expected, 3)) + + self.assertTrue(np.all(np.isfinite(out_np))) + if mode == "linear": + # Linear AA uses non-negative triangle weights, so it is + # bounded by the input range. Cubic interpolation can + # overshoot because the Keys kernel has negative lobes. + self.assertLessEqual( + out_np.max(), + in_np.max() + 0.01, + "linear AA output should not exceed input range", + ) + self.assertGreaterEqual( + out_np.min(), + in_np.min() - 0.01, + "linear AA output should not go below input range", + ) + + def test_antialias_1d_smoke(self): + """Test that antialias works on 1D spatial input (3D tensor). + + PyTorch does not support antialias on 1D tensors, so this is a + smoke test only (correct shape, non-trivial, smoother than non-AA). + """ + np.random.seed(0) + for length, scale in ((16, 0.5), (32, 0.25)): + for mode in ("linear", "cubic"): + with self.subTest(length=length, scale=scale, mode=mode): + in_np = np.random.normal(0, 1, (1, length, 3)).astype(np.float32) + in_mx = mx.array(in_np) + + out_aa = nn.Upsample( + scale_factor=scale, + mode=mode, + align_corners=False, + antialias=True, + )(in_mx) + out_no = nn.Upsample( + scale_factor=scale, + mode=mode, + align_corners=False, + antialias=False, + )(in_mx) + mx.eval(out_aa, out_no) + + expected_len = int(length * scale) + self.assertEqual(out_aa.shape, (1, expected_len, 3)) + # AA should differ from non-AA + self.assertGreater( + float(mx.abs(out_aa - out_no).max()), + 1e-6, + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 9f43f8c54a1c5d80296af3e66782113cde496f34 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 13 Aug 2026 09:21:54 +0900 Subject: [PATCH 190/222] chore: Require pytorch 2.12 for metal dlpack tests (#4206) --- python/tests/test_array.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 5c6db31482..aee2ab8872 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -22,10 +22,19 @@ except ImportError: has_tf = False + try: import torch - has_torch_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + torch_version = [int(v) for v in torch.__version__.split("+")[0].split(".")] + is_torch_212 = torch_version[0] > 2 or ( + torch_version[0] == 2 and torch_version[1] >= 12 + ) + has_torch_mps = ( + is_torch_212 + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ) except ImportError: torch = None has_torch_mps = False From bb6d960fb0353ee3c5012b7be509dc23e803dba6 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 13 Aug 2026 09:22:19 +0900 Subject: [PATCH 191/222] chore: Remove std::optional in C++ public interface (#4207) --- mlx/ops.cpp | 14 ++++++++------ mlx/ops.h | 6 +----- python/src/array.cpp | 5 +---- python/src/convert.cpp | 13 +++++++++++-- python/src/convert.h | 1 + 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index ce29cb4e3d..d90a95d3a2 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -293,12 +293,10 @@ array linspace( s); } -array astype( - array a, - Dtype dtype, - std::optional copy, - StreamOrDevice s /* = {} */) { - if (dtype == a.dtype() && !copy.value_or(false)) { +// Private API used by python bindings. +MLX_API array +astype(array a, Dtype dtype, bool force_copy, StreamOrDevice s = {}) { + if (dtype == a.dtype() && !force_copy) { return a; } auto copied_shape = a.shape(); // |a| will be moved @@ -309,6 +307,10 @@ array astype( {std::move(a)}); } +array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) { + return astype(std::move(a), dtype, false, s); +} + array as_strided( array a, Shape shape, diff --git a/mlx/ops.h b/mlx/ops.h index cacc4e28d6..01e0a99286 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -47,11 +47,7 @@ MLX_API array linspace( StreamOrDevice s = {}); /** Convert an array to the given data type. */ -MLX_API array -astype(array a, Dtype dtype, std::optional copy, StreamOrDevice s = {}); -inline array astype(array a, Dtype dtype, StreamOrDevice s = {}) { - return astype(std::move(a), dtype, std::nullopt, s); -} +MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {}); /** Create a view of an array with the given shape and strides. */ MLX_API array as_strided( diff --git a/python/src/array.cpp b/python/src/array.cpp index 6f65e14cde..fbab846182 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -517,10 +517,7 @@ void init_array(nb::module_& m) { nb::object, std::optional> dl_device, std::optional copy) { - if (copy.value_or(false)) { - return mlx_to_dlpack(mx::astype(a, a.dtype(), true), dl_device); - } - return mlx_to_dlpack(a, dl_device); + return mlx_to_dlpack(a, copy.value_or(false), dl_device); }, nb::kw_only(), "stream"_a = nb::none(), diff --git a/python/src/convert.cpp b/python/src/convert.cpp index 3b161b99e6..9941358e84 100644 --- a/python/src/convert.cpp +++ b/python/src/convert.cpp @@ -19,6 +19,11 @@ #include "mlx/ops.h" #include "mlx/utils.h" +// Defined in ops.cpp. +namespace mlx::core { +array astype(array a, Dtype dtype, bool force_copy, StreamOrDevice s = {}); +} + enum PyScalarT { pybool = 0, pyint = 1, @@ -374,7 +379,11 @@ nb::ndarray mlx_to_np_array(const mx::array& a) { nb::ndarray<> mlx_to_dlpack( const mx::array& a, + bool force_copy, std::optional> dl_device) { + if (force_copy) { + return mlx_to_nd_array<>(mx::astype(a, a.dtype(), true), dl_device); + } return mlx_to_nd_array<>(a, dl_device); } @@ -728,9 +737,9 @@ mx::array create_array( } else if (nb::isinstance(v)) { auto arr = nb::cast(v); auto dtype = t.value_or(arr.dtype()); - return mx::astype(arr, dtype, copy); + return mx::astype(arr, dtype, copy.value_or(false)); } else { auto arr = to_array_with_accessor(v); - return mx::astype(arr, t.value_or(arr.dtype()), copy); + return mx::astype(arr, t.value_or(arr.dtype()), copy.value_or(false)); } } diff --git a/python/src/convert.h b/python/src/convert.h index d9be9a8cb0..133f443ed7 100644 --- a/python/src/convert.h +++ b/python/src/convert.h @@ -70,6 +70,7 @@ mx::array nd_array_to_mlx( nb::ndarray mlx_to_np_array(const mx::array& a); nb::ndarray<> mlx_to_dlpack( const mx::array& a, + bool force_copy, std::optional> dl_device); nb::object to_scalar(mx::array& a); From 7729d5872e7825e552563c630a14bfac35e4cd28 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Wed, 12 Aug 2026 18:42:50 -0700 Subject: [PATCH 192/222] Fix SliceUpdate JVPs with one traced input (#4200) --- mlx/primitives.cpp | 17 ++++++++++------- tests/autograd_tests.cpp | 24 ++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index cec010647f..3bafd19407 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -5177,13 +5177,16 @@ std::vector SliceUpdate::jvp( // Check inputs assert(primals.size() == 2); - if (argnums.size() != 2) { - throw std::runtime_error( - "[SliceUpdate] JVP for one argument not implemented yet."); + array result_tan = zeros_like(primals[0], stream()); + array update_tan = zeros_like(primals[1], stream()); + for (int i = 0; i < argnums.size(); ++i) { + if (argnums[i] == 0) { + result_tan = tangents[i]; + } else if (argnums[i] == 1) { + update_tan = tangents[i]; + } } - auto result_tan = tangents[0]; - switch (reduce_type_) { case SliceUpdate::None: return {array( @@ -5191,14 +5194,14 @@ std::vector SliceUpdate::jvp( result_tan.dtype(), std::make_shared( stream(), reduce_type_, start_indices_, end_indices_, strides_), - {result_tan, tangents[1]})}; + {result_tan, update_tan})}; case SliceUpdate::Sum: return {array( result_tan.shape(), result_tan.dtype(), std::make_shared( stream(), reduce_type_, start_indices_, end_indices_, strides_), - {result_tan, tangents[1]})}; + {result_tan, update_tan})}; case SliceUpdate::Prod: case SliceUpdate::Max: case SliceUpdate::Min: { diff --git a/tests/autograd_tests.cpp b/tests/autograd_tests.cpp index 199af55010..5a0aca0800 100644 --- a/tests/autograd_tests.cpp +++ b/tests/autograd_tests.cpp @@ -849,6 +849,30 @@ TEST_CASE("test slice grads") { CHECK_EQ(out.size(), 0); } +TEST_CASE("test slice update jvp with one tangent") { + auto src = array({1.0f, 2.0f, 3.0f, 4.0f}); + auto update = array({5.0f, 6.0f}); + auto src_tan = array({1.0f, 2.0f, 3.0f, 4.0f}); + auto update_tan = array({7.0f, 8.0f}); + + for (bool add : {false, true}) { + auto update_fn = [&src, add](array x) { + return add ? slice_update_add(src, x, Shape{1}, Shape{3}) + : slice_update(src, x, Shape{1}, Shape{3}); + }; + auto out = jvp(update_fn, update, update_tan).second; + CHECK(array_equal(out, array({0.0f, 7.0f, 8.0f, 0.0f})).item()); + + auto src_fn = [&update, add](array x) { + return add ? slice_update_add(x, update, Shape{1}, Shape{3}) + : slice_update(x, update, Shape{1}, Shape{3}); + }; + out = jvp(src_fn, src, src_tan).second; + auto expected = add ? src_tan : array({1.0f, 0.0f, 0.0f, 4.0f}); + CHECK(array_equal(out, expected).item()); + } +} + TEST_CASE("test min and max vjp") { // Test min { From 3f2e4a32bb53a9a76e44f37b5565c09ec1119932 Mon Sep 17 00:00:00 2001 From: anchor Date: Thu, 13 Aug 2026 15:09:03 +0800 Subject: [PATCH 193/222] Raise instead of hanging when reflect/symmetric pad gets an empty axis (#4223) Co-authored-by: codeAnqiang-ma <273298913+codeAnqiang-ma@users.noreply.github.com> --- mlx/ops.cpp | 6 ++++++ python/tests/test_ops.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index d90a95d3a2..9c8db3a26d 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1489,6 +1489,12 @@ array reflect_pad( if (L == 0 && H == 0) { continue; } + if (n == 0) { + std::ostringstream msg; + msg << "[pad] Cannot pad empty axis " << ax << " using mode '" + << (include_edge ? "symmetric" : "reflect") << "'."; + throw std::invalid_argument(msg.str()); + } // reflect skips the edge value (period 2(n-1)); symmetric repeats it // (period 2n). int offset = (!include_edge && n > 1) ? 1 : 0; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d569dbd3c7..1b46237ce5 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2305,6 +2305,17 @@ def test_pad_reflect_symmetric(self): ) self.assertEqual(b_mlx.dtype, mx.float32) + # An empty axis cannot be extended; numpy raises for these too. + # Used to hang in an infinite loop rather than raise. + for mode in ("reflect", "symmetric"): + with self.assertRaises(ValueError): + mx.pad(mx.array([]), 2, mode=mode) + with self.assertRaises(ValueError): + mx.pad(mx.zeros((0, 3)), [(1, 1), (0, 0)], mode=mode) + # A zero-width pad on the empty axis stays allowed + out = mx.pad(mx.zeros((0, 3)), [(0, 0), (2, 1)], mode=mode) + self.assertEqual(out.shape, (0, 6)) + def test_as_strided(self): x_npy = np.random.randn(128).astype(np.float32) x_mlx = mx.array(x_npy) From d56ed249415b6b767d1a65b2eb19a99e1cd9fdb9 Mon Sep 17 00:00:00 2001 From: PhysicistJohn <54456354+PhysicistJohn@users.noreply.github.com> Date: Thu, 13 Aug 2026 03:13:17 -0700 Subject: [PATCH 194/222] Reuse Stockham twiddles in large batched Bluestein FFTs (#4084) Co-authored-by: Cheng --- mlx/backend/metal/fft.cpp | 64 ++++++++++++++-- mlx/backend/metal/jit_kernels.cpp | 12 +++ mlx/backend/metal/kernels.h | 5 ++ mlx/backend/metal/kernels/fft.h | 113 ++++++++++++++++++++++++---- mlx/backend/metal/nojit_kernels.cpp | 7 ++ python/tests/test_fft.py | 23 ++++++ 6 files changed, 202 insertions(+), 22 deletions(-) diff --git a/mlx/backend/metal/fft.cpp b/mlx/backend/metal/fft.cpp index 6c7a931390..3c9195f8ea 100644 --- a/mlx/backend/metal/fft.cpp +++ b/mlx/backend/metal/fft.cpp @@ -27,6 +27,10 @@ using MTLFC = std::tuple; #define MIN_THREADGROUP_MEM_SIZE 256 // For strided reads/writes, coalesce at least this many complex64s #define MIN_COALESCE_WIDTH 4 +// Precomputed radix twiddles pay off for the largest batched fused plan. +constexpr int BLUESTEIN_TWIDDLE_TABLE_FFT_SIZE = 4096; +constexpr int BLUESTEIN_TWIDDLE_TABLE_SIZE = 584; +constexpr int MIN_BLUESTEIN_TWIDDLE_TABLE_BATCH = 1024; inline const std::vector supported_radices() { // Ordered by preference in decomposition. @@ -316,7 +320,10 @@ std::tuple compute_raders_constants( } // Bluestein -std::pair compute_bluestein_constants(int n, int bluestein_n) { +std::pair compute_bluestein_constants( + int n, + int bluestein_n, + int radix_twiddle_size = 0) { // We need to calculate the Bluestein twiddle factors // in double precision for the overall numerical stability // of Bluestein's FFT algorithm to be acceptable. @@ -343,7 +350,7 @@ std::pair compute_bluestein_constants(int n, int bluestein_n) { w_k.set_data(allocator::malloc(w_k.nbytes())); std::copy(w_k_vec.begin(), w_k_vec.end(), w_k.data()); - array w_q({bluestein_n}, complex64, nullptr, {}); + array w_q({bluestein_n + radix_twiddle_size}, complex64, nullptr, {}); w_q.set_data(allocator::malloc(w_q.nbytes())); auto w_q_ptr = reinterpret_cast*>(w_q.data()); @@ -649,6 +656,15 @@ void fft_op( size = out.size(); } int total_batch_size = size / n; + static const std::vector bluestein_twiddle_table_plan = { + 0, 0, 4, 0, 0, 0, 0, 0, 0}; + bool use_bluestein_twiddle_table = !real && + plan.bluestein_n == BLUESTEIN_TWIDDLE_TABLE_FFT_SIZE && + plan.stockham == bluestein_twiddle_table_plan && + total_batch_size >= MIN_BLUESTEIN_TWIDDLE_TABLE_BATCH; + if (plan.bluestein_n > 0 && !real) { + func_consts.push_back(make_bool(&use_bluestein_twiddle_table, 22)); + } int threads_per_fft = (fft_size + elems_per_thread - 1) / elems_per_thread; // We batch among threadgroups for improved efficiency when n is small @@ -702,6 +718,9 @@ void fft_op( std::string base_name = kname.str(); // We use a specialized kernel for each FFT size kname << "_n" << fft_size << "_inv_" << inverse; + if (use_bluestein_twiddle_table) { + kname << "_precomputed"; + } std::string hash_name = kname.str(); auto template_def = func_name == "four_step_fft" ? get_template_definition( base_name, @@ -720,22 +739,45 @@ void fft_op( auto kernel = get_fft_kernel(d, base_name, hash_name, func_consts, template_def); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(in_contiguous, 0); - compute_encoder.set_output_array(out, 1); - if (plan.bluestein_n > 0) { // Precomputed twiddle factors for Bluestein's - auto [w_k, w_q] = compute_bluestein_constants(n, plan.bluestein_n); + auto [w_k, w_q] = use_bluestein_twiddle_table + ? compute_bluestein_constants( + n, plan.bluestein_n, BLUESTEIN_TWIDDLE_TABLE_SIZE) + : compute_bluestein_constants(n, plan.bluestein_n); copies.push_back(w_q); copies.push_back(w_k); + if (use_bluestein_twiddle_table) { + auto twiddle_kernel = + get_fft_twiddle_kernel(d, base_name, template_def); + compute_encoder.set_compute_pipeline_state(twiddle_kernel); + compute_encoder.set_output_array( + w_q, + 0, + plan.bluestein_n * static_cast(sizeof(complex64_t))); + auto twiddle_group_size = std::min( + static_cast(BLUESTEIN_TWIDDLE_TABLE_SIZE), + twiddle_kernel->maxTotalThreadsPerThreadgroup()); + compute_encoder.dispatch_threads( + MTL::Size(BLUESTEIN_TWIDDLE_TABLE_SIZE, 1, 1), + MTL::Size(twiddle_group_size, 1, 1)); + } + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_input_array(w_q, 2); // w_q compute_encoder.set_input_array(w_k, 3); // w_k compute_encoder.set_bytes(n, 4); compute_encoder.set_bytes(plan.bluestein_n, 5); compute_encoder.set_bytes(total_batch_size, 6); } else if (plan.rader_n > 1) { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + auto [b_q, g_q, g_minus_q] = compute_raders_constants(plan.rader_n, s); copies.push_back(b_q); copies.push_back(g_q); @@ -748,10 +790,18 @@ void fft_op( compute_encoder.set_bytes(total_batch_size, 6); compute_encoder.set_bytes(plan.rader_n, 7); } else if (four_step_params.required) { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(four_step_params.n1, 2); compute_encoder.set_bytes(four_step_params.n2, 3); compute_encoder.set_bytes(total_batch_size, 4); } else { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(n, 2); compute_encoder.set_bytes(total_batch_size, 3); } diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index d40bea85b2..c7900ecdf8 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1008,6 +1008,18 @@ MTL::ComputePipelineState* get_fft_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string& library_name, + const std::string& template_def) { + auto lib = d.get_library(library_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::fft() << template_def; + return kernel_source.str(); + }); + return d.get_kernel("generate_bluestein_twiddles", lib); +} + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 3e4258c383..21b754514c 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -296,6 +296,11 @@ MTL::ComputePipelineState* get_fft_kernel( const metal::MTLFCList& func_consts, const std::string& template_def); +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string& library_name, + const std::string& template_def); + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels/fft.h b/mlx/backend/metal/kernels/fft.h index 408fc1a6d2..1d847331d0 100644 --- a/mlx/backend/metal/kernels/fft.h +++ b/mlx/backend/metal/kernels/fft.h @@ -45,19 +45,47 @@ STEEL_CONST int rader_5_steps_ [[function_constant(18)]]; STEEL_CONST int rader_4_steps_ [[function_constant(19)]]; STEEL_CONST int rader_3_steps_ [[function_constant(20)]]; STEEL_CONST int rader_2_steps_ [[function_constant(21)]]; +STEEL_CONST bool use_bluestein_twiddle_table_ [[function_constant(22)]]; + +// Generate the compact radix-8 Stockham twiddle table used by the fused +// 4096-point Bluestein convolution. The first p=1 stage needs no twiddles; +// the remaining stages contribute 8, 64, and 512 entries. +[[kernel]] void generate_bluestein_twiddles( + device float2* twiddles [[buffer(0)]], + uint index [[thread_position_in_grid]]) { + uint p; + uint k; + if (index < 8) { + p = 8; + k = index; + } else if (index < 72) { + p = 64; + k = index - 8; + } else { + p = 512; + k = index - 72; + } + twiddles[index] = get_twiddle(k, 8 * p); +} // See "radix.h" for radix codelets. template using RadixFunc = void (*)(thread vec*, thread vec*); // Perform a single radix n butterfly with appropriate twiddles -template radix_func> +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_butterfly( int i, int p, + int twiddle_offset, thread vec* x, thread short* indices, - thread vec* y) { + thread vec* y, + const device vec* twiddles) { // i: the index in the overall DFT that we're processing. // p: the size of the DFTs we're merging at this step. // m: how many threads are working on this DFT. @@ -76,7 +104,12 @@ METAL_FUNC void radix_butterfly( // Apply twiddles if (p > 1) { - vec twiddle_1 = get_twiddle(k, radix * p); + vec twiddle_1; + if constexpr (use_twiddle_table) { + twiddle_1 = twiddles[twiddle_offset + k]; + } else { + twiddle_1 = get_twiddle(k, radix * p); + } vec twiddle = twiddle_1; x[1] = complex_mul(x[1], twiddle); @@ -97,17 +130,23 @@ METAL_FUNC void radix_butterfly( // Perform all the radix steps required for a // particular radix size n. -template radix_func> +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_n_steps( int i, thread int* p, + thread int* twiddle_offset, int m, int n, int num_steps, thread vec* inputs, thread short* indices, thread vec* values, - threadgroup vec* buf) { + threadgroup vec* buf, + const device vec* twiddles) { int m_r = n / radix; // When combining different sized radices, we have to do // multiple butterflies in a single thread. @@ -127,8 +166,14 @@ METAL_FUNC void radix_n_steps( for (int r = 0; r < radix; r++) { inputs[r] = buf[index + r * m_r]; } - radix_butterfly( - index, *p, inputs, indices + t * radix, values + t * radix); + radix_butterfly( + index, + *p, + *twiddle_offset, + inputs, + indices + t * radix, + values + t * radix, + twiddles); } } @@ -147,24 +192,41 @@ METAL_FUNC void radix_n_steps( // Wait until all threads have written back to threadgroup mem threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (use_twiddle_table) { + if (*p > 1) { + *twiddle_offset += *p; + } + } *p *= radix; } } -#define RADIX_STEP(radix, radix_func, num_steps) \ - radix_n_steps>( \ - fft_idx, p, m, n, num_steps, inputs, indices, values, buf); - -template +#define RADIX_STEP(radix, radix_func, num_steps) \ + radix_n_steps>( \ + fft_idx, \ + p, \ + &twiddle_offset, \ + m, \ + n, \ + num_steps, \ + inputs, \ + indices, \ + values, \ + buf, \ + twiddles); + +template METAL_FUNC void perform_fft( int fft_idx, thread int* p, int m, int n, - threadgroup vec* buf) { + threadgroup vec* buf, + const device vec* twiddles = nullptr) { vec inputs[MAX_RADIX]; short indices[MAX_OUTPUT_SIZE]; vec values[MAX_OUTPUT_SIZE]; + int twiddle_offset = 0; RADIX_STEP(2, radix2, rader ? rader_2_steps_ : radix_2_steps_); RADIX_STEP(3, radix3, rader ? rader_3_steps_ : radix_3_steps_); @@ -437,9 +499,20 @@ template int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup threadgroup vec* buf = &shared_in[tg_idx]; + const device vec* twiddles = w_q + n; // fft - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } vec convolution_factor = {1.0f, -1.0f}; if constexpr (!metal::is_same_v) { @@ -458,7 +531,17 @@ template // ifft p = 1; - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } read_writer.write_padded(length, w_k); } diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index e0a3ff935d..5da78db8a3 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -353,6 +353,13 @@ MTL::ComputePipelineState* get_fft_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string&, + const std::string&) { + return d.get_kernel("generate_bluestein_twiddles"); +} + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 26f473f2ca..9358ede794 100644 --- a/python/tests/test_fft.py +++ b/python/tests/test_fft.py @@ -168,6 +168,29 @@ def test_fft_shared_mem(self): atol = 1e-4 if num < 1025 else 1e-3 self._run_ffts((batch_size, num), atol=atol) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_batched_bluestein_twiddle_table(self): + for batch, n in ((1023, 1031), (1024, 1031), (1024, 1531)): + with self.subTest(batch=batch, n=n), mx.stream(mx.gpu): + index = mx.arange(n, dtype=mx.float32) + base = mx.cos(index * 0.017) + 0.25 * mx.sin(index * 0.031) + base = base + 1j * ( + mx.sin(index * 0.023) - 0.125 * mx.cos(index * 0.047) + ) + scale_index = mx.arange(batch, dtype=mx.float32) + scales = (0.75 + 0.25 * mx.cos(scale_index * 0.013)) * ( + mx.cos(scale_index * 0.019) + 1j * mx.sin(scale_index * 0.019) + ) + signal = scales[:, None] * base[None, :] + + for transform, atol in ((mx.fft.fft, 1e-3), (mx.fft.ifft, 1e-4)): + expected = scales[:, None] * transform(base)[None, :] + output = transform(signal) + mx.eval(expected, output) + self.assertTrue( + mx.allclose(output, expected, atol=atol, rtol=1e-4).item() + ) + @unittest.skip("Too slow for CI but useful for local testing.") def test_fft_exhaustive(self): nums = range(2, 4097) From c6ef1eba84e001a605ec280e6bd8fa032f53e801 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Thu, 13 Aug 2026 03:16:52 -0700 Subject: [PATCH 195/222] chore: Validate decay_steps and step_size in the schedulers (#4217) --- python/mlx/optimizers/schedulers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/mlx/optimizers/schedulers.py b/python/mlx/optimizers/schedulers.py index 67e4e29cd2..99712097cb 100644 --- a/python/mlx/optimizers/schedulers.py +++ b/python/mlx/optimizers/schedulers.py @@ -51,6 +51,8 @@ def step_decay(init: float, decay_rate: float, step_size: int) -> Callable: >>> optimizer.learning_rate array(0.081, dtype=float32) """ + if step_size < 1: + raise ValueError(f"step_size must be greater than 0, but got {step_size}.") def schedule(step): return init * (decay_rate ** (step // step_size)) @@ -79,6 +81,8 @@ def cosine_decay(init: float, decay_steps: int, end: float = 0.0) -> Callable: >>> optimizer.learning_rate array(0.0999961, dtype=float32) """ + if decay_steps < 1: + raise ValueError(f"decay_steps must be greater than 0, but got {decay_steps}.") def schedule(step): s = mx.minimum(step, decay_steps) From fd7f0232dc84382f3f8356307f25ebb679607275 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Thu, 13 Aug 2026 06:19:10 -0400 Subject: [PATCH 196/222] chore: Add runnable distributed examples (#4219) --- examples/python/distributed_data_parallel.py | 102 ++++++++++++++++++ .../python/distributed_tensor_parallel.py | 72 +++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 examples/python/distributed_data_parallel.py create mode 100644 examples/python/distributed_tensor_parallel.py diff --git a/examples/python/distributed_data_parallel.py b/examples/python/distributed_data_parallel.py new file mode 100644 index 0000000000..6f19766ddf --- /dev/null +++ b/examples/python/distributed_data_parallel.py @@ -0,0 +1,102 @@ +# Copyright © 2026 Apple Inc. + +""" +Data parallel training of a small MLP. + +Every rank holds a different slice of the same dataset and averages the +gradients each step, so the training is equivalent to running on one process +and the final loss is the same at any number of ranks, up to the order the +floating point additions happen in: + + python examples/python/distributed_data_parallel.py + mlx.launch -n 2 python examples/python/distributed_data_parallel.py + mlx.launch -n 4 python examples/python/distributed_data_parallel.py + +The model is replicated and the batch is split, which is what you reach for +when the model fits on one machine but the data is large. +""" + +import time + +import mlx.core as mx +import mlx.nn as nn +import mlx.optimizers as optim + +num_features = 100 +num_examples = 1_000 +hidden = 64 +num_iters = 200 +lr = 0.05 + +world = mx.distributed.init() + +if num_examples % world.size() != 0: + raise ValueError( + f"Cannot split {num_examples} examples evenly over {world.size()} ranks." + ) + +# Fixed keys, so every rank draws the same dataset rather than one of its own. +w_star = mx.random.normal((num_features,), key=mx.random.key(0)) +X = mx.random.normal((num_examples, num_features), key=mx.random.key(1)) +y = X @ w_star + 1e-2 * mx.random.normal((num_examples,), key=mx.random.key(2)) + +# Keep this rank's slice and drop the rest. +examples_per_rank = num_examples // world.size() +start = world.rank() * examples_per_rank +X = X[start : start + examples_per_rank] +y = y[start : start + examples_per_rank] + + +class MLP(nn.Module): + def __init__(self, dims: int, hidden: int): + super().__init__() + self.layers = [nn.Linear(dims, hidden), nn.Linear(hidden, 1)] + + def __call__(self, x): + return self.layers[1](nn.relu(self.layers[0](x))).squeeze(-1) + + +# Seeding the global rng starts every rank from the same weights, which the +# averaged gradient then keeps in step. +mx.random.seed(0) +model = MLP(num_features, hidden) +mx.eval(model.parameters()) + +optimizer = optim.SGD(learning_rate=lr) + + +def loss_fn(model, X, y): + return 0.5 * mx.mean(mx.square(model(X) - y)) + + +loss_and_grad_fn = nn.value_and_grad(model, loss_fn) + +tic = time.perf_counter() +for _ in range(num_iters): + loss, grads = loss_and_grad_fn(model, X, y) + + # Each rank has gradients for its own slice. Averaging them gives the + # gradients of the whole dataset, which is what keeps this equivalent to + # single process training. One call handles the whole parameter tree. + grads = nn.average_gradients(grads, group=world) + + optimizer.update(model, grads) + mx.eval(model.parameters(), optimizer.state) +toc = time.perf_counter() + +# Every slice is the same size, so averaging the per rank losses gives the loss +# over the whole dataset. +loss = mx.distributed.all_sum(loss_fn(model, X, y), group=world) / world.size() + +# Only rank 0 prints the loss, but every rank has to evaluate it. Arrays are +# lazy, so leaving this to the print below would mean the other ranks never +# join the all_sum and everyone waits forever. +mx.eval(loss) + +throughput = num_iters / (toc - tic) + +if world.rank() == 0: + print( + f"Loss {loss.item():.6f}, Throughput {throughput:.2f} (it/s) " + f"over {world.size()} rank(s), {examples_per_rank} examples per rank" + ) diff --git a/examples/python/distributed_tensor_parallel.py b/examples/python/distributed_tensor_parallel.py new file mode 100644 index 0000000000..db66c4c65b --- /dev/null +++ b/examples/python/distributed_tensor_parallel.py @@ -0,0 +1,72 @@ +# Copyright © 2026 Apple Inc. + +""" +Tensor parallel inference for a small MLP. + +The two linear layers are sharded across ranks: the first splits its output +features so each rank computes part of the hidden state, the second splits its +input features and sums the partial results. That costs one all reduce per +block, and the sharded model returns what the full model returns: + + python examples/python/distributed_tensor_parallel.py + mlx.launch -n 2 python examples/python/distributed_tensor_parallel.py + mlx.launch -n 4 python examples/python/distributed_tensor_parallel.py + +Unlike data parallelism this splits the model rather than the batch, so it is +what you reach for when the weights are too big for one machine. +""" + +import mlx.core as mx +import mlx.nn as nn + +dims = 256 +hidden = 1024 +num_tokens = 8 + +world = mx.distributed.init() + +if hidden % world.size() != 0: + raise ValueError( + f"Cannot split {hidden} hidden features evenly over {world.size()} ranks." + ) + + +class MLP(nn.Module): + def __init__(self, dims: int, hidden: int): + super().__init__() + self.up = nn.Linear(dims, hidden) + self.down = nn.Linear(hidden, dims) + + def __call__(self, x): + return self.down(nn.silu(self.up(x))) + + +# Seeding the global rng gives every rank the same weights to shard. +mx.random.seed(0) +model = MLP(dims, hidden) +mx.eval(model.parameters()) + +x = mx.random.normal((num_tokens, dims), key=mx.random.key(1)) +expected = model(x) +mx.eval(expected) + +# Each rank keeps a slice of each weight. from_linear takes the slice belonging +# to this rank, so the layers never hold the full weight afterwards. +model.up = nn.AllToShardedLinear.from_linear(model.up, group=world) +model.down = nn.ShardedToAllLinear.from_linear(model.down, group=world) +mx.eval(model.parameters()) + +y = model(x) + +# Every rank evaluates: the down projection ends in an all reduce, so a rank +# that skipped this would leave the others waiting. +mx.eval(y) + +difference = mx.abs(y - expected).max().item() +hidden_per_rank = hidden // world.size() + +if world.rank() == 0: + print( + f"Max |sharded - full| = {difference:.3e} over {world.size()} rank(s), " + f"{hidden_per_rank} of {hidden} hidden features per rank" + ) From 09ebe730ba7882f07739958789f579e9d2dfda4a Mon Sep 17 00:00:00 2001 From: JamesMcCarthy44 Date: Thu, 13 Aug 2026 12:18:06 +0100 Subject: [PATCH 197/222] Break monolithic MTLResidencySet into smaller sets (#4211) --- mlx/backend/metal/allocator.cpp | 16 +- mlx/backend/metal/allocator.h | 2 +- mlx/backend/metal/device.cpp | 14 +- mlx/backend/metal/device.h | 12 +- mlx/backend/metal/eval.cpp | 4 +- mlx/backend/metal/resident.cpp | 250 +++++++++++++++++++++------ mlx/backend/metal/resident.h | 113 ++++++++++-- mlx/utils.h | 15 ++ tests/CMakeLists.txt | 4 + tests/residency_tests.cpp | 296 ++++++++++++++++++++++++++++++++ 10 files changed, 640 insertions(+), 86 deletions(-) create mode 100644 tests/residency_tests.cpp diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index 60459c67c2..903cce2187 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -44,13 +44,13 @@ namespace metal { MetalAllocator::MetalAllocator(Device& d) : device_(d.mtl_device()), - residency_set_(d.residency_set()), + residency_sets_(d.residency_sets()), buffer_cache_( vm_page_size, [](MTL::Buffer* buf) { return buf->length(); }, [this](MTL::Buffer* buf) { if (!buf->heap()) { - residency_set_.erase(buf); + residency_sets_.erase(buf); } auto pool = metal::new_scoped_memory_pool(); buf->release(); @@ -73,7 +73,7 @@ MetalAllocator::MetalAllocator(Device& d) heap_desc->setResourceOptions(resource_options); heap_desc->setSize(heap_size_); heap_ = NS::TransferPtr(device_->newHeap(heap_desc)); - residency_set_.insert(heap_.get()); + residency_sets_.insert(heap_.get()); } MetalAllocator::~MetalAllocator() = default; @@ -100,7 +100,7 @@ size_t MetalAllocator::get_memory_limit() { size_t MetalAllocator::set_wired_limit(size_t limit) { std::unique_lock lk(mutex_); std::swap(limit, wired_limit_); - residency_set_.resize(wired_limit_); + residency_sets_.resize(wired_limit_); return limit; }; @@ -159,7 +159,7 @@ Buffer MetalAllocator::malloc(size_t size) { lk.lock(); num_resources_++; if (!buf->heap()) { - residency_set_.insert(buf); + residency_sets_.insert(buf); } } @@ -192,7 +192,7 @@ void MetalAllocator::free(Buffer buffer) { } else { num_resources_--; if (!buf->heap()) { - residency_set_.erase(buf); + residency_sets_.erase(buf); } lk.unlock(); auto pool = metal::new_scoped_memory_pool(); @@ -210,7 +210,7 @@ Buffer MetalAllocator::make_buffer(void* ptr, size_t size) { return Buffer{nullptr}; } std::unique_lock lk(mutex_); - residency_set_.insert(buf); + residency_sets_.insert(buf); active_memory_ += buf->length(); peak_memory_ = std::max(peak_memory_, active_memory_); num_resources_++; @@ -225,7 +225,7 @@ void MetalAllocator::release(Buffer buffer) { std::unique_lock lk(mutex_); active_memory_ -= buf->length(); num_resources_--; - residency_set_.erase(buf); + residency_sets_.erase(buf); lk.unlock(); auto pool = metal::new_scoped_memory_pool(); buf->release(); diff --git a/mlx/backend/metal/allocator.h b/mlx/backend/metal/allocator.h index 4cbbfb0adc..885951364e 100644 --- a/mlx/backend/metal/allocator.h +++ b/mlx/backend/metal/allocator.h @@ -57,7 +57,7 @@ class MetalAllocator : public allocator::Allocator { friend MetalAllocator& allocator(); NS::SharedPtr heap_; - ResidencySet& residency_set_; + ResidencySets& residency_sets_; // Caching allocator BufferCache buffer_cache_; diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index cd601dd319..65df5c108c 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -309,17 +309,16 @@ MTL::Library* load_library( CommandEncoder::CommandEncoder( Device& d, int index, - ResidencySet& residency_set) - : device_(d) { + ResidencySets& residency_sets) + : device_(d), residency_sets_(residency_sets) { auto pool = new_scoped_memory_pool(); queue_ = NS::TransferPtr(device_.mtl_device()->newCommandQueue()); if (!queue_) { throw std::runtime_error( "[metal::CommandEncoder] Failed to make new command queue."); } - if (residency_set.mtl_residency_set()) { - queue_->addResidencySet(residency_set.mtl_residency_set()); - } + // Sets created later are attached in commit(). + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); debug_set_stream_queue_label(queue_.get(), index); buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); } @@ -519,6 +518,9 @@ bool CommandEncoder::needs_commit() const { } void CommandEncoder::commit(std::function completion) { + // Metal locks a command buffer's residency at commit time, so attach any + // sets created since the last commit first. + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); buffer_->addCompletedHandler( [&error_ = error_, wait_events = std::move(wait_events_), @@ -586,7 +588,7 @@ MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() { return encoder_.get(); } -Device::Device() : device_(load_device()), residency_set_(device_.get()) { +Device::Device() : device_(load_device()), residency_sets_(device_.get()) { auto pool = new_scoped_memory_pool(); default_library_ = NS::TransferPtr(load_default_library(device_.get())); arch_ = env::metal_gpu_arch(); diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 871e95ccff..3bb1e9e3b3 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -25,7 +25,7 @@ class EventImpl; class MLX_API CommandEncoder { public: - CommandEncoder(Device& d, int index, ResidencySet& residency_set); + CommandEncoder(Device& d, int index, ResidencySets& residency_sets); ~CommandEncoder(); CommandEncoder(const CommandEncoder&) = delete; @@ -114,6 +114,10 @@ class MLX_API CommandEncoder { int buffer_ops_{0}; size_t buffer_sizes_{0}; + // The residency set and how many of its sets this queue has attached. + ResidencySets& residency_sets_; + uint64_t sets_attached_{0}; + // The events hooked to current command buffer. std::vector> wait_events_; std::vector, uint64_t>> signal_events_; @@ -193,8 +197,8 @@ class MLX_API Device { const MTLFCList& func_consts = {}, const std::vector& linked_functions = {}); - ResidencySet& residency_set() { - return residency_set_; + ResidencySets& residency_sets() { + return residency_sets_; } private: @@ -230,7 +234,7 @@ class MLX_API Device { const std::vector& linked_functions = {}); NS::SharedPtr device_; - ResidencySet residency_set_; + ResidencySets residency_sets_; std::shared_mutex kernel_mtx_; std::shared_mutex library_mtx_; diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp index e6826253de..43a6fb9932 100644 --- a/mlx/backend/metal/eval.cpp +++ b/mlx/backend/metal/eval.cpp @@ -16,14 +16,14 @@ void new_stream(Stream s) { assert(s.device == Device::gpu); auto& encoders = metal::get_command_encoders(); auto& d = metal::device(s.device); - encoders.try_emplace(s.index, d, s.index, d.residency_set()); + encoders.try_emplace(s.index, d, s.index, d.residency_sets()); } void new_thread_unsafe_stream(Stream s) { assert(s.device == Device::gpu); auto& encoders = metal::get_global_command_encoders(); auto& d = metal::device(s.device); - encoders.try_emplace(s.index, d, s.index, d.residency_set()); + encoders.try_emplace(s.index, d, s.index, d.residency_sets()); } void eval(array& arr) { diff --git a/mlx/backend/metal/resident.cpp b/mlx/backend/metal/resident.cpp index 80de4f5792..2436993637 100644 --- a/mlx/backend/metal/resident.cpp +++ b/mlx/backend/metal/resident.cpp @@ -1,19 +1,40 @@ // Copyright © 2024 Apple Inc. -#include "mlx/backend/metal/resident.h" +#include +#include +#include +#include +#include + #include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/resident.h" +#include "mlx/utils.h" namespace mlx::core::metal { -ResidencySet::ResidencySet(MTL::Device* d) { +ResidencySets::ResidencySets(MTL::Device* d) { if (!d->supportsFamily(MTL::GPUFamilyMetal3)) { - return; - } else if (__builtin_available(macOS 15, iOS 18, *)) { - auto pool = new_scoped_memory_pool(); - auto desc = MTL::ResidencySetDescriptor::alloc()->init()->autorelease(); - NS::Error* error; - wired_set_ = NS::TransferPtr(d->newResidencySet(desc, &error)); - if (!wired_set_) { + return; // enabled_ stays false, so everything below is a no-op + } + if (__builtin_available(macOS 15, iOS 18, *)) { + device_ = d; + enabled_ = true; + int pct = env::residency_set_max_pct(); + if (pct <= 0 || pct >= 100) { + max_bytes_per_set_ = 0; // a single set holds everything + } else { + size_t ws = static_cast(d->recommendedMaxWorkingSetSize()); + // Floor the cap so a small working set or a small percentage cannot ask + // for an absurd number of sets. 64 MiB comfortably exceeds the + // allocator's heap. + max_bytes_per_set_ = std::max( + (ws / 100) * static_cast(pct), size_t(64) << 20); + } + std::lock_guard lk(mtx_); + // Set 0 always exists and is the fallback when a later set cannot be + // made, so failing to create it is fatal. + NS::Error* error = nullptr; + if (!add_set_locked(&error)) { std::ostringstream msg; msg << "[metal::Device] Unable to construct residency set.\n"; if (error) { @@ -21,75 +42,198 @@ ResidencySet::ResidencySet(MTL::Device* d) { } throw std::runtime_error(msg.str()); } - wired_set_->requestResidency(); } } -void ResidencySet::insert(MTL::Allocation* buf) { - if (!wired_set_) { - return; +ResidencySets::~ResidencySets() = default; + +bool ResidencySets::add_set_locked(NS::Error** error_out) { + NS::SharedPtr set; + if (__builtin_available(macOS 15, iOS 18, *)) { + auto pool = new_scoped_memory_pool(); + auto desc = MTL::ResidencySetDescriptor::alloc()->init()->autorelease(); + NS::Error* error = nullptr; + set = NS::TransferPtr(device_->newResidencySet(desc, &error)); + if (set) { + // A standing request, so allocations added to this set later are + // covered without requesting residency again on every insert. + set->requestResidency(); + } else if (error_out) { + *error_out = error; + } } - if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) { - wired_set_->addAllocation(buf); - wired_set_->commit(); - } else { - unwired_set_.insert(buf); + if (!set) { + return false; + } + sets_.push_back(Set{std::move(set), 0}); + num_sets_.store(sets_.size(), std::memory_order_release); + if (env::residency_debug()) { + fprintf( + stderr, + "[residency] created residency set %zu (max_bytes_per_set=%zu MB)\n", + sets_.size() - 1, + max_bytes_per_set_ >> 20); } + return true; +} + +uint32_t ResidencySets::choose_set_locked(size_t bytes) { + if (max_bytes_per_set_ == 0) { + return 0; + } + const uint32_t none = static_cast(sets_.size()); + uint32_t empty = none; + uint32_t emptiest = 0; + for (uint32_t i = 0; i < sets_.size(); ++i) { + const size_t size = sets_[i].size; + if (size + bytes <= max_bytes_per_set_) { + return i; // first fit + } + // Only reached by an allocation larger than the cap, which needs a set of + // its own. Reusing an emptied one keeps the set count from growing. + if (i != 0 && size == 0 && empty == none) { + empty = i; + } + if (size < sets_[emptiest].size) { + emptiest = i; + } + } + if (empty != none) { + return empty; + } + if (sets_.size() < kMaxSets && add_set_locked()) { + return static_cast(sets_.size() - 1); + } + // At the cap, or the driver would not make another set. Filling the emptiest + // set keeps them evenly sized, so each one still holds a bounded fraction of + // the wired bytes. + return emptiest; +} + +uint32_t ResidencySets::add_to_set_locked( + const MTL::Allocation* buf, + Placement& at) { + uint32_t idx = choose_set_locked(at.bytes); + auto& s = sets_[idx]; + s.set->addAllocation(buf); + s.size += at.bytes; + total_wired_ += at.bytes; + at.set_id = idx; + return idx; } -void ResidencySet::erase(MTL::Allocation* buf) { - if (!wired_set_) { +void ResidencySets::remove_from_set_locked( + const MTL::Allocation* buf, + Placement& at) { + auto& s = sets_[at.set_id]; + s.set->removeAllocation(buf); + // Subtract the size recorded at insert; allocatedSize() is never read back. + s.size -= at.bytes; + total_wired_ -= at.bytes; + at.set_id = kNoSet; +} + +void ResidencySets::insert(MTL::Allocation* buf) { + if (!enabled_) { return; } - if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) { - unwired_set_.erase(it); - } else { - wired_set_->removeAllocation(buf); - wired_set_->commit(); + const size_t bytes = buf->allocatedSize(); + std::lock_guard lk(mtx_); + + auto [it, inserted] = buf_to_set_.try_emplace(buf, Placement{kNoSet, bytes}); + if (!inserted) { + assert(false && "allocation is already tracked"); + return; + } + // Stay within the wired limit. The excess is tracked but left out of any + // set, and is added to one by resize() if the limit is raised later. + if (total_wired_ + bytes > capacity_) { + return; } + uint32_t idx = add_to_set_locked(buf, it->second); + sets_[idx].set->commit(); } -void ResidencySet::resize(size_t size) { - if (!wired_set_) { +void ResidencySets::erase(MTL::Allocation* buf) { + if (!enabled_) { return; } + std::lock_guard lk(mtx_); + auto it = buf_to_set_.find(buf); + if (it == buf_to_set_.end()) { + assert(false && "erasing an allocation that was never inserted"); + return; + } + if (it->second.set_id != kNoSet) { + const uint32_t idx = it->second.set_id; + remove_from_set_locked(buf, it->second); + sets_[idx].set->commit(); + } + buf_to_set_.erase(it); +} +void ResidencySets::resize(size_t size) { + if (!enabled_) { + return; + } + std::lock_guard lk(mtx_); if (capacity_ == size) { return; } capacity_ = size; - size_t current_size = wired_set_->allocatedSize(); - - if (current_size < size) { - auto pool = new_scoped_memory_pool(); - // Add unwired allocations to the set - for (auto it = unwired_set_.begin(); it != unwired_set_.end();) { - auto buf_size = (*it)->allocatedSize(); - if (current_size + buf_size > size) { - it++; - } else { - current_size += buf_size; - wired_set_->addAllocation(*it); - unwired_set_.erase(it++); + auto pool = new_scoped_memory_pool(); + std::vector touched(sets_.size(), false); + // The loops below only mutate map values, never insert or erase, so the + // iterators stay valid across the whole walk. + if (total_wired_ < capacity_) { + // The budget grew: add allocations that now fit. + for (auto& [buf, at] : buf_to_set_) { + if (at.set_id != kNoSet || total_wired_ + at.bytes > capacity_) { + continue; } + uint32_t idx = add_to_set_locked(buf, at); + if (idx >= touched.size()) { + touched.resize(sets_.size(), false); // a new set was made + } + touched[idx] = true; } - wired_set_->commit(); - } else if (current_size > size) { - auto pool = new_scoped_memory_pool(); - // Remove wired allocations until under capacity - auto allocations = wired_set_->allAllocations(); - auto num_allocations = wired_set_->allocationCount(); - for (int i = 0; i < num_allocations && current_size > size; ++i) { - auto buf = static_cast(allocations->object(i)); - wired_set_->removeAllocation(buf); - current_size -= buf->allocatedSize(); - unwired_set_.insert(buf); + } else { + // The budget shrank: remove allocations until we are back under it. + for (auto& [buf, at] : buf_to_set_) { + if (total_wired_ <= capacity_) { + break; + } + if (at.set_id == kNoSet) { + continue; + } + touched[at.set_id] = true; + remove_from_set_locked(buf, at); + } + } + for (size_t i = 0; i < touched.size(); ++i) { + if (touched[i]) { + sets_[i].set->commit(); } - wired_set_->commit(); } } -ResidencySet::~ResidencySet() = default; +void ResidencySets::attach_new_sets(MTL::CommandQueue* q, uint64_t& attached) { + // Lock-free when there is nothing new, which is the common case. It also + // covers the disabled case, where the count stays 0. + if (num_sets_.load(std::memory_order_acquire) == attached) { + return; + } + // Attach the new sets and record how far we got under a single lock hold, + // so a set created in between cannot be missed. + std::lock_guard lk(mtx_); + std::vector sets; + sets.reserve(sets_.size() - attached); + for (uint64_t id = attached; id < sets_.size(); ++id) { + sets.push_back(sets_[id].set.get()); + } + q->addResidencySets(sets.data(), sets.size()); + attached = sets_.size(); +} } // namespace mlx::core::metal diff --git a/mlx/backend/metal/resident.h b/mlx/backend/metal/resident.h index 50b1bd03d4..23960ad134 100644 --- a/mlx/backend/metal/resident.h +++ b/mlx/backend/metal/resident.h @@ -2,33 +2,122 @@ #pragma once -#include +#include +#include +#include +#include +#include #include namespace mlx::core::metal { -class ResidencySet { +// Keeps allocations GPU-resident, up to the wired limit set by +// `set_wired_limit` (0 by default, i.e. nothing is wired unless asked for). +// +// Within that budget the allocations are distributed over several size-capped +// MTL::ResidencySets rather than one large one. macOS makes residency decisions +// per residency set, so when a set loses residency under GPU memory pressure +// only the allocations in that set have to be made resident again. Capping the +// size of each set bounds the cost of one such event. Every set holds a +// standing requestResidency() and is attached to every command queue. +// +// MLX_RESIDENCY_SET_MAX_PCT (env::residency_set_max_pct) sets the per-set cap. +// The total wired budget is unaffected by it: that is still `set_wired_limit`. +class ResidencySets { public: - ResidencySet(MTL::Device* d); - ~ResidencySet(); + ResidencySets(MTL::Device* d); + ~ResidencySets(); - ResidencySet(const ResidencySet&) = delete; - ResidencySet& operator=(const ResidencySet&) = delete; - - const MTL::ResidencySet* mtl_residency_set() { - return wired_set_.get(); - } + ResidencySets(const ResidencySets&) = delete; + ResidencySets& operator=(const ResidencySets&) = delete; + // Called with the allocator's mutex held. void insert(MTL::Allocation* buf); void erase(MTL::Allocation* buf); void resize(size_t size); + bool enabled() const { + return enabled_; + } + + // Attaches the sets this queue has not seen yet. Called from the encoder + // thread before every command-buffer commit, because Metal locks a command + // buffer's residency at commit time. + void attach_new_sets(MTL::CommandQueue* q, uint64_t& attached); + + // Total bytes currently wired across all sets. + size_t wired_size() const { + std::lock_guard lk(mtx_); + return total_wired_; + } + size_t num_sets() const { + return num_sets_.load(std::memory_order_acquire); + } + + // Testing only: sets the per-set cap for subsequent inserts, bypassing the + // size floor so tests can reach the multi-set paths cheaply. 0 selects the + // single-set layout. + void set_max_bytes_per_set(size_t bytes) { + std::lock_guard lk(mtx_); + max_bytes_per_set_ = bytes; + } + size_t max_bytes_per_set() const { + std::lock_guard lk(mtx_); + return max_bytes_per_set_; + } + private: - NS::SharedPtr wired_set_; - std::unordered_set unwired_set_; + // A set id of kNoSet means the allocation is tracked but is not in any + // set, because it did not fit in the wired limit. + static constexpr uint32_t kNoSet = UINT32_MAX; + // A command queue accepts a limited number of residency sets and every set + // is attached to every queue, so the number of sets is capped. Once the cap + // is reached allocations go to the emptiest set, which then grows past + // max_bytes_per_set_. The limit is from the Metal feature set tables: + // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf + static constexpr uint32_t kMaxSets = 32; + + struct Set { + NS::SharedPtr set; + size_t size{0}; + }; + + // The set an allocation lives in and the size it was inserted with. Sizes + // are recorded rather than read back from the allocation at erase, so the + // running totals cannot drift. + struct Placement { + uint32_t set_id; + size_t bytes; + }; + + // The following all require mtx_. add_set_locked reports whether the driver + // gave us a set; the new set is the last one. add_to_set_locked returns the + // set it used. Neither it nor remove_from_set_locked commits, so a bulk + // resize costs one commit per set instead of one per allocation. + bool add_set_locked(NS::Error** error = nullptr); + uint32_t choose_set_locked(size_t bytes); + uint32_t add_to_set_locked(const MTL::Allocation* buf, Placement& at); + void remove_from_set_locked(const MTL::Allocation* buf, Placement& at); + + MTL::Device* device_{nullptr}; + bool enabled_{false}; + + // Per-set cap in bytes (0 for a single set) and the total wired budget. + size_t max_bytes_per_set_{0}; size_t capacity_{0}; + size_t total_wired_{0}; + + // Sets indexed by id. Append-only, so ids stay valid and dense. + std::vector sets_; + // Every tracked allocation, wired or not. + std::unordered_map buf_to_set_; + + // sets_.size(), published so a queue can check for new sets without + // taking the lock. + std::atomic num_sets_{0}; + mutable std::mutex mtx_; }; } // namespace mlx::core::metal diff --git a/mlx/utils.h b/mlx/utils.h index b5b516d89c..0c62ba30d9 100644 --- a/mlx/utils.h +++ b/mlx/utils.h @@ -191,6 +191,21 @@ inline int max_mb_per_buffer(int default_value) { return max_mb_per_buffer_; } +// Per-set residency-set size, as a percentage of the device's recommended +// max working-set size. Controls only how wired memory is distributed across +// residency sets, never how much is wired; see metal::ResidencySets. A value +// <= 0 or >= 100 puts everything in a single set. +inline int residency_set_max_pct() { + static int residency_set_max_pct_ = get_var("MLX_RESIDENCY_SET_MAX_PCT", 5); + return residency_set_max_pct_; +} + +// Log each residency set as it is created. +inline bool residency_debug() { + static bool residency_debug_ = get_var("MLX_RESIDENCY_DEBUG", 0); + return residency_debug_; +} + inline bool metal_fast_synch() { static bool metal_fast_synch = get_var("MLX_METAL_FAST_SYNCH", 0); return metal_fast_synch; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2a4a41c6b6..859215acb1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -10,6 +10,10 @@ if(MLX_BUILD_METAL OR MLX_BUILD_CUDA) set(METAL_TEST_SOURCES gpu_tests.cpp) endif() +if(MLX_BUILD_METAL) + list(APPEND METAL_TEST_SOURCES residency_tests.cpp) +endif() + include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake) target_sources( diff --git a/tests/residency_tests.cpp b/tests/residency_tests.cpp new file mode 100644 index 0000000000..60cd7c90e4 --- /dev/null +++ b/tests/residency_tests.cpp @@ -0,0 +1,296 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "doctest/doctest.h" + +#include "mlx/allocator.h" +#include "mlx/backend/metal/device.h" +#include "mlx/memory.h" +#include "mlx/mlx.h" + +using namespace mlx::core; + +namespace { + +constexpr size_t MB = 1 << 20; + +metal::ResidencySets& residency() { + return metal::device(Device::gpu).residency_sets(); +} + +// Restores the wired limit, the cache limit and the set cap on scope exit so +// a failing assertion can't leak global state into the rest of the suite. A +// cache limit of 0 makes free() release (and unwire) immediately instead of +// recycling, so residency accounting is observable synchronously. +struct LimitGuard { + size_t wired; + size_t cache; + size_t max_per_set; + LimitGuard(size_t new_wired, size_t new_cache) + : cache(set_cache_limit(new_cache)), + max_per_set(residency().max_bytes_per_set()) { + clear_cache(); + synchronize(); // retire command buffers still holding temporaries + wired = set_wired_limit(new_wired); + } + ~LimitGuard() { + residency().set_max_bytes_per_set(max_per_set); + set_cache_limit(cache); + set_wired_limit(wired); + } +}; + +// Sum of a small graph, used to drive a command-buffer commit (and with it the +// residency attach path) while sets exist. Synchronizes so the command +// buffer's temporaries are released before the caller checks accounting. +void check_gpu_work() { + auto x = sum(ones({256, 256}, float32)); + eval(x); + CHECK_EQ(x.item(), 65536.0f); + synchronize(); +} + +} // namespace + +TEST_CASE("test residency set wires nothing when the wired limit is zero") { + if (!residency().enabled()) { + INFO("skipped: needs a Metal 3 GPU on macOS >= 15"); + return; + } + LimitGuard guard(0, 0); + + // The default wired limit is 0, and at 0 nothing may be wired -- not even + // the allocator's heap. + CHECK_EQ(residency().wired_size(), 0); + + std::vector bufs; + for (int i = 0; i < 4; ++i) { + bufs.push_back(allocator::malloc(4 * MB)); + CHECK_EQ(residency().wired_size(), 0); + } + for (auto buf : bufs) { + allocator::free(buf); + } + CHECK_EQ(residency().wired_size(), 0); +} + +TEST_CASE("test residency set never wires more than the wired limit") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + // Sample the baseline with a budget in place, so the allocator's heap is + // already wired and only our own allocations move the total. + set_wired_limit(256 * MB); + const size_t baseline = residency().wired_size(); + const size_t limit = baseline + 8 * MB; + set_wired_limit(limit); + + std::vector bufs; + for (int i = 0; i < 8; ++i) { // asks for 32 MB against an 8 MB budget + bufs.push_back(allocator::malloc(4 * MB)); + CHECK_LE(residency().wired_size(), limit); + } + // Over budget overall, but the budget itself is still used. + CHECK_GE(residency().wired_size(), baseline + 4 * MB); + + for (auto buf : bufs) { + allocator::free(buf); + } + // Everything of ours is unwired again, exactly. + CHECK_EQ(residency().wired_size(), baseline); +} + +TEST_CASE("test raising the wired limit wires already-allocated buffers") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + + auto buf = allocator::malloc(8 * MB); + CHECK_EQ(residency().wired_size(), 0); + + set_wired_limit(64 * MB); + CHECK_GE(residency().wired_size(), 8 * MB); + + allocator::free(buf); +} + +TEST_CASE("test lowering the wired limit unwires buffers") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(64 * MB, 0); + const size_t baseline_with_budget = residency().wired_size(); + + auto buf = allocator::malloc(8 * MB); + CHECK_GE(residency().wired_size(), baseline_with_budget + 8 * MB); + + set_wired_limit(0); + CHECK_EQ(residency().wired_size(), 0); + + // ...and comes back when the budget is restored. + set_wired_limit(64 * MB); + CHECK_GE(residency().wired_size(), baseline_with_budget + 8 * MB); + + allocator::free(buf); +} + +TEST_CASE("test residency accounting is exact across free/realloc cycles") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(64 * MB, 0); + + // Baseline is whatever is wired with nothing of ours allocated (the heap). + const size_t baseline = residency().wired_size(); + + for (int i = 0; i < 32; ++i) { + auto buf = allocator::malloc(8 * MB); + CHECK_GE(residency().wired_size(), baseline + 8 * MB); + allocator::free(buf); + // Exact: erase subtracts the bytes recorded at insert, so repeated + // cycles must not drift the running total. + CHECK_EQ(residency().wired_size(), baseline); + } +} + +TEST_CASE("test wired allocations are spread across size-capped sets") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + const size_t max_per_set = 8 * MB; + residency().set_max_bytes_per_set(max_per_set); + set_wired_limit(128 * MB); + const size_t baseline = residency().wired_size(); + const size_t baseline_sets = residency().num_sets(); + + // Each buffer is over half the cap, so no two share a set. + std::vector bufs; + for (int i = 0; i < 8; ++i) { + bufs.push_back(allocator::malloc(5 * MB)); + } + CHECK_EQ(residency().wired_size(), baseline + 8 * 5 * MB); + CHECK_GT(residency().num_sets(), baseline_sets); + // No set may exceed the cap while there is room to make more. + CHECK_GE(residency().num_sets(), 8); + + // Sets created after the queue exists must still be attached before the + // commit that could reference them. + check_gpu_work(); + + const size_t sets_before = residency().num_sets(); + for (auto buf : bufs) { + allocator::free(buf); + } + CHECK_EQ(residency().wired_size(), baseline); + + // Emptied sets are reused rather than leaked, so cycling the same + // allocations must not keep growing the set count. + for (int cycle = 0; cycle < 4; ++cycle) { + std::vector again; + for (int i = 0; i < 8; ++i) { + again.push_back(allocator::malloc(5 * MB)); + } + CHECK_EQ(residency().num_sets(), sets_before); + for (auto buf : again) { + allocator::free(buf); + } + } + CHECK_EQ(residency().wired_size(), baseline); +} + +TEST_CASE("test the set count stays within the command queue limit") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + // A tiny cap against a large budget asks for far more sets than a command + // queue accepts; the count must saturate instead. + residency().set_max_bytes_per_set(2 * MB); + set_wired_limit(256 * MB); + const size_t baseline = residency().wired_size(); + + std::vector bufs; + for (int i = 0; i < 64; ++i) { + bufs.push_back(allocator::malloc(2 * MB)); + } + CHECK_LE(residency().num_sets(), 32); + CHECK_EQ(residency().wired_size(), baseline + 64 * 2 * MB); + + // Attaching a saturated set set to a queue must not fail. + check_gpu_work(); + + for (auto buf : bufs) { + allocator::free(buf); + } + CHECK_EQ(residency().wired_size(), baseline); +} + +TEST_CASE("test an allocation larger than the set cap is still wired") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + residency().set_max_bytes_per_set(4 * MB); + set_wired_limit(128 * MB); + const size_t baseline = residency().wired_size(); + + auto big = allocator::malloc(32 * MB); + CHECK_EQ(residency().wired_size(), baseline + 32 * MB); + check_gpu_work(); + + allocator::free(big); + CHECK_EQ(residency().wired_size(), baseline); +} + +TEST_CASE("test a set cap of zero keeps everything in one set") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + residency().set_max_bytes_per_set(0); + set_wired_limit(128 * MB); + const size_t baseline = residency().wired_size(); + const size_t baseline_sets = residency().num_sets(); + + std::vector bufs; + for (int i = 0; i < 16; ++i) { + bufs.push_back(allocator::malloc(4 * MB)); + } + CHECK_EQ(residency().num_sets(), baseline_sets); + CHECK_EQ(residency().wired_size(), baseline + 16 * 4 * MB); + + for (auto buf : bufs) { + allocator::free(buf); + } + CHECK_EQ(residency().wired_size(), baseline); +} + +TEST_CASE("test the wired limit is used up to its boundary") { + if (!residency().enabled()) { + return; + } + LimitGuard guard(0, 0); + // Sample the baseline with a budget in place, then leave room for exactly one + // more 4 MB allocation. + set_wired_limit(256 * MB); + const size_t baseline = residency().wired_size(); + set_wired_limit(baseline + 4 * MB); + + auto first = allocator::malloc(4 * MB); + CHECK_EQ(residency().wired_size(), baseline + 4 * MB); + + // The budget is exactly full: the next allocation is tracked but not wired. + auto second = allocator::malloc(4 * MB); + CHECK_EQ(residency().wired_size(), baseline + 4 * MB); + + // Freeing the wired one does not promote the pending one. + allocator::free(first); + CHECK_EQ(residency().wired_size(), baseline); + + allocator::free(second); + CHECK_EQ(residency().wired_size(), baseline); +} From a8e24f2029e4b0a976288cc917f3990e47184244 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 13 Aug 2026 15:12:08 +0200 Subject: [PATCH 198/222] [CUDA][Improvement] RMSNorm backward (#3881) --- mlx/backend/cuda/rms_norm.cu | 277 ++++++++++++++++++++++++++++++++++- 1 file changed, 275 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 97f3e94a5e..e4c0702a9f 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -3,6 +3,7 @@ #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/cuda/reduce/reduce.cuh" +#include "mlx/backend/cuda/steel/utils.cuh" #include "mlx/backend/gpu/copy.h" #include "mlx/dtype_utils.h" #include "mlx/fast_primitives.h" @@ -180,7 +181,7 @@ template < int BLOCK_DIM, int REDUCE_DIM, int N_READS = 4> -__global__ void rms_norm_vjp_small( +__global__ void rms_norm_vjp_small_fallback( const T* x, const T* w, const T* g, @@ -242,6 +243,154 @@ __global__ void rms_norm_vjp_small( } } +template +__global__ void rms_norm_vjp_small( + const T* x, + const T* w, + const T* g, + T* gx, + float* gw, // accumulate in float always + float eps, + int32_t axis_size, + int32_t n_rows, + int64_t w_stride) { + // persistent kernel, numblocks = number of sms * 2; + // each block is responsible for a row. 128*2 blocks = 256 stride. + // we pipeline loads and computation in shared memory: + // loading the row while doing dw, dx computation for another row + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceF2 = BlockBroadcastReduce; + __shared__ typename BlockReduceF2::TempStorage temp; + + // double buffering + constexpr int STAGES = 2; + const int num_blocks = static_cast(grid.num_blocks()); + const int num_tiles = cuda::ceil_div(n_rows, num_blocks); + auto tid = block.thread_index().x; + auto bid = grid.block_rank(); + constexpr int buffer_size = N_READS * BLOCK_SIZE * N_CHUNKS; // axis_size + + // shared memory is dymanic because we need > 48 kb + extern __shared__ char smem_raw[]; + T* smem_x = reinterpret_cast(smem_raw); + T* smem_dy = smem_x + STAGES * buffer_size; + constexpr int BYTES_PER_READ = N_READS * static_cast(sizeof(T)); + constexpr int STAGE_BYTES = buffer_size * static_cast(sizeof(T)); + uint32_t smem_x_addr = __cvta_generic_to_shared(smem_x); + uint32_t smem_dy_addr = __cvta_generic_to_shared(smem_dy); + + AlignedVector wn[N_CHUNKS]; + AlignedVector xn[N_CHUNKS]; + AlignedVector gn[N_CHUNKS]; + // gw_block is summed over every row this block visits + AlignedVector gw_block[N_CHUNKS]; + // zero init + for (int j = 0; j < N_CHUNKS; j++) { + for (int k = 0; k < N_READS; k++) { + gw_block[j][k] = 0.f; + } + } + // shift global pointer for each block + x += axis_size * bid; + g += axis_size * bid; + gx += axis_size * bid; + gw += axis_size * bid; + + // prefetch first row for each block + // and load weights to registers in the same loop + for (int j = 0; j < N_CHUNKS; j++) { + int offset = (j * BLOCK_SIZE + tid) * N_READS; + if (bid < n_rows) { + cp_async( + smem_x_addr + BYTES_PER_READ * (j * BLOCK_SIZE + tid), &x[offset]); + cp_async( + smem_dy_addr + BYTES_PER_READ * (j * BLOCK_SIZE + tid), &g[offset]); + } + wn[j] = load_vector( + w, j * BLOCK_SIZE + tid, axis_size, w_stride, T(0)); + } + // commit all N_CHUNKS 128 byte loads for the first row + cp_async_commit(); + // pipelineing + for (int tile = 0; tile < num_tiles; tile++) { + x += axis_size * num_blocks; + g += axis_size * num_blocks; + + int next = tile + 1; + int index = next % STAGES; + int64_t next_row = + static_cast(bid) + static_cast(next) * num_blocks; + + if (next < num_tiles && next_row < n_rows) { + for (int j = 0; j < N_CHUNKS; j++) { + int offset = (j * BLOCK_SIZE + tid) * N_READS; + cp_async( + smem_x_addr + index * STAGE_BYTES + + BYTES_PER_READ * (j * BLOCK_SIZE + tid), + &x[offset]); + cp_async( + smem_dy_addr + index * STAGE_BYTES + + BYTES_PER_READ * (j * BLOCK_SIZE + tid), + &g[offset]); + } + } + cp_async_commit(); // always commit, empty at the tail + cp_async_wait<1>(); // always wait for 1, the tail is empty + + int64_t cur_row = + static_cast(bid) + static_cast(tile) * num_blocks; + if (cur_row < n_rows) { + // load x and g from shared to registers + // compute the reduction per row + float2 factors = {}; + for (int j = 0; j < N_CHUNKS; j++) { + xn[j] = unsafe_load_vector( + smem_x + (tile % STAGES) * buffer_size, j * BLOCK_SIZE + tid); + gn[j] = unsafe_load_vector( + smem_dy + (tile % STAGES) * buffer_size, j * BLOCK_SIZE + tid); + for (int k = 0; k < N_READS; k++) { + float t = static_cast(xn[j][k]); + float wi = wn[j][k]; + float gi = gn[j][k]; + float wg = wi * gi; + factors = plus_f2(factors, {wg * t, t * t}); + } + } + factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {}); + float meangwx = factors.x / axis_size; + float normalizer = rsqrt(factors.y / axis_size + eps); + float normalizer3 = normalizer * normalizer * normalizer; + + // we store dx after processing each row, accumulate dw + T* gx_row = gx + static_cast(tile) * num_blocks * axis_size; + for (int j = 0; j < N_CHUNKS; j++) { + int offset = j * BLOCK_SIZE + tid; + for (int k = 0; k < N_READS; k++) { + float xi = static_cast(xn[j][k]); + float wi = wn[j][k]; + float gi = gn[j][k]; + if constexpr (HAS_W) { + gw_block[j][k] += gi * xi * normalizer; + } + xn[j][k] = + static_cast(normalizer * wi * gi - xi * meangwx * normalizer3); + } + store_vector(gx_row, offset, xn[j], axis_size); + } + } + } + // store this block's fp32 partial + if constexpr (HAS_W) { + for (int j = 0; j < N_CHUNKS; j++) { + int offset = j * BLOCK_SIZE + tid; + store_vector(gw, offset, gw_block[j], axis_size); + } + } +} + template __global__ void rms_norm_vjp( const T* x, @@ -481,6 +630,45 @@ void RMSNorm::eval_gpu( }); } +template +inline bool use_rmsnorm_vjp_fast( + const array& x, + const array& w, + const array& g, + const array& gx, + bool has_w, + int32_t axis_size, + int64_t w_stride) { + if (!has_w || w_stride != 1) { + return false; + } + if (axis_size != N_READS * BLOCK_SIZE * N_CHUNKS) { + return false; + } + auto aligned = [](const array& a) { + return reinterpret_cast(gpu_ptr(a)) % 16 == 0; + }; + return aligned(x) && aligned(w) && aligned(g) && aligned(gx); +} + +template +inline int rmsnorm_vjp_num_blocks( + Kernel kernel, + const Stream& s, + int block_size, + size_t smem_bytes, + int32_t n_rows) { + int dev = cu::device(s.device).cuda_device(); + int sm_count = 0; + CHECK_CUDA_ERROR( + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev)); + int blocks_per_sm = 1; + CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_sm, kernel, block_size, smem_bytes)); + int64_t want = static_cast(sm_count) * std::max(blocks_per_sm, 1); + return std::max(1, static_cast(std::min(want, n_rows))); +} + void RMSNormVJP::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -532,6 +720,91 @@ void RMSNormVJP::eval_gpu( int32_t n_rows = x.data_size() / axis_size; int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; + bool handled = false; + + if (has_w && cu::device(s.device).compute_capability_major() >= 8) { + dispatch_float_types(gx.dtype(), "rms_norm_vjp", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_num_chunks( + axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + constexpr int N_CHUNKS = n_chunks(); + if (!use_rmsnorm_vjp_fast( + x, w, g, gx, has_w, axis_size, w_stride)) { + return; + } + constexpr int STAGES = 2; + size_t smem_bytes = + size_t(2) * STAGES * axis_size * sizeof(DataType); + int dev = cu::device(s.device).cuda_device(); + int smem_max = 0; + CHECK_CUDA_ERROR(cudaDeviceGetAttribute( + &smem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); + // we should not be here often, is never true on Ampere and later + if (smem_bytes > static_cast(smem_max)) { + return; + } + auto kernel = cu::rms_norm_vjp_small< + DataType, + /*HAS_W=*/true, + BLOCK_SIZE, + N_CHUNKS, + N_READS>; + if (smem_bytes > 48000) { + CHECK_CUDA_ERROR(cudaFuncSetAttribute( + reinterpret_cast(kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes)); + } + int num_blocks = rmsnorm_vjp_num_blocks( + kernel, s, BLOCK_SIZE, smem_bytes, n_rows); + array gw_temp({num_blocks, axis_size}, float32, nullptr, {}); + gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder)); + encoder.add_temporary(gw_temp); + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(g); + encoder.set_output_array(gx); + encoder.set_output_array(gw_temp); + encoder.add_kernel_node_ex( + kernel, + dim3{static_cast(num_blocks)}, + dim3{static_cast(BLOCK_SIZE)}, + {}, // no cluster + static_cast(smem_bytes), + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(g), + gpu_ptr(gx), + gpu_ptr(gw_temp), + eps_, + axis_size, + n_rows, + w_stride); + + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, + {num_blocks}, + {axis_size}); + if (gw.dtype() == float32) { + col_reduce( + encoder, gw_temp, gw, Reduce::ReduceType::Sum, {0}, plan); + } else { + array gw_f32({axis_size}, float32, nullptr, {}); + col_reduce( + encoder, gw_temp, gw_f32, Reduce::ReduceType::Sum, {0}, plan); + encoder.add_temporary(gw_f32); + copy_gpu(gw_f32, gw, CopyType::General, s); + } + handled = true; + }); + }); + } + if (handled) { + return; + } + // Allocate a temporary to store the gradients for w and allocate the output // gradient accumulators. array gw_temp = @@ -560,7 +833,7 @@ void RMSNormVJP::eval_gpu( [&](auto group_dim, auto n_groups, auto groups_per_block) { constexpr int block_dim = group_dim() * n_groups(); static_assert(block_dim <= 32 || groups_per_block() == 1); - auto kernel = cu::rms_norm_vjp_small< + auto kernel = cu::rms_norm_vjp_small_fallback< DataType, has_w_constant.value, block_dim, From 77a0c1e8ae16eb781a79e2b7134ba84b2ee7b597 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eren=20Menge=C5=9F?= <76516591+erenmenges@users.noreply.github.com> Date: Thu, 13 Aug 2026 18:40:56 -0400 Subject: [PATCH 199/222] Fix MultiOptimizer on models containing empty modules (#4215) Co-authored-by: Cheng --- python/mlx/utils.py | 17 ++++++++++++++--- python/tests/test_optimizers.py | 21 ++++++++++++++++++++ python/tests/test_tree.py | 34 +++++++++++++++++++++++++++++++++ 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/python/mlx/utils.py b/python/mlx/utils.py index 66f132b7ef..f8595ef9a9 100644 --- a/python/mlx/utils.py +++ b/python/mlx/utils.py @@ -294,7 +294,8 @@ def tree_reduce(fn, tree, initializer=None, is_leaf=None): def tree_merge(tree_a, tree_b, merge_fn=None): """Merge two Python trees in one containing the values of both. It can be - thought of as a deep dict.update method. + thought of as a deep dict.update method. Empty containers are treated as + empty subtrees. Args: tree_a (Any): The first Python tree. @@ -305,9 +306,19 @@ def tree_merge(tree_a, tree_b, merge_fn=None): The Python tree containing the values of both ``tree_a`` and ``tree_b``. """ - if isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0: + empty_a = isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0 + empty_b = isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0 + + if empty_a and empty_b: + if type(tree_a) is not type(tree_b): + raise ValueError( + f"Cannot merge {type(tree_a).__name__} with {type(tree_b).__name__}" + ) + return type(tree_a)() + + if empty_a: tree_a = None - if isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0: + if empty_b: tree_b = None if tree_a is None and tree_b is not None: return tree_b diff --git a/python/tests/test_optimizers.py b/python/tests/test_optimizers.py index fec3d70a21..3d51fedf32 100644 --- a/python/tests/test_optimizers.py +++ b/python/tests/test_optimizers.py @@ -613,6 +613,27 @@ def __init__(self): self.assertFalse(any("bias" in k for k, v in adam_states)) self.assertFalse(any("weight" in k for k, v in sgd_states)) + def test_multi_optimizer_with_parameterless_layers(self): + mx.random.seed(0) + # test a sequential that has a no parameter module like ReLU + model = nn.Sequential(nn.Linear(4, 4), nn.ReLU(), nn.Linear(4, 4)) + mx.eval(model.parameters()) + + optimizer = opt.MultiOptimizer( + [opt.Muon(learning_rate=0.01), opt.AdamW(learning_rate=0.01)], + [lambda _, w: w.ndim >= 2], + ) + + loss_and_grad = nn.value_and_grad(model, lambda m, x: m(x).sum()) + _, grads = loss_and_grad(model, mx.ones((1, 4))) + optimizer.update(model, grads) + + w, b = model.layers[0].weight, model.layers[0].bias + optimizer.update(model, grads) + mx.eval(model.parameters()) + self.assertFalse(mx.array_equal(w, model.layers[0].weight)) + self.assertFalse(mx.array_equal(b, model.layers[0].bias)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/python/tests/test_tree.py b/python/tests/test_tree.py index c6f31981b1..96184bd233 100644 --- a/python/tests/test_tree.py +++ b/python/tests/test_tree.py @@ -46,6 +46,40 @@ def test_merge(self): self.assertEqual(k1, k2) self.assertTrue(mx.array_equal(v1, v2)) + def test_empty_subtree_merge(self): + # make sure mlx pytrees treat empty dict {} as an empty node + self.assertEqual([], mlx.utils.tree_flatten({"a": {}})) + + # empty dict merging + self.assertEqual({}, mlx.utils.tree_merge({}, {})) + self.assertEqual( + [{"a": 1, "b": 2}, {}], mlx.utils.tree_merge([{"a": 1}, {}], [{"b": 2}, {}]) + ) + self.assertEqual({"a": {}}, mlx.utils.tree_merge({"a": {}}, {"a": {}})) + self.assertEqual( + {"a": 1, "b": {}, "c": 2}, + mlx.utils.tree_merge( + {"a": 1, "b": {}, "c": {}}, {"a": {}, "b": {}, "c": 2} + ), + ) + + # empty list merging + self.assertEqual({"a": []}, mlx.utils.tree_merge({"a": []}, {"a": []})) + + # empty tuple merging + self.assertEqual({"a": ()}, mlx.utils.tree_merge({"a": ()}, {"a": ()})) + + # merging different empty structures + with self.assertRaises(ValueError): + mlx.utils.tree_merge({}, []) + + def merge_called(a, b): + raise AssertionError("merge_fn called on empty subtrees") + + self.assertEqual( + {"a": {}}, mlx.utils.tree_merge({"a": {}}, {"a": {}}, merge_called) + ) + def test_supported_trees(self): from typing import NamedTuple From 9b6575c35434efa588fff7e408ad55bb7617c62d Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:15:06 +0800 Subject: [PATCH 200/222] Return tuple in meshgrid (#4229) --- python/src/ops.cpp | 6 +++--- python/tests/test_ops.py | 5 +++++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index d8892a45d2..bdcf3dde4b 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3386,14 +3386,14 @@ void init_ops(nb::module_& m) { mx::StreamOrDevice s) { std::vector arrays = nb::cast>(arrays_); - return mx::meshgrid(arrays, sparse, indexing, s); + return nb::tuple(nb::cast(mx::meshgrid(arrays, sparse, indexing, s))); }, "arrays"_a, "sparse"_a = false, "indexing"_a = "xy", "stream"_a = nb::none(), nb::sig( - "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> array"), + "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Generate multidimensional coordinate grids from 1-D coordinate arrays @@ -3406,7 +3406,7 @@ void init_ops(nb::module_& m) { Defaults to ``'xy'``. Returns: - list(array): The output arrays. + tuple(array): The output arrays. )pbdoc"); m.def( "repeat", diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 1b46237ce5..83f95ea5fc 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2132,6 +2132,11 @@ def test_meshgrid(self): x = mx.array([1, 2, 3], dtype=mx.int32) y = np.array([1, 2, 3], dtype=np.int32) + # Test return type is a tuple + self.assertIsInstance(mx.meshgrid(x), tuple) + self.assertIsInstance(mx.meshgrid(x, x), tuple) + self.assertIsInstance(mx.meshgrid(x, x, x, sparse=True), tuple) + # Test single input a_mlx = mx.meshgrid(x) a_np = np.meshgrid(y) From 1d717bd3c562a45e6c0f6e195d413ea16d22898f Mon Sep 17 00:00:00 2001 From: AK <144495202+AKnassa@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:38:31 -0400 Subject: [PATCH 201/222] Add endpoint parameter to linspace (#4184) Co-authored-by: Cheng --- mlx/ops.cpp | 6 ++++- mlx/ops.h | 16 +++++++++-- python/src/ops.cpp | 8 +++++- python/tests/test_double.py | 8 +++++- python/tests/test_ops.py | 53 ++++++++++++++++++++++++++++++++++++- tests/ops_tests.cpp | 18 ++++++++++++- 6 files changed, 102 insertions(+), 7 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 9c8db3a26d..ef229fe87f 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -271,6 +271,7 @@ array linspace( double start, double stop, int num /* = 50 */, + bool endpoint /* = true */, Dtype dtype /* = float32 */, StreamOrDevice s /* = {} */) { if (num < 0) { @@ -282,8 +283,11 @@ array linspace( return astype(array({start}), dtype, s); } auto inner_type = dtype == float64 ? float64 : float32; + // Without the endpoint the samples are spaced so that `stop` would be the + // next one after the last, i.e. the step is (stop - start) / num. + auto denominator = endpoint ? num - 1 : num; array t = - divide(arange(0, num, inner_type, s), array(num - 1, inner_type), s); + divide(arange(0, num, inner_type, s), array(denominator, inner_type), s); array t_bar = subtract(array(1, inner_type), t, s); return astype( add(multiply(t_bar, array(start, inner_type), s), diff --git a/mlx/ops.h b/mlx/ops.h index 01e0a99286..f597753b1e 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -38,13 +38,25 @@ MLX_API array arange(int start, int stop, int step, StreamOrDevice s = {}); MLX_API array arange(int start, int stop, StreamOrDevice s = {}); MLX_API array arange(int stop, StreamOrDevice s = {}); -/** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */ +/** + * A 1D array of `num` evenly spaced numbers in the range `[start, stop]`, or + * in the half-open range `[start, stop)` when `endpoint` is false. + */ MLX_API array linspace( double start, double stop, - int num = 50, + int num, + bool endpoint, Dtype dtype = float32, StreamOrDevice s = {}); +inline array linspace( + double start, + double stop, + int num = 50, + Dtype dtype = float32, + StreamOrDevice s = {}) { + return linspace(start, stop, num, true, dtype, s); +} /** Convert an array to the given data type. */ MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {}); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index bdcf3dde4b..4dc5114bf7 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1644,22 +1644,25 @@ void init_ops(nb::module_& m) { [](Scalar start, Scalar stop, int num, + bool endpoint, std::optional dtype, mx::StreamOrDevice s) { return mx::linspace( scalar_to_double(start), scalar_to_double(stop), num, + endpoint, dtype.value_or(mx::float32), s); }, "start"_a, "stop"_a, "num"_a = 50, + "endpoint"_a = true, "dtype"_a.none() = mx::float32, "stream"_a = nb::none(), nb::sig( - "def linspace(start: scalar, stop: scalar, num: int | None = 50, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), + "def linspace(start: scalar, stop: scalar, num: int | None = 50, endpoint: bool = True, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. @@ -1667,6 +1670,9 @@ void init_ops(nb::module_& m) { start (scalar): Starting value. stop (scalar): Stopping value. num (int, optional): Number of samples, defaults to ``50``. + endpoint (bool, optional): If ``True``, ``stop`` is the last + sample. Otherwise it is not included and the samples are spaced + over the half-open interval ``[start, stop)``. Default: ``True``. dtype (Dtype, optional): Specifies the data type of the output, default to ``float32``. diff --git a/python/tests/test_double.py b/python/tests/test_double.py index 65603cd937..3186e7e573 100644 --- a/python/tests/test_double.py +++ b/python/tests/test_double.py @@ -336,9 +336,15 @@ def test_python_float_keeps_double_precision(self): def test_linspace(self): with mx.stream(mx.cpu): - vals = mx.linspace(0, math.pi, 2, mx.float64) + vals = mx.linspace(0, math.pi, 2, dtype=mx.float64) self.assertEqual(vals.tolist()[1], math.pi) + vals = mx.linspace(0, math.pi, 4, endpoint=False, dtype=mx.float64) + self.assertEqual(vals.dtype, mx.float64) + self.assertTrue( + np.allclose(vals.tolist(), np.linspace(0, math.pi, 4, endpoint=False)) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 83f95ea5fc..cdf0bc1b22 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2950,7 +2950,7 @@ def test_linspace(self): self.assertEqualArray(a, expected) # Test int64 dtype - b = mx.linspace(0, 10, 5, mx.int64) + b = mx.linspace(0, 10, 5, dtype=mx.int64) expected = mx.array(np.linspace(0, 10, 5, dtype=int)) self.assertEqualArray(b, expected) @@ -2977,6 +2977,57 @@ def test_linspace(self): self.assertEqual(d[0], a) self.assertEqual(d[-1], b) + def test_linspace_endpoint(self): + # endpoint=True is the default and matches the old behaviour + a = mx.linspace(0, 1, 5, endpoint=True) + self.assertEqualArray(a, mx.array(np.linspace(0, 1, 5, endpoint=True))) + self.assertEqualArray(a, mx.linspace(0, 1, 5)) + + # endpoint=False drops the stop value and uses a step of + # (stop - start) / num instead of (stop - start) / (num - 1) + for num in [0, 1, 2, 5, 50]: + b = mx.linspace(0, 10, num, endpoint=False) + expected = mx.array(np.linspace(0, 10, num, endpoint=False)) + self.assertEqualArray(b, expected) + + c = mx.linspace(-2.7, -0.7, 7, endpoint=False) + self.assertEqualArray(c, mx.array(np.linspace(-2.7, -0.7, 7, endpoint=False))) + + # endpoint is the fourth positional argument, before dtype, as in numpy + self.assertEqualArray( + mx.linspace(0, 10, 5, False), mx.array(np.linspace(0, 10, 5, False)) + ) + + # dtype still applies + d = mx.linspace(0, 10, 5, False, mx.int64) + self.assertEqual(d.dtype, mx.int64) + self.assertEqualArray( + d, mx.array(np.linspace(0, 10, 5, endpoint=False, dtype=int)) + ) + + # the start is kept and the stop is excluded + e = mx.linspace(3.0, 4.0, 4, endpoint=False).tolist() + self.assertEqual(e[0], 3.0) + self.assertNotIn(4.0, e) + + # decreasing ranges drop the stop value too + f = mx.linspace(10, 0, 5, endpoint=False) + self.assertEqualArray(f, mx.array(np.linspace(10, 0, 5, endpoint=False))) + + # start == stop keeps every sample at that value + g = mx.linspace(5, 5, 4, endpoint=False) + self.assertEqualArray(g, mx.array(np.linspace(5, 5, 4, endpoint=False))) + + # integer dtype truncates fractional steps, as in numpy + h = mx.linspace(0, 10, 3, endpoint=False, dtype=mx.int32) + self.assertEqualArray( + h, mx.array(np.linspace(0, 10, 3, endpoint=False, dtype=np.int32)) + ) + + # num must still be non-negative + with self.assertRaises(ValueError): + mx.linspace(0, 1, -1, endpoint=False) + def test_repeat(self): # Setup data for the tests data = mx.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]]) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index f7a2b8ab92..3da0a2950b 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3348,13 +3348,29 @@ TEST_CASE("test linspace") { auto expected = array({0.0f, 2.5f, 5.0f, 7.5f, 10.0f}, {5}); CHECK(array_equal(x, expected).item()); - x = linspace(0, 10, 5, int32); + x = linspace(0, 10, 5, true, int32); expected = array({0, 2, 5, 7, 10}, {5}); CHECK(array_equal(x, expected).item()); x = linspace(0, 1, 0); expected = array(std::initializer_list{}, {0}); CHECK(array_equal(x, expected).item()); + + x = linspace(0, 10, 5, false); + expected = array({0.0f, 2.0f, 4.0f, 6.0f, 8.0f}, {5}); + CHECK(array_equal(x, expected).item()); + + x = linspace(0, 10, 5, false, int32); + expected = array({0, 2, 4, 6, 8}, {5}); + CHECK(array_equal(x, expected).item()); + + x = linspace(1, 10, 1, false); + expected = array({1.0f}, {1}); + CHECK(array_equal(x, expected).item()); + + x = linspace(0, 1, 0, false); + expected = array(std::initializer_list{}, {0}); + CHECK(array_equal(x, expected).item()); } TEST_CASE("test quantize dequantize") { From 306bdcd18dce9b6734ea76b2b842abf3aa4af1f6 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Thu, 13 Aug 2026 22:39:17 -0700 Subject: [PATCH 202/222] Fix vmap of partition/argpartition dropping the kth argument (#4116) --- mlx/primitives.cpp | 4 +- python/tests/test_vmap.py | 79 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 3bafd19407..3c3d4fc604 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -616,7 +616,7 @@ std::pair, std::vector> ArgPartition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{argpartition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{argpartition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } std::vector ArgPartition::vjp( @@ -3421,7 +3421,7 @@ std::pair, std::vector> Partition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{partition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{partition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } bool Partition::is_equivalent(const Primitive& other) const { diff --git a/python/tests/test_vmap.py b/python/tests/test_vmap.py index 99c30a2dc2..ec97ed6fdf 100644 --- a/python/tests/test_vmap.py +++ b/python/tests/test_vmap.py @@ -252,6 +252,85 @@ def test_vmap_argreduce(self): expected = mx.array([2, 1]) self.assertTrue(mx.array_equal(out, expected)) + def _unstack(self, x, axis): + return [s.squeeze(axis) for s in mx.split(x, x.shape[axis], axis=axis)] + + def test_vmap_partition(self): + # Distinct values so each lane has a single valid kth element + a = mx.random.permutation(2 * 3 * 4).reshape(2, 3, 4).astype(mx.float32) + + for in_axis in (0, 1, 2): + slices = self._unstack(a, in_axis) + # Axis of the batched output that the inner axis maps onto + out_axes_map = [d for d in range(a.ndim) if d != in_axis] + for axis in (0, 1, -1): + oaxis = out_axes_map[axis if axis >= 0 else axis + 2] + for kth in range(slices[0].shape[axis]): + expected = mx.stack( + [mx.partition(x, kth, axis=axis) for x in slices], + axis=in_axis, + ) + pivot = mx.take(expected, mx.array([kth]), axis=oaxis) + + out = mx.vmap( + lambda x: mx.partition(x, kth, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + self.assertEqual(out.shape, expected.shape) + # partition only pins the kth element; the two sides are + # an arbitrary permutation, so compare against the sorted + # input rather than element-wise. + self.assertTrue( + mx.array_equal(mx.sort(out, axis=oaxis), mx.sort(a, axis=oaxis)) + ) + self.assertTrue( + mx.array_equal(mx.take(out, mx.array([kth]), axis=oaxis), pivot) + ) + + idx = mx.vmap( + lambda x: mx.argpartition(x, kth, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + self.assertEqual(idx.shape, expected.shape) + gathered = mx.take_along_axis(a, idx, axis=oaxis) + self.assertTrue( + mx.array_equal( + mx.sort(gathered, axis=oaxis), mx.sort(a, axis=oaxis) + ) + ) + self.assertTrue( + mx.array_equal( + mx.take(gathered, mx.array([kth]), axis=oaxis), pivot + ) + ) + + def test_vmap_topk(self): + a = mx.random.permutation(2 * 3 * 4).reshape(2, 3, 4).astype(mx.float32) + + for in_axis in (0, 1, 2): + slices = self._unstack(a, in_axis) + out_axes_map = [d for d in range(a.ndim) if d != in_axis] + for axis in (0, 1, -1): + oaxis = out_axes_map[axis if axis >= 0 else axis + 2] + for k in range(1, slices[0].shape[axis] + 1): + out = mx.vmap( + lambda x: mx.topk(x, k, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + expected = mx.stack( + [mx.topk(x, k, axis=axis) for x in slices], axis=in_axis + ) + self.assertEqual(out.shape, expected.shape) + # topk does not promise an order within the k elements + self.assertTrue( + mx.array_equal( + mx.sort(out, axis=oaxis), mx.sort(expected, axis=oaxis) + ) + ) + def test_vmap_mean(self): a = mx.arange(8).reshape(2, 4) out = mx.vmap(mx.mean)(a) From d9ad465542b4c33f46acd89c028053a10b652577 Mon Sep 17 00:00:00 2001 From: anchor Date: Fri, 14 Aug 2026 16:40:16 +0800 Subject: [PATCH 203/222] Fix nan_to_num replacing inf with 0 for float16 and bfloat16 (#4222) Co-authored-by: codeAnqiang-ma <273298913+codeAnqiang-ma@users.noreply.github.com> Co-authored-by: Cheng --- mlx/ops.cpp | 7 ++++--- python/tests/test_ops.py | 12 +++++++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index ef229fe87f..7a9afcfcc7 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -16,6 +16,7 @@ #include "mlx/primitives.h" #include "mlx/transforms.h" #include "mlx/transforms_impl.h" +#include "mlx/types/limits.h" #include "mlx/utils.h" namespace mlx::core { @@ -2111,11 +2112,11 @@ array nan_to_num( auto type_to_max = [](const auto& dtype) -> float { if (dtype == float32) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == bfloat16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == float16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else { std::ostringstream msg; msg << "[nan_to_num] Does not yet support given type: " << dtype << "."; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index cdf0bc1b22..2b6f9324f5 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2267,7 +2267,7 @@ def test_nan_to_num(self): self.assertTrue(np.allclose(out_mx, out_np)) for t in [mx.float32, mx.float16]: - a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]) + a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]).astype(t) out_mx = mx.nan_to_num(a) out_np = np.nan_to_num(a) self.assertTrue(np.allclose(out_mx, out_np)) @@ -2277,6 +2277,16 @@ def test_nan_to_num(self): out_mx = mx.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000) self.assertTrue(np.allclose(out_mx, out_np)) + # bfloat16 has no numpy analogue; infinities should clamp to the + # dtype's largest finite value, not 0 + a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]).astype( + mx.bfloat16 + ) + out_mx = mx.nan_to_num(a) + bf_max = mx.finfo(mx.bfloat16).max + expected = mx.array([bf_max, 6.9, 0.0, -bf_max]).astype(mx.bfloat16) + self.assertTrue(mx.array_equal(out_mx, expected)) + def test_pad_reflect_symmetric(self): # mx.pad reflect/symmetric must match numpy.pad exactly. Covers # in-bounds, multi-reflect (pad larger than the axis, exercising the From bd5c3a2b170bb95340482e35b2a49fb08aea4de3 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:42:06 -0700 Subject: [PATCH 204/222] Fix einsum not broadcasting batch dimensions in batched tensordot (#4125) Co-authored-by: Cheng --- mlx/einsum.cpp | 7 ++++++- python/tests/test_einsum.py | 35 ++++++++++++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/mlx/einsum.cpp b/mlx/einsum.cpp index b683733d0b..705a5fc2cb 100644 --- a/mlx/einsum.cpp +++ b/mlx/einsum.cpp @@ -356,7 +356,7 @@ array batch_tensordot( std::vector b_batch, std::vector b_concat, StreamOrDevice s) { - // Broadcast contracting dimensions + // Broadcast contracting and batch dimensions. { auto a_shape = a.shape(); auto b_shape = b.shape(); @@ -365,6 +365,11 @@ array batch_tensordot( a_shape[a_contract[i]] = d; b_shape[b_contract[i]] = d; } + for (int i = 0; i < a_batch.size(); ++i) { + auto d = std::max(a.shape(a_batch[i]), b.shape(b_batch[i])); + a_shape[a_batch[i]] = d; + b_shape[b_batch[i]] = d; + } a = broadcast_to(a, a_shape, s); b = broadcast_to(b, b_shape, s); } diff --git a/python/tests/test_einsum.py b/python/tests/test_einsum.py index a73ea38187..c87a6dc45f 100644 --- a/python/tests/test_einsum.py +++ b/python/tests/test_einsum.py @@ -188,7 +188,6 @@ def test_broadcasting(self): a = mx.full((5, 1), 1.0) b = mx.full((8, 2), 1.0) a_mx = mx.einsum("ab,bc->c", a, b) - return a_np = np.einsum("ab,bc->c", a, b) self.assertTrue(np.array_equal(a_mx, a_np)) @@ -358,6 +357,40 @@ def inputs_for_case(test_case): with self.assertRaises(ValueError): mx.einsum(test_case[1], *inputs) + def test_ellipses_broadcast(self): + # Size 1 batch dimensions covered by an ellipsis have to broadcast + # against the other operands, including when the smaller operand + # comes first. + shape_pairs = [ + ((1, 3, 4), (2, 4, 5)), + ((2, 3, 4), (1, 4, 5)), + ((1, 1, 3, 4), (5, 2, 4, 5)), + ((5, 1, 3, 4), (1, 2, 4, 5)), + ] + for sa, sb in shape_pairs: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...ij,...jk->...ik", a, b) + np_out = np.einsum("...ij,...jk->...ik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + for sa, sb in [((1, 4), (5, 4)), ((5, 4), (1, 4))]: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...i,...i->...", a, b) + np_out = np.einsum("...i,...i->...", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + # Same thing with explicit labels rather than an ellipsis + a = mx.random.uniform(shape=(1, 3, 4)) + b = mx.random.uniform(shape=(2, 4, 5)) + mx_out = mx.einsum("bij,bjk->bik", a, b) + np_out = np.einsum("bij,bjk->bik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From d2d7138cf371ce34cd832fe50389d5333f909e0a Mon Sep 17 00:00:00 2001 From: Pasha Khosravi Date: Mon, 23 Feb 2026 19:30:09 -0800 Subject: [PATCH 205/222] Add 1-bit affine quantization support --- benchmarks/python/comparative/bench_mlx.py | 31 ++- benchmarks/python/comparative/compare.py | 35 ++++ mlx/backend/cpu/quantized.cpp | 34 +++- mlx/backend/metal/kernels/quantized.h | 177 ++++++++++++++---- mlx/backend/metal/kernels/quantized.metal | 1 + mlx/backend/metal/kernels/quantized_nax.h | 162 +++++++++++++--- mlx/backend/metal/kernels/quantized_nax.metal | 1 + mlx/ops.cpp | 28 ++- python/src/ops.cpp | 16 +- python/tests/test_quantized.py | 96 +++++++++- 10 files changed, 483 insertions(+), 98 deletions(-) diff --git a/benchmarks/python/comparative/bench_mlx.py b/benchmarks/python/comparative/bench_mlx.py index 4e6ba04f8e..1f12064ebe 100644 --- a/benchmarks/python/comparative/bench_mlx.py +++ b/benchmarks/python/comparative/bench_mlx.py @@ -72,12 +72,17 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits): quant_matmul = { + "quant_matmul_32_1": partial(_quant_matmul, transpose=False, group_size=32, bits=1), "quant_matmul_32_2": partial(_quant_matmul, transpose=False, group_size=32, bits=2), "quant_matmul_32_4": partial(_quant_matmul, transpose=False, group_size=32, bits=4), "quant_matmul_32_8": partial(_quant_matmul, transpose=False, group_size=32, bits=8), + "quant_matmul_64_1": partial(_quant_matmul, transpose=False, group_size=64, bits=1), "quant_matmul_64_2": partial(_quant_matmul, transpose=False, group_size=64, bits=2), "quant_matmul_64_4": partial(_quant_matmul, transpose=False, group_size=64, bits=4), "quant_matmul_64_8": partial(_quant_matmul, transpose=False, group_size=64, bits=8), + "quant_matmul_128_1": partial( + _quant_matmul, transpose=False, group_size=128, bits=1 + ), "quant_matmul_128_2": partial( _quant_matmul, transpose=False, group_size=128, bits=2 ), @@ -87,6 +92,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits): "quant_matmul_128_8": partial( _quant_matmul, transpose=False, group_size=128, bits=8 ), + "quant_matmul_t_32_1": partial( + _quant_matmul, transpose=True, group_size=32, bits=1 + ), "quant_matmul_t_32_2": partial( _quant_matmul, transpose=True, group_size=32, bits=2 ), @@ -96,6 +104,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits): "quant_matmul_t_32_8": partial( _quant_matmul, transpose=True, group_size=32, bits=8 ), + "quant_matmul_t_64_1": partial( + _quant_matmul, transpose=True, group_size=64, bits=1 + ), "quant_matmul_t_64_2": partial( _quant_matmul, transpose=True, group_size=64, bits=2 ), @@ -105,6 +116,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits): "quant_matmul_t_64_8": partial( _quant_matmul, transpose=True, group_size=64, bits=8 ), + "quant_matmul_t_128_1": partial( + _quant_matmul, transpose=True, group_size=128, bits=1 + ), "quant_matmul_t_128_2": partial( _quant_matmul, transpose=True, group_size=128, bits=2 ), @@ -420,7 +434,22 @@ def selu(x): print(bench(matmul, *xs)) elif args.benchmark.startswith("quant_matmul"): - print(bench(quant_matmul[args.benchmark], *xs)) + # Parse group_size and bits from the benchmark name, e.g. + # "quant_matmul_128_4" or "quant_matmul_t_128_4" + fn = quant_matmul[args.benchmark] + gs = fn.keywords["group_size"] + bits = fn.keywords["bits"] + transpose = fn.keywords["transpose"] + + # xs[0] = activation x, xs[1] = original (float) weight matrix + # Quantize the weight internally so the caller only needs: + # --size MxK --size NxK (transpose=True) or --size MxK --size KxN + w_float = xs[1].astype(mx.float16) + w_q, scales, biases = mx.quantize(w_float, group_size=gs, bits=bits) + mx.eval(w_q, scales, biases) + x_input = xs[0].astype(mx.float16) + mx.eval(x_input) + print(bench(_quant_matmul, x_input, w_q, scales, biases, transpose, gs, bits)) elif args.benchmark == "linear": if args.fused: diff --git a/benchmarks/python/comparative/compare.py b/benchmarks/python/comparative/compare.py index aa81369fe2..6f3697bb23 100644 --- a/benchmarks/python/comparative/compare.py +++ b/benchmarks/python/comparative/compare.py @@ -29,6 +29,18 @@ def compare(args): print((t_torch - t_mlx) / t_torch, " ".join(args), sep="\t") +def compare_mlx_quant(args_base, bits_list): + """Compare quantized matmul across bit widths (MLX only, no PyTorch).""" + results = {} + for bits in bits_list: + bench_args = args_base.replace("{bits}", str(bits)).split() + results[bits] = run_or_raise(["python", BENCH_MLX] + bench_args) + baseline = max(results.values()) + for bits in bits_list: + speedup = (baseline - results[bits]) / baseline if baseline > 0 else 0 + print(f"{speedup:.4f}\t{args_base.replace('{bits}', str(bits))}") + + def compare_mlx_dtypes(args, dt1, dt2): t_mlx_dt1 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt1]) t_mlx_dt2 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt2]) @@ -283,3 +295,26 @@ def predicate(x): compare_filtered("topk --size 32768x128 --axis 1") compare_filtered("topk --size 128x128 --axis 0 --cpu") compare_filtered("topk --size 128x128 --axis 1 --cpu") + + # Quantized matmul ops (MLX only — compare across bit widths) + # qmv path (M=1, token generation, memory-bandwidth bound) + for gs in [64, 128]: + compare_mlx_quant( + f"quant_matmul_t_{gs}_{{bits}} --size 1x4096 --size 4096x4096", + [1, 2, 4, 8], + ) + compare_mlx_quant( + f"quant_matmul_t_{gs}_{{bits}} --size 1x4096 --size 11008x4096", + [1, 2, 4, 8], + ) + # qmm path (prompt processing, more compute bound) + for gs in [64, 128]: + for M in [32, 512]: + compare_mlx_quant( + f"quant_matmul_t_{gs}_{{bits}} --size {M}x4096 --size 4096x4096", + [1, 2, 4, 8], + ) + compare_mlx_quant( + f"quant_matmul_t_{gs}_{{bits}} --size {M}x4096 --size 11008x4096", + [1, 2, 4, 8], + ) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index 3469d99788..968d2a8aff 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -351,6 +351,10 @@ void _qmm_dispatch_typed( int bits, bool transposed_w) { switch (bits) { + case 1: + _qmm_dispatch_group( + result, x, w, scales, biases, M, N, K, group_size, transposed_w); + break; case 2: _qmm_dispatch_group( result, x, w, scales, biases, M, N, K, group_size, transposed_w); @@ -376,7 +380,8 @@ void _qmm_dispatch_typed( result, x, w, scales, biases, M, N, K, group_size, transposed_w); break; default: - throw std::invalid_argument("Quantization bits must be 2, 3, 4, 6 or 8."); + throw std::invalid_argument( + "Quantization bits must be 1, 2, 3, 4, 5, 6 or 8."); } } @@ -1172,15 +1177,24 @@ void quantize( w_min = std::min(w_min, (float)w[w_idx + j]); } bool mask = std::abs(w_min) > std::abs(w_max); - float scale = std::max((w_max - w_min) / n_bins, eps); - scale = mask ? scale : -scale; - - float edge = mask ? w_min : w_max; - float q0 = std::rint(edge / scale); - float bias = 0; - if (q0 != 0) { - scale = edge / q0; - bias = edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = std::max(w_max - w_min, eps); + bias = w_min; + } else { + scale = std::max((w_max - w_min) / n_bins, eps); + scale = mask ? scale : -scale; + + float edge = mask ? w_min : w_max; + float q0 = std::rint(edge / scale); + bias = 0; + if (q0 != 0) { + scale = edge / q0; + bias = edge; + } } size_t out_idx = i * int_per_group; for (int j = 0; j < int_per_group / bytes_per_pack; ++j) { diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6d87dc770f..83f83c9d94 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -28,13 +28,28 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; x_thread[i] = x[i]; @@ -107,13 +122,28 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; x_thread[i] = x[i]; @@ -196,13 +226,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -298,13 +342,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -395,11 +453,25 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + result[8 * i] += x * (select(U(0), scale, bool(wb & 0x01)) + bias); + result[8 * i + 1] += x * (select(U(0), scale, bool(wb & 0x02)) + bias); + result[8 * i + 2] += x * (select(U(0), scale, bool(wb & 0x04)) + bias); + result[8 * i + 3] += x * (select(U(0), scale, bool(wb & 0x08)) + bias); + result[8 * i + 4] += x * (select(U(0), scale, bool(wb & 0x10)) + bias); + result[8 * i + 5] += x * (select(U(0), scale, bool(wb & 0x20)) + bias); + result[8 * i + 6] += x * (select(U(0), scale, bool(wb & 0x40)) + bias); + result[8 * i + 7] += x * (select(U(0), scale, bool(wb & 0x80)) + bias); + } + } + + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -485,11 +557,33 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { template inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + U s[8] = { + scale, + scale / static_cast(2.0f), + scale / static_cast(4.0f), + scale / static_cast(8.0f), + scale / static_cast(16.0f), + scale / static_cast(32.0f), + scale / static_cast(64.0f), + scale / static_cast(128.0f)}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = s[0] * (w[i] & 0x01) + bias; + w_local[8 * i + 1] = s[1] * (w[i] & 0x02) + bias; + w_local[8 * i + 2] = s[2] * (w[i] & 0x04) + bias; + w_local[8 * i + 3] = s[3] * (w[i] & 0x08) + bias; + w_local[8 * i + 4] = s[4] * (w[i] & 0x10) + bias; + w_local[8 * i + 5] = s[5] * (w[i] & 0x20) + bias; + w_local[8 * i + 6] = s[6] * (w[i] & 0x40) + bias; + w_local[8 * i + 7] = s[7] * (w[i] & 0x80) + bias; + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = { scale, scale / static_cast(4.0f), @@ -578,9 +672,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -2653,14 +2747,23 @@ template w_min = simd_min(w_min); w_max = simd_max(w_max); - float scale = max((w_max - w_min) / n_bins, eps); - bool side = abs(w_min) > abs(w_max); - scale = side ? scale : -scale; - float edge = side ? w_min : w_max; - float q0 = round(edge / scale); - bool at_zero = q0 == 0.0f; - scale = at_zero ? scale : edge / q0; - float bias = at_zero ? 0 : edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = max(w_max - w_min, eps); + bias = w_min; + } else { + scale = max((w_max - w_min) / n_bins, eps); + bool side = abs(w_min) > abs(w_max); + scale = side ? scale : -scale; + float edge = side ? w_min : w_max; + float q0 = round(edge / scale); + bool at_zero = q0 == 0.0f; + scale = at_zero ? scale : edge / q0; + bias = at_zero ? 0 : edge; + } // Write out the scales and biases size_t gindex = in_index / group_size; @@ -2764,7 +2867,9 @@ template #pragma clang loop unroll(full) for (int i = 0; i < pack_factor; i++) { uint8_t d; - if (bits == 2) { + if (bits == 1) { + d = (val >> i) & 0x01; + } else if (bits == 2) { d = (val >> (bits * i)) & 0x03; } else if (bits == 4) { d = (val >> (bits * i)) & 0x0f; diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 069482cbaf..e2e50dcbb4 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -172,6 +172,7 @@ instantiate_quantized_types(32, bits) #define instantiate_quantized_all() \ + instantiate_quantized_groups(1) \ instantiate_quantized_groups(2) \ instantiate_quantized_groups(3) \ instantiate_quantized_groups(4) \ diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 31e51a5b7e..910fe73189 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -31,13 +31,28 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; x_thread[i] = x[i]; @@ -110,13 +125,28 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; x_thread[i] = x[i]; @@ -199,13 +229,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -301,13 +345,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -398,11 +456,33 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); - if (bits == 2) { + if (bits == 1) { + U s[8] = { + scale, + scale / 2.0f, + scale / 4.0f, + scale / 8.0f, + scale / 16.0f, + scale / 32.0f, + scale / 64.0f, + scale / 128.0f}; + for (int i = 0; i < (values_per_thread / 8); i++) { + result[8 * i] += x * (s[0] * (w[i] & 0x01) + bias); + result[8 * i + 1] += x * (s[1] * (w[i] & 0x02) + bias); + result[8 * i + 2] += x * (s[2] * (w[i] & 0x04) + bias); + result[8 * i + 3] += x * (s[3] * (w[i] & 0x08) + bias); + result[8 * i + 4] += x * (s[4] * (w[i] & 0x10) + bias); + result[8 * i + 5] += x * (s[5] * (w[i] & 0x20) + bias); + result[8 * i + 6] += x * (s[6] * (w[i] & 0x40) + bias); + result[8 * i + 7] += x * (s[7] * (w[i] & 0x80) + bias); + } + } + + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -487,11 +567,33 @@ template inline void dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + U s[8] = { + scale, + scale / static_cast(2.0f), + scale / static_cast(4.0f), + scale / static_cast(8.0f), + scale / static_cast(16.0f), + scale / static_cast(32.0f), + scale / static_cast(64.0f), + scale / static_cast(128.0f)}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = s[0] * (w[i] & 0x01) + bias; + w_local[8 * i + 1] = s[1] * (w[i] & 0x02) + bias; + w_local[8 * i + 2] = s[2] * (w[i] & 0x04) + bias; + w_local[8 * i + 3] = s[3] * (w[i] & 0x08) + bias; + w_local[8 * i + 4] = s[4] * (w[i] & 0x10) + bias; + w_local[8 * i + 5] = s[5] * (w[i] & 0x20) + bias; + w_local[8 * i + 6] = s[6] * (w[i] & 0x40) + bias; + w_local[8 * i + 7] = s[7] * (w[i] & 0x80) + bias; + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = { scale, scale / static_cast(4.0f), @@ -580,9 +682,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -715,9 +817,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 27302ecb5f..2cc1aeea84 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -98,6 +98,7 @@ instantiate_quantized_types(32, bits) #define instantiate_quantized_all() \ + instantiate_quantized_groups(1) \ instantiate_quantized_groups(2) \ instantiate_quantized_groups(3) \ instantiate_quantized_groups(4) \ diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 7a9afcfcc7..0dd10f3d24 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -5049,10 +5049,10 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { throw std::invalid_argument(msg.str()); } - if (bits < 2 || bits > 8 || bits == 7) { + if (bits < 1 || bits > 8 || bits == 7) { std::ostringstream msg; msg << "[quantize] The requested number of bits " << bits - << " is not supported. The supported bits are 2, 3, 4, 5, 6 and 8."; + << " is not supported. The supported bits are 1, 2, 3, 4, 5, 6 and 8."; throw std::invalid_argument(msg.str()); } @@ -5073,14 +5073,22 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { w_max = astype(w_max, float32, s); w_min = astype(w_min, float32, s); - array mask = greater(abs(w_min, s), abs(w_max, s), s); - array scales = - maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); - scales = where(mask, scales, negative(scales, s), s); - array edge = where(mask, w_min, w_max, s); - array q0 = round(divide(edge, scales, s), s); - scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); - array biases = where(equal(q0, zero, s), zero, edge, s); + array scales(0, float32); + array biases(0, float32); + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scales = maximum(subtract(w_max, w_min, s), eps, s); + biases = w_min; + } else { + array mask = greater(abs(w_min, s), abs(w_max, s), s); + scales = maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); + scales = where(mask, scales, negative(scales, s), s); + array edge = where(mask, w_min, w_max, s); + array q0 = round(divide(edge, scales, s), s); + scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); + biases = where(equal(q0, zero, s), zero, edge, s); + } packed_w = pack_and_quantize(packed_w, scales, biases, bits, s); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 4dc5114bf7..a680ef4961 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4707,14 +4707,14 @@ void init_ops(nb::module_& m) { .. table:: Quantization modes - ====== ====================== ========================== ============= ===== - mode group size bits scale type bias - ====== ====================== ========================== ============= ===== - affine 32, 64\ :sup:`*`, 128 2, 3, 4\ :sup:`*`, 5, 6, 8 same as input yes - mxfp4 32\ :sup:`*` 4\ :sup:`*` e8m0 no - mxfp8 32\ :sup:`*` 8\ :sup:`*` e8m0 no - nvfp4 16\ :sup:`*` 4\ :sup:`*` e4m3 no - ====== ====================== ========================== ============= ===== + ====== ====================== ============================== ============= ===== + mode group size bits scale type bias + ====== ====================== ============================== ============= ===== + affine 32, 64\ :sup:`*`, 128 1, 2, 3, 4\ :sup:`*`, 5, 6, 8 same as input yes + mxfp4 32\ :sup:`*` 4\ :sup:`*` e8m0 no + mxfp8 32\ :sup:`*` 8\ :sup:`*` e8m0 no + nvfp4 16\ :sup:`*` 4\ :sup:`*` e4m3 no + ====== ====================== ============================== ============= ===== :sup:`*` indicates the default value when unspecified. diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 15bc892bd8..e3f01b727d 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -22,7 +22,7 @@ class TestQuantized(mlx_tests.MLXTestCase): def test_quantize_dequantize(self): w = mx.random.normal(shape=(128, 512)) for gs in [32, 64, 128]: - for b in [2, 3, 5, 6, 4, 8]: + for b in [1, 2, 3, 5, 6, 4, 8]: with self.subTest(gs=gs, b=b): w_q, scales, biases = mx.quantize(w, group_size=gs, bits=b) w_hat = mx.dequantize(w_q, scales, biases, gs, b) @@ -33,7 +33,7 @@ def test_quantize_dequantize(self): # test quantize/dequantize 0s a = mx.zeros((256, 512)) for gs in [32, 64, 128]: - for b in [2, 3, 4, 5, 6, 8]: + for b in [1, 2, 3, 4, 5, 6, 8]: w_q, scales, biases = mx.quantize(a, gs, b) a_hat = mx.dequantize(w_q, scales, biases, gs, b) self.assertTrue(mx.all(a_hat == 0)) @@ -189,6 +189,96 @@ def test_nvfp4_quantize_dequantize(self): ) self.assertTrue(mx.allclose(w, w_hat, rtol=1e-5, atol=1e-5)) + def test_1bit_quantize_dequantize(self): + """Test 1-bit affine quantization.""" + + # Symmetric binary weights {-0.5, +0.5} should round-trip perfectly + # (affine formula gives scale=1.0, bias=-0.5) + for gs in [32, 64, 128]: + with self.subTest(gs=gs, case="pack_symmetric_weights"): + signs = (mx.random.uniform(shape=(128, 512)) > 0.5).astype(mx.float32) + w = signs * 1.0 - (1 - signs) * 1.0 # {-1.0, +1.0} + w = w * 0.5 # {-0.5, +0.5} + + w_q, scales, biases = mx.quantize(w, group_size=gs, bits=1) + w_hat = mx.dequantize(w_q, scales, biases, gs, 1) + + self.assertLess((w - w_hat).abs().max(), 1e-5) + + # Asymmetric binary weights {0.1, 0.9} should round-trip perfectly + # (affine formula gives scale=0.8, bias=0.1) + for gs in [32, 64, 128]: + with self.subTest(gs=gs, case="pack_asymmetric_weights"): + bits = (mx.random.uniform(shape=(128, 512)) > 0.5).astype(mx.float32) + w = bits * 0.9 + (1 - bits) * 0.1 # {0.1, 0.9} + + w_q, scales, biases = mx.quantize(w, group_size=gs, bits=1) + w_hat = mx.dequantize(w_q, scales, biases, gs, 1) + + self.assertLess((w - w_hat).abs().max(), 1e-5) + + # Verify dequantized values are exactly {bias, bias + scale} + w = mx.random.normal(shape=(64, 256)) + for gs in [32, 64, 128]: + with self.subTest(gs=gs, case="dequant_values"): + w_q, scales, biases = mx.quantize(w, group_size=gs, bits=1) + w_hat = mx.dequantize(w_q, scales, biases, gs, 1) + + for i in range(scales.shape[0]): + for j in range(scales.shape[1]): + s = scales[i, j].item() + b = biases[i, j].item() + row_start = j * gs + row_end = row_start + gs + vals = w_hat[i, row_start:row_end] + mx.eval(vals) + for v in vals.tolist(): + self.assertTrue( + abs(v - b) < 1e-5 or abs(v - (b + s)) < 1e-5, + f"Value {v} not in {{bias={b}, bias+scale={b+s}}}", + ) + + # 1-bit quantize/dequantize zeros — scale floors to eps, bias=0 + a = mx.zeros((256, 512)) + for gs in [32, 64, 128]: + w_q, scales, biases = mx.quantize(a, gs, 1) + a_hat = mx.dequantize(w_q, scales, biases, gs, 1) + self.assertLess(a_hat.abs().max(), 1e-5) + + # Quantized matmul with symmetric binary weights + key = mx.random.key(42) + k1, k2 = mx.random.split(key) + for gs in [32, 64, 128]: + with self.subTest(gs=gs, case="quantized_matmul_symmetric"): + x = mx.random.normal(shape=(4, 256), key=k1) + signs = (mx.random.uniform(shape=(128, 256), key=k2) > 0.5).astype( + mx.float32 + ) + w = signs * 0.3 - (1 - signs) * 0.3 # {-0.3, +0.3} + + w_q, scales, biases = mx.quantize(w, gs, 1) + w_hat = mx.dequantize(w_q, scales, biases, gs, 1) + y_q = mx.quantized_matmul(x, w_q, scales, biases, True, gs, 1) + y_hat = x @ w_hat.T + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-5) + + # Quantized matmul with asymmetric binary weights + for gs in [32, 64, 128]: + with self.subTest(gs=gs, case="quantized_matmul_asymmetric"): + x = mx.random.normal(shape=(4, 256), key=k1) + bits = (mx.random.uniform(shape=(128, 256), key=k2) > 0.5).astype( + mx.float32 + ) + w = bits * 0.7 + (1 - bits) * 0.1 # {0.1, 0.7} + + w_q, scales, biases = mx.quantize(w, gs, 1) + w_hat = mx.dequantize(w_q, scales, biases, gs, 1) + y_q = mx.quantized_matmul(x, w_q, scales, biases, True, gs, 1) + y_hat = x @ w_hat.T + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-5) + def test_qqmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -291,7 +381,7 @@ def test_qmm(self): dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 tests = product( [128, 64, 32], # group_size - [2, 4, 8], # bits + [1, 2, 4, 8], # bits [8, 32, 33, 64], # M [128, 256], # N [128, 256], # K From 5743cb8a79ee43c7763a3fa75d5c0da740af286b Mon Sep 17 00:00:00 2001 From: Pasha Khosravi Date: Fri, 6 Mar 2026 18:06:31 -0800 Subject: [PATCH 206/222] Guard fast-path Metal kernel dispatch for 1-bit quantization --- mlx/backend/metal/quantized.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index f659e16c93..4f324f808a 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1751,7 +1751,8 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { + if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && + !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } From 3e0fa1fa13c2371b3c7beeefb41582f7d622c279 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Fri, 5 Jun 2026 07:49:09 -0700 Subject: [PATCH 207/222] metal: 1-bit qmv_fast use 1 pack/thread for occupancy (#3) affine qmv_fast set packs_per_thread=2 for all bits except 2-bit, so 1-bit got values_per_thread=64 (x_thread[64], ~256B/thread) -> low occupancy -> 1-bit decode saturates only ~75% of M5 DRAM BW vs ~90/96% for 2/4-bit. Use 1 pack/thread for bits<=2 (values_per_thread=32), matching 2-bit's register footprint. Measured on M5 Pro (distinct-weight DRAM-bound, 2-bit as drift control): 1-bit 24.0 -> 21.9 us/matvec (~9%, 75->82% BW), 2-bit control 33.2->32.9 (0.8% drift); correct, rel_err 2.7e-4. Also makes scale_step_per_thread (=group_size/values_per_thread) well-defined for group_size=32 at 1-bit. --- mlx/backend/metal/kernels/quantized.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 83f83c9d94..ac315b867f 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -853,7 +853,7 @@ METAL_FUNC void qmv_fast_impl( uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { - constexpr int packs_per_thread = bits == 2 ? 1 : 2; + constexpr int packs_per_thread = bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor(); From 564c1160fe24efea8b2e0380fc67a88d208278a7 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Thu, 23 Jul 2026 06:32:40 -0400 Subject: [PATCH 208/222] Disable NAX on generation 17 for low-bit correctness --- mlx/backend/metal/device.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 65df5c108c..252f72bd9a 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -960,9 +960,10 @@ bool is_nax_available() { can_use_nax = true; } auto& d = metal::device(mlx::core::Device::gpu); - auto arch = d.get_architecture().back(); auto gen = d.get_architecture_gen(); - can_use_nax &= gen >= (arch == 'p' ? 18 : 17); + // Generation 17 advertises NAX support, but produces incorrect results for + // some low-bit affine quantized matrix-vector shapes. + can_use_nax &= gen >= 18; return can_use_nax; }; static bool is_nax_available_ = _check_nax(); From 629f09185a0c2dbd7b892c577f9a113a87e033ee Mon Sep 17 00:00:00 2001 From: Brian <288398250+bri-prism@users.noreply.github.com> Date: Mon, 13 Jul 2026 23:28:44 +0100 Subject: [PATCH 209/222] Route affine qmv_wide by bit-width and batch size The qmv_wide kernel dequantizes a full group of weights into registers and reuses them across the M input rows. That amortization only pays off for 2-bit once three or more rows share a group; at M=2 it breaks even, and for 1-bit the weight traffic is small enough that the per-row dequant dominates and the specialized qmv is faster. Gate the dispatch accordingly: 1-bit and 2-bit M<2 keep the specialized qmv; 2-bit routes to qmv_wide only at M>=3 on gen-15+; fp modes are unchanged. Add 1-bit affine coverage (full sweep + tiny shapes) to test_qmv_wide. Measured on affine 2-bit matvecs across a range of projection shapes: qmv_wide wins at widths >=3 (more so as N grows), with no regression at width <=2 or for 1-bit. --- mlx/backend/metal/quantized.cpp | 16 +++++++++++++--- python/tests/test_quantized.py | 9 +++++++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 4f324f808a..96aaa96532 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -535,8 +535,18 @@ void qmv( } // affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. -inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { - return mode != "affine" || d.get_architecture_gen() >= 15; +// The 1-bit path has so little weight traffic that unpacking into registers +// dominates even when several input rows reuse the result. The 2-bit path +// breaks even after two rows and wins once three or more rows share a block. +inline bool +use_qmv_wide(const std::string& mode, int bits, int M, metal::Device& d) { + if (mode != "affine") { + return true; + } + if (bits == 1 || (bits == 2 && M < 3)) { + return false; + } + return d.get_architecture_gen() >= 15; } // Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): @@ -1759,7 +1769,7 @@ void dispatch_qmv( // Small batch so route to qmv_wide, which reuses each weight group across the // M vectors. - if (M >= 2 && use_qmv_wide(mode, d) && !global_scale) { + if (M >= 2 && use_qmv_wide(mode, bits, M, d) && !global_scale) { qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index e3f01b727d..768be86a45 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -744,7 +744,7 @@ def test_qmv_wide(self): # Affine: every bit-width and group size. for group_size, bits, K in product( - [32, 64, 128], [2, 3, 4, 5, 6, 8], [128, 512] + [32, 64, 128], [1, 2, 3, 4, 5, 6, 8], [128, 512] ): for M, N, B in product(Ms, Ns, Bs): with self.subTest(M=M, N=N, K=K, B=B, group_size=group_size, bits=bits): @@ -778,7 +778,12 @@ def test_qmv_wide(self): # Tiny shapes (M, K, N): small K and non-multiple output rows. tiny = [(2, 32, 10), (4, 32, 7), (3, 64, 5), (5, 64, 3)] - settings = [(4, 32, "affine"), (6, 32, "affine"), (4, 16, "nvfp4")] + settings = [ + (1, 32, "affine"), + (4, 32, "affine"), + (6, 32, "affine"), + (4, 16, "nvfp4"), + ] for M, K, N in tiny: for bits, group_size, mode in settings: with self.subTest( From 3ebf6207f7577a7ff684ed7924dd936644ca97e7 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Thu, 16 Jul 2026 21:54:39 -0300 Subject: [PATCH 210/222] Add native affine 1-bit CUDA QMV --- mlx/backend/cuda/device/cute_dequant.cuh | 28 ++++++++++ mlx/backend/cuda/quantized/affine_quantize.cu | 32 ++++++++--- mlx/backend/cuda/quantized/qmm/qmm.cu | 12 ++++ mlx/backend/cuda/quantized/qmm/qmv.cu | 7 ++- python/tests/test_quantized.py | 56 ++++++++++++++++++- 5 files changed, 123 insertions(+), 12 deletions(-) diff --git a/mlx/backend/cuda/device/cute_dequant.cuh b/mlx/backend/cuda/device/cute_dequant.cuh index 6416c5b87a..d72500c4da 100644 --- a/mlx/backend/cuda/device/cute_dequant.cuh +++ b/mlx/backend/cuda/device/cute_dequant.cuh @@ -9,9 +9,37 @@ namespace cutlass { +using uint1b_mlx_t = integer_subbyte<1, false>; using uint3b_t = integer_subbyte<3, false>; using uint5b_t = integer_subbyte<5, false>; +template +struct NumericArrayConverter { + static_assert(N % 8 == 0); + + using result_type = Array; + using source_type = Array; + + CUTLASS_HOST_DEVICE + static result_type convert(const source_type& source) { + result_type result; + auto* packed = reinterpret_cast(&source); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 8; ++i) { + CUTLASS_PRAGMA_UNROLL + for (int bit = 0; bit < 8; ++bit) { + result[i * 8 + bit] = T((packed[i] >> bit) & 0x01); + } + } + return result; + } + + CUTLASS_HOST_DEVICE + result_type operator()(const source_type& source) const { + return convert(source); + } +}; + template struct NumericArrayConverter { static_assert(N % 8 == 0); diff --git a/mlx/backend/cuda/quantized/affine_quantize.cu b/mlx/backend/cuda/quantized/affine_quantize.cu index 13bd35dade..728f31e731 100644 --- a/mlx/backend/cuda/quantized/affine_quantize.cu +++ b/mlx/backend/cuda/quantized/affine_quantize.cu @@ -63,14 +63,23 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) { w_min = cg::reduce(warp, w_min, min_op); w_max = cg::reduce(warp, w_max, max_op); - float scale = max((w_max - w_min) / n_bins, eps); - bool side = abs(w_min) > abs(w_max); - scale = side ? scale : -scale; - float edge = side ? w_min : w_max; - float q0 = round(edge / scale); - bool at_zero = q0 == 0.0f; - scale = at_zero ? scale : edge / q0; - float bias = at_zero ? 0 : edge; + float scale; + float bias; + + if constexpr (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = max(w_max - w_min, eps); + bias = w_min; + } else { + scale = max((w_max - w_min) / n_bins, eps); + bool side = abs(w_min) > abs(w_max); + scale = side ? scale : -scale; + float edge = side ? w_min : w_max; + float q0 = round(edge / scale); + bool at_zero = q0 == 0.0f; + scale = at_zero ? scale : edge / q0; + bias = at_zero ? 0 : edge; + } // Write out the scales and biases size_t gindex = in_index / group_size; @@ -212,7 +221,9 @@ __global__ void affine_dequantize( #pragma clang loop unroll(full) for (int i = 0; i < pack_factor; i++) { uint8_t d; - if (bits == 2) { + if (bits == 1) { + d = (val >> i) & 0x01; + } else if (bits == 2) { d = (val >> (bits * i)) & 0x03; } else if (bits == 4) { d = (val >> (bits * i)) & 0x0f; @@ -244,6 +255,9 @@ void dispatch_groups(int group_size, F&& f) { template void dispatch_bits(int bits, F&& f) { switch (bits) { + case 1: + f(std::integral_constant{}); + break; case 2: f(std::integral_constant{}); break; diff --git a/mlx/backend/cuda/quantized/qmm/qmm.cu b/mlx/backend/cuda/quantized/qmm/qmm.cu index 5c1d4f76b4..62d5616e30 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm.cu @@ -168,6 +168,18 @@ bool supports_qmv( QuantizationMode mode, cu::Device& device) { int k = x.shape(-1); + if (mode == QuantizationMode::Affine) { + if (bits != 1 && bits != 2 && bits != 3 && bits != 4 && bits != 5 && + bits != 6 && bits != 8) { + return false; + } + if (group_size != 32 && group_size != 64 && group_size != 128) { + return false; + } + if (k % group_size != 0 || !biases) { + return false; + } + } if (k % 8 != 0) { return false; } diff --git a/mlx/backend/cuda/quantized/qmm/qmv.cu b/mlx/backend/cuda/quantized/qmm/qmv.cu index 4ec6b95a7a..ee4943b285 100644 --- a/mlx/backend/cuda/quantized/qmm/qmv.cu +++ b/mlx/backend/cuda/quantized/qmm/qmv.cu @@ -118,7 +118,8 @@ __device__ __forceinline__ void qmv_kernel_impl( } // Accumulations of current row. - cuda::std::conditional_t<(bits >= 8), float, T> sums[elems_per_thread] = {}; + cuda::std::conditional_t<(bits == 1 || bits >= 8), float, T> + sums[elems_per_thread] = {}; auto dequant_fma_tile = [&](int idx) { S scale = scales[idx / group_size]; @@ -383,7 +384,9 @@ inline void dispatch_quant_types( f.template operator()(); } else { dispatch_groups(group_size, tag, [&]() { - if (bits == 2) { + if (bits == 1) { + f.template operator()(); + } else if (bits == 2) { f.template operator()(); } else if (bits == 3) { f.template operator()(); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 768be86a45..8d03884de4 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -279,6 +279,60 @@ def test_1bit_quantize_dequantize(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-5) + def test_1bit_qmv_packed_bit_order(self): + bit_positions = [0, 1, 7, 8, 15, 16, 31] + packed_word = sum(1 << bit for bit in bit_positions) + packed = mx.array([[packed_word, 0, packed_word, 0]], dtype=mx.uint32) + scales = mx.ones((1, 4), dtype=mx.float32) + biases = mx.zeros((1, 4), dtype=mx.float32) + x = mx.arange(1, 129, dtype=mx.float32).reshape(1, 128) + + actual = mx.quantized_matmul( + x, + packed, + scales, + biases, + transpose=True, + group_size=32, + bits=1, + ) + expected = sum(position + 1 for position in bit_positions) + sum( + position + 65 for position in bit_positions + ) + self.assertEqual(actual.shape, (1, 1)) + self.assertEqual(actual.item(), expected) + + def test_1bit_qmv_dtypes_groups_and_odd_rows(self): + key = mx.random.key(1729) + k1, k2 = mx.random.split(key) + for dtype, tol in [ + (mx.float16, 2e-2), + (mx.bfloat16, 1e-1), + (mx.float32, 1e-5), + ]: + for group_size in [32, 64, 128]: + with self.subTest(dtype=dtype, group_size=group_size): + x = mx.random.normal((3, 256), key=k1).astype(dtype) + signs = ( + mx.random.uniform(shape=(67, 256), key=k2) > 0.5 + ).astype(dtype) + w = signs * 0.6 - (1 - signs) * 0.2 + w_q, scales, biases = mx.quantize(w, group_size, 1) + w_hat = mx.dequantize( + w_q, + scales, + biases, + group_size, + 1, + dtype=dtype, + ) + actual = mx.quantized_matmul( + x, w_q, scales, biases, True, group_size, 1 + ) + expected = x @ w_hat.T + self.assertEqual(actual.shape, (3, 67)) + self.assertLess((actual - expected).abs().max(), tol) + def test_qqmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -620,7 +674,7 @@ def test_qmv(self): k1, k2 = mx.random.split(key) tests = product( [128, 64, 32], # group_size - [2, 3, 4, 5, 6, 8], # bits + [1, 2, 3, 4, 5, 6, 8], # bits [256, 512, 67], # M [64, 256], # N [0, 1, 3, 8], # B From dde9c8359305e29c96256351bf04488780a77033 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Thu, 23 Jul 2026 06:38:33 -0400 Subject: [PATCH 211/222] Format rebased low-bit changes --- mlx/backend/metal/kernels/quantized.h | 3 ++- python/tests/test_quantized.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index ac315b867f..c7e15ef6b1 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -853,7 +853,8 @@ METAL_FUNC void qmv_fast_impl( uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { - constexpr int packs_per_thread = bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy + constexpr int packs_per_thread = + bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor(); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 8d03884de4..24039557c7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -313,9 +313,9 @@ def test_1bit_qmv_dtypes_groups_and_odd_rows(self): for group_size in [32, 64, 128]: with self.subTest(dtype=dtype, group_size=group_size): x = mx.random.normal((3, 256), key=k1).astype(dtype) - signs = ( - mx.random.uniform(shape=(67, 256), key=k2) > 0.5 - ).astype(dtype) + signs = (mx.random.uniform(shape=(67, 256), key=k2) > 0.5).astype( + dtype + ) w = signs * 0.6 - (1 - signs) * 0.2 w_q, scales, biases = mx.quantize(w, group_size, 1) w_hat = mx.dequantize( From f8afc8b3e2bb5bcdf6e8d48e53c8fa630637d2fd Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Sat, 25 Jul 2026 23:35:31 -0400 Subject: [PATCH 212/222] Route low-bit CUDA matmuls through native QMV --- mlx/backend/cuda/quantized/quantized.cpp | 9 +++++++++ python/tests/test_quantized.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 4d25f3c3e0..18326712ff 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -103,6 +103,15 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = x.shape(-1); int B = out.size() / (M * N); + // The affine 1-bit and 2-bit implementations live in qmv. The generic + // matrix kernels either do not support these packed weights (1-bit) or can + // produce an invalid large-M launch on post-sm90 devices (2-bit). + if (can_use_qmv && mode_ == QuantizationMode::Affine && + (bits_ == 1 || bits_ == 2)) { + call_qmv(); + return; + } + if (can_use_qmm_sm90) { if (can_use_qmv && (M == 1 && B == 1 && N <= 16384 && K <= 16384)) { call_qmv(); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 24039557c7..f43f00d8a4 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -333,6 +333,21 @@ def test_1bit_qmv_dtypes_groups_and_odd_rows(self): self.assertEqual(actual.shape, (3, 67)) self.assertLess((actual - expected).abs().max(), tol) + def test_low_bit_affine_large_m(self): + key = mx.random.key(2026) + k1, k2 = mx.random.split(key) + x = mx.random.normal(shape=(16, 256), key=k1) + w = mx.random.normal(shape=(67, 256), key=k2) + + for bits in [1, 2]: + with self.subTest(bits=bits): + w_q, scales, biases = mx.quantize(w, 128, bits) + w_hat = mx.dequantize(w_q, scales, biases, 128, bits) + actual = mx.quantized_matmul(x, w_q, scales, biases, True, 128, bits) + expected = x @ w_hat.T + self.assertEqual(actual.shape, expected.shape) + self.assertLess((actual - expected).abs().max(), 1e-3) + def test_qqmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From d285f7248ad6572a5cf46de4e4fc68825ce53963 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Sun, 26 Jul 2026 00:05:55 -0400 Subject: [PATCH 213/222] Keep global CUDA streams alive through process exit --- mlx/backend/cuda/device.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 30248f5568..076cc4f519 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -583,8 +583,13 @@ std::unordered_map& get_command_encoders() { } std::unordered_map& get_global_command_encoders() { - static std::unordered_map encoders; - return encoders; + // Keep process-global encoders alive through static destruction. The CUDA + // runtime may already be shutting down when ordinary static destructors run, + // and CommandEncoder::~CommandEncoder synchronizes its CUDA stream. The OS + // reclaims these process-lifetime resources on exit; explicit clear_streams() + // still releases them when callers need deterministic teardown. + static auto* encoders = new std::unordered_map; + return *encoders; } } // namespace mlx::core::cu From 9190776d7f56dcb38edbc2ac248b7304ac9a8103 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Mon, 27 Jul 2026 12:57:39 -0300 Subject: [PATCH 214/222] Expose quantized headers to custom Metal kernels --- mlx/backend/metal/custom_kernel.cpp | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index b73dd72def..5e7cc461a8 100644 --- a/mlx/backend/metal/custom_kernel.cpp +++ b/mlx/backend/metal/custom_kernel.cpp @@ -49,8 +49,33 @@ void CustomKernel::eval_gpu( std::string lib_name = fmt::format( "{}_{:x}_{}", name_, std::hash{}(source_), compile_options_); - auto lib = d.get_library( - lib_name, compile_options_, [this] { return metal::utils() + source_; }); + auto lib = d.get_library(lib_name, compile_options_, [this] { + if (source_.find("MLX_INCLUDE_AFFINE_QUANTIZED_HEADERS") != + std::string::npos) { + std::string kernel_source; + concatenate( + kernel_source, + metal::utils(), + metal::quantized_utils(), + metal::gemm(), + metal::quantized(), + source_); + return kernel_source; + } + if (source_.find("MLX_INCLUDE_FP_QUANTIZED_HEADERS") != + std::string::npos) { + std::string kernel_source; + concatenate( + kernel_source, + metal::utils(), + metal::quantized_utils(), + metal::gemm(), + metal::fp_quantized(), + source_); + return kernel_source; + } + return metal::utils() + source_; + }); auto kernel = d.get_kernel(name_, lib); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); From 4d88c9d87f9714d36e736dab762a0b91a3fc7bc7 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Wed, 29 Jul 2026 01:37:31 -0300 Subject: [PATCH 215/222] Speed up NVFP4 block-loader staging Adapt the bit-exact NVFP4 loader rewrite from mlx.fast submission 4228f9e1-18a5-4a1a-8047-96dc4e9289ca (validated commit b8fe3af). Fold fp4's exact 2^14 renormalization into the e4m3 group scale and decode eight packed nibbles per uint32 in the regular and NAX loaders. Other quantization modes keep the scalar path.\n\nThe source submission exhaustively checked 404,226,048 staged values with zero bit mismatches. This adaptation adds a local 524,288-value BF16 equivalence gate in the parent mlx-swift repository.\n\nCo-authored-by: anupsv <6407789+anupsv@users.noreply.github.com> --- mlx/backend/metal/kernels/fp_quantized.h | 76 ++++++++++++++--- mlx/backend/metal/kernels/fp_quantized_nax.h | 85 +++++++++++++++----- 2 files changed, 133 insertions(+), 28 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 6e77569f56..9d398ab1fa 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -145,6 +145,48 @@ inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { } } +// NVFP4 block-loader staging fast path. fp4_e2m1 embeds its three magnitude +// bits in a half value that is smaller by exactly 2^14, then renormalizes each +// value. For e4m3 group scales, moving that exact power-of-two factor to the +// once-per-group scale preserves every staged bit while removing the +// per-value renormalization. Decoding four packed bytes at once constructs the +// same eight half bit patterns as the scalar nibble walk. The source version +// of this rewrite was exhaustively checked for bfloat, half, and float across +// 404,226,048 staged values with zero mismatches. +static inline float fp4nv_scale_x16384(uint8_t s) { + return float(*(thread fp8_e4m3*)(&s)) * 16384.0f; +} + +// packed_uchar4 has byte alignment, so this does not add an alignment +// precondition to the original byte-at-a-time access. +static inline uint32_t fp4nv_pack4(const device uint8_t* p) { + return as_type(uchar4(*(const device packed_uchar4*)p)); +} + +template +static inline void fp4nv_decode8(uint32_t c, float scale, thread T* out) { + const float2 v0 = float2(as_type( + ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * + scale; + const float2 v1 = float2(as_type( + ((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * + scale; + const float2 v2 = float2(as_type( + ((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * + scale; + const float2 v3 = + float2(as_type(((c & 0x70007000u) >> 3) | (c & 0x80008000u))) * + scale; + out[0] = T(v0.x); + out[1] = T(v1.x); + out[2] = T(v2.x); + out[3] = T(v3.x); + out[4] = T(v0.y); + out[5] = T(v1.y); + out[6] = T(v2.y); + out[7] = T(v3.y); +} + template < typename T, short BROWS, @@ -203,16 +245,34 @@ struct QuantizedBlockLoader { scales_ + bi * src_ld / group_size + (bj * pack_factor) / group_size) {} + MLX_MTL_CONST bool fp4nv_fast = (bits == 4) && (group_size == 16) && + (bytes_per_pack == 1) && (n_reads >= 4) && ((n_reads % 4) == 0); + + void stage() const thread { + if constexpr (fp4nv_fast) { + const float scale = fp4nv_scale_x16384(*scales); + for (int i = 0; i < n_reads / 4; i++) { + T vals[8]; + fp4nv_decode8(fp4nv_pack4(src + i * 4), scale, vals); + for (int j = 0; j < 8; j++) { + dst[i * 8 + j] = vals[j]; + } + } + } else { + T scale = dequantize_scale(*scales); + for (int i = 0; i < n_reads; i++) { + dequantize( + src[i * bytes_per_pack], scale, dst + i * pack_factor); + } + } + } + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); - for (int i = 0; i < n_reads; i++) { - dequantize( - src[i * bytes_per_pack], scale, dst + i * pack_factor); - } + stage(); } void load_safe(short2 src_tile_dim) const thread { @@ -234,11 +294,7 @@ struct QuantizedBlockLoader { return; } - T scale = dequantize_scale(*scales); - for (int i = 0; i < n_reads; i++) { - dequantize( - src[i * bytes_per_pack], scale, dst + i * pack_factor); - } + stage(); } void next() thread { diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index cf64ff7f46..89c453a7cb 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -69,6 +69,40 @@ inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { } } +// NVFP4 block-loader staging fast path. This is the NAX twin of the exact +// power-of-two scale fold and packed-nibble decode in fp_quantized.h. +static inline float fp4nv_scale_x16384(uint8_t s) { + return float(*(thread fp8_e4m3*)(&s)) * 16384.0f; +} + +static inline uint32_t fp4nv_pack4(const device uint8_t* p) { + return as_type(uchar4(*(const device packed_uchar4*)p)); +} + +template +static inline void fp4nv_decode8(uint32_t c, float scale, thread T* out) { + const float2 v0 = float2(as_type( + ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * + scale; + const float2 v1 = float2(as_type( + ((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * + scale; + const float2 v2 = float2(as_type( + ((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * + scale; + const float2 v3 = + float2(as_type(((c & 0x70007000u) >> 3) | (c & 0x80008000u))) * + scale; + out[0] = T(v0.x); + out[1] = T(v1.x); + out[2] = T(v2.x); + out[3] = T(v3.x); + out[4] = T(v0.y); + out[5] = T(v1.y); + out[6] = T(v2.y); + out[7] = T(v3.y); +} + template < typename T, short BROWS, @@ -127,20 +161,43 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id) {} + MLX_MTL_CONST bool fp4nv_fast = (bits == 4) && (group_size == 16) && + (bytes_per_pack == 1) && (n_reads_per_scale >= 4) && + ((n_reads_per_scale % 4) == 0); + + void stage() const thread { + if constexpr (fp4nv_fast) { + int k = 0; + for (int i = 0; i < n_steps_per_read; i++) { + const float scale = fp4nv_scale_x16384(scales[i]); + for (int j = 0; j < n_reads_per_scale / 4; j++) { + T vals[8]; + fp4nv_decode8(fp4nv_pack4(src + k), scale, vals); + for (int e = 0; e < 8; e++) { + dst[k * pack_factor + e] = vals[e]; + } + k += 4; + } + } + } else { + int k = 0; + for (int i = 0; i < n_steps_per_read; i++) { + T scale = dequantize_scale(scales[i]); + for (int j = 0; j < n_reads_per_scale; j++) { + dequantize( + src[k * bytes_per_pack], scale, dst + k * pack_factor); + k++; + } + } + } + } + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - int k = 0; - for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); - for (int j = 0; j < n_reads_per_scale; j++) { - dequantize( - src[k * bytes_per_pack], scale, dst + k * pack_factor); - k++; - } - } + stage(); } void load_safe(short2 src_tile_dim) const thread { @@ -162,15 +219,7 @@ struct QuantizedBlockLoader { return; } - int k = 0; - for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); - for (int j = 0; j < n_reads_per_scale; j++) { - dequantize( - src[k * bytes_per_pack], scale, dst + k * pack_factor); - k++; - } - } + stage(); } void next() thread { From 0dfe6cbd9e47abae408fa818af3b1e2a4cf14dce Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Wed, 5 Aug 2026 09:19:13 -0300 Subject: [PATCH 216/222] Tune large H3 GEMMs on M4 Max --- mlx/backend/metal/matmul.cpp | 38 +++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 8c0b46d6de..3a9920d73f 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -373,6 +373,40 @@ void steel_matmul_regular_axpby( char devc = d.get_architecture().back(); GEMM_TPARAM_MACRO(devc) + if (int value = env::get_var("MLX_GEMM_BM", 0); value > 0) { + bm = value; + } + if (int value = env::get_var("MLX_GEMM_BN", 0); value > 0) { + bn = value; + } + if (int value = env::get_var("MLX_GEMM_BK", 0); value > 0) { + bk = value; + } + if (int value = env::get_var("MLX_GEMM_WM", 0); value > 0) { + wm = value; + } + if (int value = env::get_var("MLX_GEMM_WN", 0); value > 0) { + wn = value; + } + const bool is_h3_projection_shape = M >= 32768 && + ((N == 21504 && K == 5376) || (N == 5376 && K == 7168) || + (N == 28672 && K == 5376) || (N == 5376 && K == 14336)); + const int h3_tuning_override = env::get_var("MLX_GEMM_H3_TUNED", -1); + const bool use_h3_tuned_schedule = out.dtype() != float32 && + ((h3_tuning_override < 0 && is_h3_projection_shape) || + (h3_tuning_override == 1 && M >= 32768 && N >= 4096)); + if (use_h3_tuned_schedule) { + bn = 64; + bk = 16; + wn = 2; + if (N >= 24576 || (N <= 8192 && K < 10000)) { + bm = 32; + wm = 1; + } else { + bm = 64; + wm = 2; + } + } // Prepare kernel name std::ostringstream kname; @@ -437,7 +471,9 @@ void steel_matmul_regular_axpby( int tm = (M + bm - 1) / bm; // TODO: Explore device-based tuning for swizzle - int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2); + int swizzle_log = use_h3_tuned_schedule + ? 2 + : env::get_var("MLX_GEMM_SWIZZLE_LOG", 0); // Prepare steel matmul params GEMMParams params{/* const int M = */ M, From 3a0382c22809e7dfee24cc279df325d7fcab9fd7 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 07:20:54 -0300 Subject: [PATCH 217/222] Align 1-bit qmv fast dispatch geometry --- mlx/backend/metal/quantized.cpp | 4 ++-- python/tests/test_quantized.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 96aaa96532..22a4ff0ae3 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -143,9 +143,9 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } // Must match the K step in qmv_fast_impl (kernels/quantized.h): -// pack_factor() * (bits == 2 ? 1 : 2) * SIMD_SIZE +// pack_factor() * (bits <= 2 ? 1 : 2) * SIMD_SIZE inline int qmv_fast_k_alignment(int bits) { - return get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32; + return get_pack_factor(bits, 32) * (bits <= 2 ? 1 : 2) * 32; } inline int add_strides_and_shapes( diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index f43f00d8a4..65086d815f 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -711,6 +711,24 @@ def test_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmv_fast_1bit_alignment(self): + # The Metal host dispatch and qmv_fast kernel must agree that 1-bit + # weights use one 32-value pack per lane: 32 * 1 * 32 = 1024 values. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + K = 1024 + N = 67 + x = mx.random.normal(shape=(1, K), key=k1) / K**0.5 + w = mx.random.normal(shape=(N, K), key=k2) / K**0.5 + w_q, scales, biases = mx.quantize(w, group_size=128, bits=1) + w_hat = mx.dequantize(w_q, scales, biases, group_size=128, bits=1) + y_q = mx.quantized_matmul( + x, w_q, scales, biases, True, group_size=128, bits=1 + ) + y_hat = x @ w_hat.T + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From b34c2353ba02a6034f1c21efce72d09cf51e0d08 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 07:21:55 -0300 Subject: [PATCH 218/222] Format refreshed MLX patches --- mlx/backend/metal/custom_kernel.cpp | 3 +-- mlx/backend/metal/kernels/fp_quantized.h | 16 ++++++++++------ mlx/backend/metal/kernels/fp_quantized_nax.h | 16 ++++++++++------ mlx/backend/metal/matmul.cpp | 5 ++--- python/tests/test_quantized.py | 4 +--- 5 files changed, 24 insertions(+), 20 deletions(-) diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index 5e7cc461a8..f2784a9c05 100644 --- a/mlx/backend/metal/custom_kernel.cpp +++ b/mlx/backend/metal/custom_kernel.cpp @@ -62,8 +62,7 @@ void CustomKernel::eval_gpu( source_); return kernel_source; } - if (source_.find("MLX_INCLUDE_FP_QUANTIZED_HEADERS") != - std::string::npos) { + if (source_.find("MLX_INCLUDE_FP_QUANTIZED_HEADERS") != std::string::npos) { std::string kernel_source; concatenate( kernel_source, diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 9d398ab1fa..fc4a580339 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -165,14 +165,18 @@ static inline uint32_t fp4nv_pack4(const device uint8_t* p) { template static inline void fp4nv_decode8(uint32_t c, float scale, thread T* out) { - const float2 v0 = float2(as_type( - ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * + const float2 v0 = + float2( + as_type( + ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * scale; - const float2 v1 = float2(as_type( - ((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * + const float2 v1 = + float2( + as_type(((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * scale; - const float2 v2 = float2(as_type( - ((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * + const float2 v2 = + float2( + as_type(((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * scale; const float2 v3 = float2(as_type(((c & 0x70007000u) >> 3) | (c & 0x80008000u))) * diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 89c453a7cb..09179846b1 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -81,14 +81,18 @@ static inline uint32_t fp4nv_pack4(const device uint8_t* p) { template static inline void fp4nv_decode8(uint32_t c, float scale, thread T* out) { - const float2 v0 = float2(as_type( - ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * + const float2 v0 = + float2( + as_type( + ((c & 0x00070007u) << 9) | ((c & 0x00080008u) << 12))) * scale; - const float2 v1 = float2(as_type( - ((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * + const float2 v1 = + float2( + as_type(((c & 0x00700070u) << 5) | ((c & 0x00800080u) << 8))) * scale; - const float2 v2 = float2(as_type( - ((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * + const float2 v2 = + float2( + as_type(((c & 0x07000700u) << 1) | ((c & 0x08000800u) << 4))) * scale; const float2 v3 = float2(as_type(((c & 0x70007000u) >> 3) | (c & 0x80008000u))) * diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 3a9920d73f..2e5fc06b33 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -471,9 +471,8 @@ void steel_matmul_regular_axpby( int tm = (M + bm - 1) / bm; // TODO: Explore device-based tuning for swizzle - int swizzle_log = use_h3_tuned_schedule - ? 2 - : env::get_var("MLX_GEMM_SWIZZLE_LOG", 0); + int swizzle_log = + use_h3_tuned_schedule ? 2 : env::get_var("MLX_GEMM_SWIZZLE_LOG", 0); // Prepare steel matmul params GEMMParams params{/* const int M = */ M, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 65086d815f..8ccda7f868 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -722,9 +722,7 @@ def test_qmv_fast_1bit_alignment(self): w = mx.random.normal(shape=(N, K), key=k2) / K**0.5 w_q, scales, biases = mx.quantize(w, group_size=128, bits=1) w_hat = mx.dequantize(w_q, scales, biases, group_size=128, bits=1) - y_q = mx.quantized_matmul( - x, w_q, scales, biases, True, group_size=128, bits=1 - ) + y_q = mx.quantized_matmul(x, w_q, scales, biases, True, group_size=128, bits=1) y_hat = x @ w_hat.T self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) From 2ccca66fcf65fb681ce754add31f9d5fbcc1bf61 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 07:25:09 -0300 Subject: [PATCH 219/222] Link quantized custom-kernel sources in AOT builds --- mlx/backend/metal/CMakeLists.txt | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index ea4a995ade..e9bc749307 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -35,6 +35,19 @@ make_jit_source(indexing/gather_axis) make_jit_source(indexing/scatter_axis) make_jit_source(hadamard) +# Custom Metal kernels can opt into the quantized helper headers even when the +# built-in kernels use the precompiled metallib. Keep the corresponding source +# strings linked in both AOT and JIT builds so CustomKernel::eval_gpu can build +# those marked kernels at runtime. +make_jit_source( + steel/gemm/gemm kernels/steel/utils.h kernels/steel/gemm/loader.h + kernels/steel/gemm/mma.h kernels/steel/gemm/params.h + kernels/steel/gemm/transforms.h) +make_jit_source(quantized_utils) +make_jit_source(quantized kernels/quantized_utils.h) +make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h + kernels/fp4.h) + if(MLX_METAL_JIT) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jit_kernels.cpp) make_jit_source(arange) @@ -52,10 +65,6 @@ if(MLX_METAL_JIT) make_jit_source( reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h kernels/reduction/reduce_row.h kernels/reduction/reduce_init.h) - make_jit_source( - steel/gemm/gemm kernels/steel/utils.h kernels/steel/gemm/loader.h - kernels/steel/gemm/mma.h kernels/steel/gemm/params.h - kernels/steel/gemm/transforms.h) make_jit_source(steel/gemm/kernels/steel_gemm_fused) make_jit_source(steel/gemm/kernels/steel_gemm_masked kernels/steel/defines.h) make_jit_source(steel/gemm/kernels/steel_gemm_gather) @@ -76,10 +85,6 @@ if(MLX_METAL_JIT) make_jit_source(steel/conv/kernels/steel_conv_general kernels/steel/defines.h kernels/steel/conv/loaders/loader_general.h) - make_jit_source(quantized_utils) - make_jit_source(quantized kernels/quantized_utils.h) - make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h - kernels/fp4.h) make_jit_source(gemv) make_jit_source(gemv_masked) From 1fb504ca7162f2fe4ab5202dadb530a56c26b1d6 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 07:27:27 -0300 Subject: [PATCH 220/222] Make generated NAX optimizations source-reproducible --- mlx/backend/metal/kernels/fp_quantized_nax.h | 21 ++++++++++++------- .../steel/attn/kernels/steel_attention_nax.h | 12 ++++++++++- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 09179846b1..59f59047ce 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -966,6 +966,18 @@ template < dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { for (int k = 0; k < K_it; k++) { + // Load this immutable activation tile before staging the weights so + // its device reads overlap the two unchanged threadgroup barriers. + // The MMA traversal and accumulator chain below remain identical. + NAXTile Atile[BK / SK]; + STEEL_PRAGMA_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if constexpr (kAlignedM.value) { + Atile[kk1 / SK].load(xn + kk1, K); + } else { + Atile[kk1 / SK].load_rows(xn + kk1, K, sgp_sm); + } + } threadgroup_barrier(mem_flags::mem_threadgroup); if constexpr (kAlignedN.value) { loader_w.load_unsafe(); @@ -978,17 +990,10 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; NAXTile Btile; volatile int compiler_barrier; - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - if constexpr (transpose) { Btile.template load( Ws + tn * BK_padded + kk1); @@ -999,7 +1004,7 @@ template < tile_matmad_nax( Dtile, - Atile, + Atile[kk1 / SK], metal::bool_constant{}, Btile, metal::bool_constant{}); diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index b48a9a942d..058bc7288b 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -186,6 +186,16 @@ template < kb_min_causal = (q_min / BK); } + // K blocks ending at or below this simdgroup's first query row need no + // causal predicate: every select would keep the already-computed score. + // Restrict the tighter bound to the no-array-mask path so its proof rests + // only on the causal relation. + int sg_kb_min_causal = kb_min_causal; + if (do_causal && !has_mask) { + int sg_q_min = tid.x * BQ + params->qL_off + int(tm); + sg_kb_min_causal = max(0, sg_q_min + 1) / BK; + } + const bool is_last_bq = int(tid.x) == (params->NQ_aligned); // const bool is_last_tq = int(simd_group_id) >= (params->qL_rem / UQ); const bool is_last_q = is_last_bq; @@ -280,7 +290,7 @@ template < } // Mask out if causal - if (do_causal && kb >= kb_min_causal) { + if (do_causal && kb >= sg_kb_min_causal) { constexpr auto neg_inf = Limits::finite_min; const int base_row = tid.x * BQ + params->qL_off + tm; From 3f6e7728eb1921d0ae98b16ca207f1655d8ec266 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 07:44:47 -0300 Subject: [PATCH 221/222] Fix affine and fallback QMV template dispatch --- mlx/backend/metal/quantized.cpp | 36 +++++++++++++++++++++++---------- python/tests/test_quantized.py | 20 ++++++++++++++++++ 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 22a4ff0ae3..8acdb24e71 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -502,17 +502,31 @@ void qmv( use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", global_scale ? "_hgs" : ""); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value(), - results_per_simdgroup); + MTL::ComputePipelineState* kernel; + if (fast && mode != "affine") { + kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmv_fast", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value(), + results_per_simdgroup); + } else { + kernel = get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); + } auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 8ccda7f868..5165ce09a4 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -727,6 +727,26 @@ def test_qmv_fast_1bit_alignment(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmv_kernel_template_arity(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + + # Affine qmv_fast has no results-per-simdgroup template parameter. + x = mx.random.normal(shape=(1, 256), key=k1) + w = mx.random.normal(shape=(64, 256), key=k2) + w_q, scales, biases = mx.quantize(w, group_size=32, bits=8) + w_hat = mx.dequantize(w_q, scales, biases, group_size=32, bits=8) + y_q = mx.quantized_matmul(x, w_q, scales, biases, True, group_size=32, bits=8) + self.assertTrue(mx.allclose(y_q, x @ w_hat.T, atol=1e-3, rtol=1e-3)) + + # The non-fast floating-point qmv template does not have one either. + x = mx.random.normal(shape=(1, 96), key=k1) + w = mx.random.normal(shape=(64, 96), key=k2) + w_q, scales = mx.quantize(w, mode="nvfp4") + w_hat = mx.dequantize(w_q, scales, mode="nvfp4") + y_q = mx.quantized_matmul(x, w_q, scales, transpose=True, mode="nvfp4") + self.assertTrue(mx.allclose(y_q, x @ w_hat.T, atol=1e-3, rtol=1e-3)) + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From b57bd7640f3f7c743b76a58478faaf1e8ee084f2 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Fri, 14 Aug 2026 08:23:07 -0300 Subject: [PATCH 222/222] Fix Swift 6.0 Linux matmul compilation --- mlx/backend/cpu/matmul.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cpu/matmul.cpp b/mlx/backend/cpu/matmul.cpp index 12781dc1ab..c72e855ddf 100644 --- a/mlx/backend/cpu/matmul.cpp +++ b/mlx/backend/cpu/matmul.cpp @@ -90,8 +90,17 @@ void matmul_general( } }; - auto [a_transposed, lda, a] = check_transpose(a_pre); - auto [b_transposed, ldb, b] = check_transpose(b_pre); + // Keep these as ordinary local variables rather than structured bindings. + // Swift 6.0's Linux clang rejects captures of structured bindings here even + // when the package is compiled as C++20. + auto a_layout = check_transpose(a_pre); + auto b_layout = check_transpose(b_pre); + auto a_transposed = std::get<0>(a_layout); + auto lda = std::get<1>(a_layout); + auto a = std::get<2>(a_layout); + auto b_transposed = std::get<0>(b_layout); + auto ldb = std::get<1>(b_layout); + auto b = std::get<2>(b_layout); size_t M = a.shape(-2); size_t N = b.shape(-1); if (M == 0 || N == 0) {