diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 4f7b14105f..4a1a27cf9c 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -52,8 +52,10 @@ runs: shell: bash run: | echo "::group::Install macOS dependencies" + brew update brew install openmpi xcodebuild -showComponent MetalToolchain + sysctl -a | grep machdep.cpu echo "::endgroup::" - name: Setup Windows environment @@ -90,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. diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 186908f09c..7005396e71 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -20,6 +20,8 @@ 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. +- Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting. diff --git a/CMakeLists.txt b/CMakeLists.txt index f8ee65ad4b..9d7b5baa30 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}) @@ -261,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) @@ -296,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( @@ -417,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/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() { 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 68b4a5bd32..6f3697bb23 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,27 @@ 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_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(["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") @@ -282,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/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/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/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 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/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), 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 --------- 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/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 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 ------ diff --git a/docs/src/index.rst b/docs/src/index.rst index 46d069929f..8b3028be70 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -39,10 +39,13 @@ 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 usage/numpy + usage/precision + usage/environment_variables usage/distributed usage/using_streams usage/export 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 9d0d15f2a5..7743cd5d0f 100644 --- a/docs/src/python/devices_and_streams.rst +++ b/docs/src/python/devices_and_streams.rst @@ -10,11 +10,13 @@ Devices and Streams Device Stream + ThreadLocalStream default_device set_default_device default_stream new_stream new_thread_local_stream + new_thread_unsafe_stream set_default_stream stream synchronize 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/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 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/docs/src/usage/compile.rst b/docs/src/usage/compile.rst index bfe7f09661..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 @@ -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/distributed.rst b/docs/src/usage/distributed.rst index 1f271835cc..431aef7102 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -335,17 +335,48 @@ 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 ``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`` + 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 +^^^^^^^^^^^^^^^^^^^ + +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: @@ -479,10 +510,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 @@ -494,19 +525,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: @@ -523,10 +554,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/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/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/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. 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``. diff --git a/docs/src/usage/precision.rst b/docs/src/usage/precision.rst new file mode 100644 index 0000000000..440a1f705d --- /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 +:envvar:`MLX_ENABLE_TF32` to ``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. 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]) diff --git a/examples/cmake_project/CMakeLists.txt b/examples/cmake_project/CMakeLists.txt index 7a73ce1d9d..a955584782 100644 --- a/examples/cmake_project/CMakeLists.txt +++ b/examples/cmake_project/CMakeLists.txt @@ -2,13 +2,13 @@ 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 # 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/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/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", ) 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" + ) 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.pc.in b/mlx.pc.in index c4e2515d71..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) @@ -44,7 +63,6 @@ if (@MLX_BUILD_METAL@) endif() set_target_properties(mlx PROPERTIES - CXX_STANDARD 17 INTERFACE_COMPILE_OPTIONS "${MLX_CXX_FLAGS}" ) 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) { 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) { diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index b1e722b83f..7a2eef51fc 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" @@ -38,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."); @@ -197,6 +198,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( @@ -263,38 +283,26 @@ 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_); 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; } 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_); + }); }); } diff --git a/mlx/backend/cpu/binary.cpp b/mlx/backend/cpu/binary.cpp index 5c64f03129..9cca16d869 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); + } + }); }); } @@ -129,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()); diff --git a/mlx/backend/cpu/binary.h b/mlx/backend/cpu/binary.h index acaca50e1b..5333aff9a1 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); + }); }); } @@ -374,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); + }); }); } @@ -439,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); + }); }); } 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(), 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."); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/mlx/backend/cpu/copy.cpp b/mlx/backend/cpu/copy.cpp index d14736c15e..381719eb0c 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 @@ -248,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 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/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 d457c1fd99..ff5a4805a3 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,20 @@ 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. + 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, @@ -234,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."); } } diff --git a/mlx/backend/cpu/indexing.cpp b/mlx/backend/cpu/indexing.cpp index d668b56adb..15a2af3cb9 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) { @@ -281,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) { @@ -510,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_); + }); }); } @@ -675,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_); + }); }); } @@ -782,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); + }); }); } 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/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()); + }); }); } diff --git a/mlx/backend/cpu/matmul.cpp b/mlx/backend/cpu/matmul.cpp index 7df331671e..c72e855ddf 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 { @@ -89,32 +90,28 @@ 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) { 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)); } 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) { 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/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index c0f1a3c315..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) { @@ -1359,4 +1373,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/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/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/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); + }); }); } 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..d69e69ecf3 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; @@ -54,13 +56,6 @@ void store(T* dst, Simd x) { *(Simd*)dst = x; } -template -constexpr bool is_complex = false; - -template -constexpr bool is_complex().real())>> = - true; - template Simd rint(Simd in) { if constexpr (is_complex) { @@ -210,7 +205,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; } @@ -250,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/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/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 diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index 090cbb7a87..767e7ef1d3 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)) @@ -315,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) { @@ -330,36 +363,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_); + }); }); } @@ -401,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_); + }); }); } @@ -449,36 +428,38 @@ 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_); + }); + }); +} + +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); + } + }); }); } diff --git a/mlx/backend/cpu/svd.cpp b/mlx/backend/cpu/svd.cpp index ca01a0a65a..80101e9072 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" @@ -89,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()); } } @@ -179,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()); } } @@ -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/mlx/backend/cpu/unary.h b/mlx/backend/cpu/unary.h index 4fab6a7544..096521df89 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); + }); }); } @@ -146,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); + }); }); } 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) 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); 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 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/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/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( 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); diff --git a/mlx/backend/cuda/ptx.cuh b/mlx/backend/cuda/ptx.cuh new file mode 100644 index 0000000000..b7b87e671e --- /dev/null +++ b/mlx/backend/cuda/ptx.cuh @@ -0,0 +1,125 @@ +#pragma once + +#include + +namespace mlx::core { + +namespace ptx { + +#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); + 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 // defined(CUTE_ARCH_TMA_SM90_ENABLED) + +} // namespace ptx +} // namespace mlx::core 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/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..bea4e21979 --- /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(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); + // 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 // defined(CUTE_ARCH_TMA_SM90_ENABLED) +} + +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(CUTE_ARCH_TMA_SM90_ENABLED) + 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 // defined(CUTE_ARCH_TMA_SM90_ENABLED) +} + +// 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/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/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/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/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/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, diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 2a1a268c91..18326712ff 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_, @@ -102,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(); @@ -211,6 +221,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/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/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 365eb57b1c..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" @@ -51,7 +52,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 +75,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, @@ -148,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, @@ -210,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, @@ -318,6 +499,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 +582,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( @@ -402,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) { @@ -453,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 = @@ -481,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, 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 5eb8f543c7..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) @@ -48,13 +61,10 @@ 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) - 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) @@ -75,28 +85,37 @@ 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) 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( + 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(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/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/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 diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 49f87d9a5b..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( @@ -977,6 +1044,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 +1123,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/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index 0648ed221e..f2784a9c05 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,34 @@ 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(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; } - } - - auto lib = d.get_library( - name_, compile_options_, [this] { return metal::utils() + 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); diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 29eecf5562..252f72bd9a 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; @@ -306,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()); } @@ -392,8 +394,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()); @@ -515,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_), @@ -582,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(); @@ -882,9 +888,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); @@ -951,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(); 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/fft.cpp b/mlx/backend/metal/fft.cpp index 61eb02dac9..3c9195f8ea 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" @@ -26,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. @@ -127,6 +132,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 @@ -300,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. @@ -327,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()); @@ -633,18 +656,28 @@ 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 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 @@ -685,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, @@ -703,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); @@ -731,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/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/mlx/backend/metal/indexing.cpp b/mlx/backend/metal/indexing.cpp index a3df0e6e28..117e402f7b 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()); @@ -519,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/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 9ed9ce43e0..c7900ecdf8 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, @@ -399,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, @@ -803,6 +856,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, @@ -909,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 cb289fc1e7..21b754514c 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, @@ -253,6 +259,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, @@ -272,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/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 2bd99d2375..ecabbda550 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() @@ -48,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) @@ -140,6 +142,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/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..fe78e0050f 100644 --- a/mlx/backend/metal/kernels/complex.h +++ b/mlx/backend/metal/kernels/complex.h @@ -4,163 +4,286 @@ #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) : real(real), imag(imag) {}; - constexpr complex64_t() : 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)) {} + + template < + 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 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 to complex64_t + // Conversions between complex types 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_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) threadgroup : 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)) {} 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) device : 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) 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 from complex64_t + // Conversions to scalar types template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const thread { - return static_cast(real); + 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 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(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)) { @@ -171,3 +294,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/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/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/kernels/fft.h b/mlx/backend/metal/kernels/fft.h index 3cce29c574..1d847331d0 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 @@ -45,18 +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 -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 < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_butterfly( int i, int p, - thread float2* x, + int twiddle_offset, + thread vec* x, thread short* indices, - thread float2* 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. @@ -75,14 +104,19 @@ 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; + 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); 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 +130,23 @@ METAL_FUNC void radix_butterfly( // Perform all the radix steps required for a // particular radix size n. -template +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 float2* inputs, + thread vec* inputs, thread short* indices, - thread float2* values, - threadgroup float2* buf) { + thread vec* values, + 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. @@ -126,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); } } @@ -146,20 +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 -METAL_FUNC void -perform_fft(int fft_idx, thread int* p, int m, int n, threadgroup float2* buf) { - float2 inputs[MAX_RADIX]; +#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, + const device vec* twiddles = nullptr) { + vec inputs[MAX_RADIX]; short indices[MAX_OUTPUT_SIZE]; - float2 values[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_); @@ -184,7 +251,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 +276,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 +287,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 +296,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 +320,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 +345,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 +357,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 +379,31 @@ 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}; + // 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 = 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)]); } @@ -330,25 +412,24 @@ 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); // Rader IFFT on x[rader_m:] p = 1; - 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)}; + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); 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 - 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 +446,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 +455,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 +474,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 +498,50 @@ 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]; + 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); + } - float2 inv = float2(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); // 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); } @@ -448,6 +560,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 +570,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 +593,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.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 bd61eef6d7..9d742eb33c 100644 --- a/mlx/backend/metal/kernels/fft/radix.h +++ b/mlx/backend/metal/kernels/fft/radix.h @@ -16,49 +16,64 @@ 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, 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 {T(metal::fast::cos(theta)), T(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 +81,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 +109,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 +136,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 +167,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]; - w[0] = {0.8090169943749475, -0.5877852522924731}; - w[1] = {0.30901699437494745, -0.9510565162951535}; +template +METAL_FUNC void radix10( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[4]; + 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}; 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 +261,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 +274,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 +287,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 +303,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 +360,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 0dc62992e6..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. @@ -27,14 +28,48 @@ 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 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,23 +85,22 @@ 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_, 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,110 +108,139 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const { - 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 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 float2 post_in(float elem) const { - 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 { - return inv ? float2(elem.x / n, -elem.y / n) : elem; + METAL_FUNC complex_T pre_out(complex_T elem) const thread { + return pre_out(elem, n); } - METAL_FUNC float2 pre_out(float2 elem, int length) const { - return inv ? float2(elem.x / length, -elem.y / length) : elem; + METAL_FUNC complex_T pre_out(complex_T elem, int length) const thread { + 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 { + 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(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; - 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]) * scale; + } } - 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]); + buf[index] = post_in(in[batch_idx + index]) * scale; } } - METAL_FUNC void write() const { + 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; - 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) const { + 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) const { + 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}; + 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 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); } } } // 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 @@ -198,53 +261,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step - METAL_FUNC void load_strided(int stride, int overall_n) { - 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]); + // Four-step FFT I/O + METAL_FUNC void load_strided(int stride, int overall_n) thread { + 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; + 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++) { + 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; + } } } - METAL_FUNC void write_strided(int stride, int overall_n) { - 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); + METAL_FUNC void write_strided(int stride, int overall_n) thread { + if constexpr (step == 1 && !four_step_real) { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + 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++) { + 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) { - // 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) { - 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 @@ -253,14 +335,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 +362,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 = @@ -309,7 +391,7 @@ METAL_FUNC void ReadWriter::write() const { 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); } } } @@ -317,7 +399,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; @@ -334,7 +416,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; } @@ -344,7 +426,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; @@ -367,18 +449,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); } } } @@ -389,14 +471,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; @@ -425,17 +507,17 @@ METAL_FUNC void ReadWriter::load() const { 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; } } } 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 +538,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; @@ -486,12 +568,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); @@ -504,7 +586,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; @@ -519,7 +601,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; } @@ -531,7 +614,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 +629,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 +658,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 +683,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 +697,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..af4cbb02cf 100644 --- a/mlx/backend/metal/kernels/fp4.h +++ b/mlx/backend/metal/kernels/fp4.h @@ -1,7 +1,10 @@ #pragma once +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { - fp4_e2m1(float x) { + fp4_e2m1(float x) thread { if (metal::isnan(x)) { bits = 0x7; return; @@ -30,17 +33,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 5dd81b4f1e..fc4a580339 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 @@ -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 { @@ -145,6 +145,52 @@ 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, @@ -186,15 +232,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,19 +249,37 @@ struct QuantizedBlockLoader { scales_ + bi * src_ld / group_size + (bj * pack_factor) / group_size) {} - void load_unsafe() const { + 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 { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -234,14 +298,10 @@ 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() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -321,10 +381,16 @@ 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, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -334,7 +400,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; @@ -374,18 +439,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 +496,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 +551,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 +603,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 +732,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 +779,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 +833,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 +970,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 +1226,17 @@ 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, + int results_per_simdgroup = 4> [[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 +1268,34 @@ 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< + T, + group_size, + bits, + has_global_scale, + results_per_simdgroup>( + 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 +1327,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 +1385,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 +1426,10 @@ template s_strides, tid); } - fp_qvm_impl( + fp_qvm_impl( w, scales, + global_scale, x, y, in_vec_size, @@ -1341,9 +1482,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 +1553,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 +1608,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 +1654,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 +1713,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 +1772,10 @@ template w_strides, s_strides, tid); - fp_qvm_impl( + fp_qvm_impl( w, scales, + global_scale, x, y, in_vec_size, @@ -1741,11 +1906,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,38 +2162,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]]) { + 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(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_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; @@ -2038,10 +2214,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]]) { @@ -2053,9 +2230,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) @@ -2070,31 +2255,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]]) { + 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(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_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..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" @@ -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( \ @@ -42,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, \ @@ -137,25 +176,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 +206,17 @@ 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, 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/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 0ea6af168e..59f59047ce 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 { @@ -69,6 +69,44 @@ 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, @@ -112,14 +150,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,23 +165,46 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + 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 { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -162,23 +223,15 @@ 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() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { scales += n_groups; } else { - scales += n_groups * group_stride; + scales += group_stride; } } }; @@ -913,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(); @@ -925,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); @@ -946,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/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/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/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..c7e15ef6b1 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 == 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); + } + } - if (bits == 2) { + 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 == 2) { + 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; + } + } + + 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(); @@ -610,15 +704,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 +720,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 +733,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 +763,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -759,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; + 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(); @@ -1591,7 +1686,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 +1743,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 +1859,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 +2111,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 +2170,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 +2232,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 +2294,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 +2441,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 +2704,7 @@ template < } } -template +template [[kernel]] void affine_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], @@ -2636,14 +2748,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; @@ -2697,7 +2818,7 @@ template } } -template +template [[kernel]] void affine_dequantize( const device uint8_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2747,7 +2868,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 d50af8679a..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 == 2) { + 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; + } + } + + 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(); @@ -612,15 +714,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 +730,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 +743,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 +773,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -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(); @@ -752,14 +854,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 +870,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 +883,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 +913,7 @@ struct QuantizedBlockLoader< } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { // if (group_steps > 1) { @@ -826,8 +928,8 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; + scales += group_stride; + biases += group_stride; } } }; @@ -1094,7 +1196,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 +1216,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 +1242,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 +1268,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 +1293,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 +1787,4 @@ template < }); }); } -} \ No newline at end of file +} diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 5a9c9fb874..2cc1aeea84 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) \ @@ -96,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/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/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h index 11d8e83ac6..f875e51c92 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,29 @@ 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); } + 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) { + U operator()(U a, U b) thread { return a + b; } }; @@ -142,19 +153,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 +176,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 +193,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 +214,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 +236,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 +253,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 +274,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..2e5493961b 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); } @@ -99,20 +99,29 @@ template struct CumMax { static constexpr constant U init = Limits::min; - template - U operator()(U a, T b) { + 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; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + 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; } - 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); } @@ -122,20 +131,29 @@ template struct CumMin { static constexpr constant U init = Limits::max; - template - U operator()(U a, T b) { + 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; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + 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; } - 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 +164,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 +176,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/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/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.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.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index adc9a42798..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 @@ -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; } }; @@ -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; @@ -207,7 +217,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; @@ -276,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; 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/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/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)); + } } }); }); 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..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()); }; /////////////////////////////////////////////////////////////////////////////// @@ -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; } }; 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=() diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 87d2bf52d1..2e5fc06b33 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" @@ -14,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" @@ -371,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; @@ -435,7 +471,8 @@ 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, @@ -1039,6 +1076,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, @@ -1230,6 +1315,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 +1539,43 @@ 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 && + 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 +1708,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( @@ -1897,7 +2209,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; @@ -2030,7 +2342,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 @@ -2260,6 +2572,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 +2889,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..5da78db8a3 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, @@ -273,6 +281,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, @@ -323,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/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); diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 62d48714e0..8acdb24e71 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1,5 +1,6 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 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" @@ -84,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) { @@ -125,6 +142,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, @@ -134,7 +157,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 +195,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 +463,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, @@ -250,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 % 512 == 0; + 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, @@ -266,26 +499,46 @@ void qmv( group_size, "_b_", bits, - B > 1 ? "_batch_1" : "_batch_0"); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1); + use_narrow_qmv ? "_r_2" : "", + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); + 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); - 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++); @@ -296,8 +549,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): @@ -511,6 +774,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 +803,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++); @@ -784,7 +1059,9 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { - if (metal::is_nax_available() && transpose && (K % 64 == 0) && + 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, @@ -884,11 +1161,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); + // 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 * group_size - while (split_k > 1 && (K % (split_k * group_size) != 0)) { + // 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) { @@ -1054,6 +1335,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, @@ -1075,7 +1357,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_"), @@ -1083,7 +1365,8 @@ void gather_qmv( "_gs_", group_size, "_b_", - bits); + bits, + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( d, @@ -1092,17 +1375,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++); @@ -1120,6 +1406,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, @@ -1149,18 +1436,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++); @@ -1213,9 +1510,15 @@ 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; + // 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; @@ -1291,9 +1594,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++); @@ -1364,6 +1666,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; @@ -1441,9 +1747,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++); @@ -1459,6 +1764,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, @@ -1469,18 +1775,32 @@ 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 && 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; } // 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, bits, M, 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) { @@ -1536,13 +1856,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; } @@ -1628,6 +1974,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1647,6 +1994,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1660,194 +2008,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/metal/reduce.cpp b/mlx/backend/metal/reduce.cpp index 562f966297..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) { @@ -292,6 +307,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 +975,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/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/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index d6f74bc2a9..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 @@ -475,7 +476,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]; @@ -624,7 +627,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/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); 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/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 ae51dd9b2f..e522307156 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) @@ -106,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/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() {} diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 4819ed2724..b7d7a19467 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) @@ -133,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/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/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..99a4ac3be3 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 { @@ -145,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 { @@ -170,4 +182,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/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/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..5b749f2bd9 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; @@ -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 2ced24416f..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 { @@ -23,7 +22,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_; @@ -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/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index 7dc23043d2..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()); } } @@ -181,18 +188,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; @@ -276,6 +295,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,16 +305,26 @@ 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); 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 +359,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/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 541e825b4b..a08035f3f3 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -11,13 +11,14 @@ 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)) { + 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() { @@ -166,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_; @@ -204,17 +217,42 @@ 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_) { + + if (n_bytes <= 32768 || 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); + 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_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, n_conns_, reduce_op); + 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 016c6af645..c901c15a85 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_; @@ -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 @@ -78,6 +88,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..15b81a2177 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,299 +84,435 @@ 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) { + // 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 wire_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; + 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; 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])); + } + + // 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, + wire_end, + send_offset, + recv_offset, + in_ptr, + out_ptr, + ReduceRecvOp{reduce_op}); } 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])); + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + wire_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, + wire_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); + } - // 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); + void + all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes, int n_wires) { + // Copy our data to the appropriate place + std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } + // 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); + dispatch_wires(n_wires, [&](int lw) { + all_gather_wire( + out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + }); + } - // 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; + // 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}; + // 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, + ((rank_ + 1) % size_) * n_bytes}; + + ring_pass<2, char>( + lw, + n_bytes, + n_bytes, + n_bytes * size_, + n_bytes_per_wire, + wire_offset, + wire_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); + } - in_flight--; + // 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); + }); + } - 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]++; - } + // 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 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 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; + 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); + } - 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++; - } - } + // 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, wire_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++; } } + }; - 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; + 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; } - } - // 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]++; - } + 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, wire_end[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); } - post_send_all(sz, buff, n_wires); buff++; - in_flight += 2 * MAX_DIR * n_wires; + in_flight += 2 * MAX_DIR; } // 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); + 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 * 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]++; + 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, wire_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) { - 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++; - } + // 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); } } } - 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; + // 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; } } } - void - all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes, int n_wires) { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - + // 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, 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, + int64_t size, + int64_t stride, + int64_t total, + int64_t size_per_wire, + const int64_t (&wire_offset)[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, + T* out_ptr, + RecvOp recv_op) { constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * 2; - 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; + 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( + wire_end[lr], std::max(0, size - send_offset[lr])); + recv_limits[lr] = std::min( + wire_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 +520,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 +558,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 +578,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 +631,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 +657,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 +705,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 +715,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 +730,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 +749,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 +808,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 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/mlx/distributed/ops.cpp b/mlx/distributed/ops.cpp index 1762f0e6bc..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( @@ -165,11 +149,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()); } diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index ea40042844..3e0c2a3221 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -21,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; @@ -238,6 +180,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 +197,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 +207,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; } } @@ -475,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 { @@ -525,7 +490,7 @@ class RingGroup : public GroupImpl { (i % 2) ? -1 : 1))); } for (auto& f : all_gathers) { - f.wait(); + f.get(); } }); } @@ -650,7 +615,7 @@ class RingGroup : public GroupImpl { reduce_op))); } for (auto& f : all_sums) { - f.wait(); + f.get(); } }); } @@ -737,8 +702,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 +714,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 +754,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 +774,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 +792,7 @@ class RingGroup : public GroupImpl { std::min(data_size, (i + 1) * segment_size) - i * segment_size)); } for (auto& f : recvs) { - f.wait(); + f.get(); } } 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); 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/mlx/export.cpp b/mlx/export.cpp index 4777fd3048..bcc0ad161c 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); @@ -414,6 +430,7 @@ struct PrimitiveFactory { "CumLogaddexp"), SERIALIZE_PRIMITIVE(Scatter), SERIALIZE_PRIMITIVE(ScatterAxis), + SERIALIZE_PRIMITIVE(SearchSorted), SERIALIZE_PRIMITIVE(Select), SERIALIZE_PRIMITIVE(Sigmoid), SERIALIZE_PRIMITIVE(Sign), @@ -521,10 +538,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 +696,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 +843,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 +922,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 +974,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 +1086,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/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/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..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()); } } @@ -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/mlx/linalg.cpp b/mlx/linalg.cpp index bdf05fb80b..0d06b56469 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); @@ -172,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) { @@ -368,6 +369,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); @@ -437,12 +447,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 849b8081dd..0dd10f3d24 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -16,15 +16,21 @@ #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 { 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 +52,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]); @@ -65,11 +77,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); @@ -245,6 +272,7 @@ array linspace( double start, double stop, int num /* = 50 */, + bool endpoint /* = true */, Dtype dtype /* = float32 */, StreamOrDevice s /* = {} */) { if (num < 0) { @@ -256,8 +284,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), @@ -267,12 +298,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 @@ -283,12 +312,20 @@ 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, 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( @@ -348,7 +385,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 /* = {} */) { @@ -356,14 +397,23 @@ 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 /* = {} */) { - 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; } @@ -542,26 +592,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), @@ -577,19 +651,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 /* = {} */) { @@ -602,21 +668,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), @@ -626,7 +682,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( @@ -636,23 +692,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( @@ -1427,6 +1467,91 @@ 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; + } + 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; + 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, @@ -1519,6 +1644,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"; @@ -1674,6 +1805,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; } @@ -1919,6 +2055,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); } @@ -1973,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 << "."; @@ -2397,7 +2536,7 @@ array prod( out_type = int32; } auto out = (is_noop) - ? a + ? astype(a, out_type, s) : array( std::move(out_shape), out_type, @@ -2428,11 +2567,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( @@ -2465,14 +2601,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( @@ -2768,6 +2901,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(); @@ -2844,7 +3008,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; } @@ -3081,7 +3245,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}); } @@ -3352,9 +3516,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); } @@ -3366,7 +3535,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); @@ -3583,12 +3753,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; @@ -3597,9 +3763,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); @@ -3621,12 +3784,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; @@ -3651,9 +3810,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); @@ -3761,7 +3917,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."; @@ -4318,16 +4475,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(); @@ -4478,7 +4640,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; } @@ -4626,10 +4789,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 */, @@ -4678,13 +4841,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) { @@ -4738,9 +4901,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) { @@ -4768,24 +4931,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 @@ -4893,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()); } @@ -4917,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); @@ -5088,12 +5252,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); @@ -5353,13 +5511,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) { @@ -5452,30 +5603,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)); @@ -5511,6 +5643,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, @@ -6009,28 +6191,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); @@ -6165,6 +6333,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/mlx/ops.h b/mlx/ops.h index ed617c3441..f597753b1e 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -38,20 +38,28 @@ 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, 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( @@ -93,14 +101,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. */ @@ -853,6 +865,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 = {}); @@ -1546,10 +1565,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 +1597,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 +1629,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 6ac88eef6c..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( @@ -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()); @@ -2481,9 +2497,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 +2515,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 +2563,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 +2584,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())}; } @@ -3372,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, @@ -3387,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}; } @@ -3399,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 { @@ -3825,6 +3847,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) { @@ -4356,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_) { @@ -4483,7 +4525,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 +4638,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 +4652,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 +4761,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; @@ -5131,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( @@ -5145,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: { @@ -5376,6 +5425,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) { @@ -5394,7 +5484,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, @@ -5409,7 +5501,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}; } @@ -5748,7 +5840,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 +6038,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 +6104,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/mlx/primitives.h b/mlx/primitives.h index 5b8517c56d..3a3d0ba5e5 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) @@ -2176,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/mlx/random.cpp b/mlx/random.cpp index def3169cb5..2932b3eaf8 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; @@ -182,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); @@ -280,7 +284,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( @@ -379,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)); @@ -449,8 +488,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); 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/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/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/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..0c62ba30d9 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_); @@ -187,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/python/mlx/_distributed_utils/common.py b/python/mlx/_distributed_utils/common.py index d57747a320..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 @@ -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) diff --git a/python/mlx/_distributed_utils/config.py b/python/mlx/_distributed_utils/config.py index d9262a0191..a4dabed629 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 @@ -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 6838897da4..4771c1fb5b 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 @@ -119,11 +121,12 @@ def make_launch_script(rank, 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 @@ -146,6 +149,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[@]}"' @@ -326,7 +333,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 @@ -365,13 +372,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 @@ -514,13 +521,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: @@ -529,10 +538,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: diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt index c05f879d13..974ce0c7a5 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 | complex + 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/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 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/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 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/activations.py b/python/mlx/nn/layers/activations.py index 4f36ae2b38..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) @@ -225,7 +229,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 +239,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 +611,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 +621,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): diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 9802c6c311..e980f8343b 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 @@ -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 @@ -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/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index c4799bd073..16979b5097 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:`~mlx.nn.layers.distributed.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/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): 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/normalization.py b/python/mlx/nn/layers/normalization.py index 56bcc3ccba..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: @@ -336,13 +337,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 +352,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 +371,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/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 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/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 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__( 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/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/mlx/nn/losses.py b/python/mlx/nn/losses.py index 61b6adb047..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: @@ -416,8 +421,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 @@ -504,8 +510,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: @@ -540,8 +545,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 +554,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/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/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index f1ef46773b..65efab222d 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() @@ -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`` """ @@ -967,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/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) diff --git a/python/mlx/utils.py b/python/mlx/utils.py index bfd167c796..f8595ef9a9 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()} @@ -286,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. @@ -297,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/src/array.cpp b/python/src/array.cpp index 1fcbbd25f0..fbab846182 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, @@ -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(), @@ -539,6 +536,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__", @@ -1025,6 +1029,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/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/src/convert.cpp b/python/src/convert.cpp index 1b8145094f..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, @@ -172,6 +177,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 +262,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."); } } @@ -313,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); } @@ -667,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); 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 9f4a7cb59e..0ad91c6e48 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: Callable[[int, int], Callable[[bytes, int], bytes]] | None = 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. @@ -116,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. @@ -145,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. @@ -174,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. @@ -203,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. @@ -236,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. @@ -264,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``. @@ -297,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``. @@ -330,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. @@ -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/src/export.cpp b/python/src/export.cpp index 51c8a8965c..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_; }; @@ -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: str | Callable, fun: Callable, *args, shapeless: bool = False, metadata: str | None = 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) -> 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_( @@ -249,7 +277,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; }) @@ -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/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/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/linalg.cpp b/python/src/linalg.cpp index bfebe3bb47..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(L: 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. @@ -347,7 +346,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`` @@ -362,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. @@ -389,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. @@ -448,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. @@ -526,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. @@ -572,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``. @@ -603,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. @@ -623,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``. @@ -645,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``. @@ -667,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. @@ -700,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 19b19bd78f..a680ef4961 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. @@ -1473,21 +1472,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 +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[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. @@ -1508,7 +1511,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``. @@ -1544,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, @@ -1595,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. @@ -1620,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. @@ -1642,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: Optional[int] = 50, dtype: Optional[Dtype] = float32, stream: Union[None, Stream, Device] = 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]``. @@ -1665,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``. @@ -1679,14 +1687,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: @@ -1723,7 +1731,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. @@ -1760,7 +1768,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. @@ -1799,7 +1807,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. @@ -1831,7 +1839,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. @@ -1862,7 +1870,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. @@ -1888,7 +1896,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. @@ -1910,8 +1918,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. @@ -1938,7 +1946,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. @@ -1955,17 +1963,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: Dtype | None = None, *, stream: StreamOrDevice = 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. @@ -1983,7 +1998,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. @@ -1997,17 +2012,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: Dtype | None = None, *, stream: StreamOrDevice = 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. @@ -2028,7 +2050,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. @@ -2052,7 +2074,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. @@ -2080,7 +2102,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. @@ -2102,7 +2124,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. @@ -2122,7 +2144,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. @@ -2145,7 +2167,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. @@ -2182,7 +2204,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. @@ -2223,7 +2245,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. @@ -2252,7 +2274,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. @@ -2280,7 +2302,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. @@ -2307,7 +2329,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. @@ -2330,7 +2352,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. @@ -2349,7 +2371,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. @@ -2367,8 +2389,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. @@ -2386,8 +2407,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. @@ -2405,8 +2425,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. @@ -2426,14 +2445,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. @@ -2446,14 +2464,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. @@ -2467,7 +2484,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. @@ -2488,7 +2505,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. @@ -2516,7 +2533,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. @@ -2544,7 +2561,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"); @@ -2562,7 +2579,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. @@ -2598,7 +2615,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. @@ -2626,7 +2643,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. @@ -2655,7 +2672,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. @@ -2684,7 +2701,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. @@ -2720,7 +2737,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. @@ -2750,7 +2767,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. @@ -2785,7 +2802,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. @@ -2814,7 +2831,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. @@ -2845,7 +2862,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. @@ -2878,7 +2895,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. @@ -2914,7 +2931,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. @@ -2958,7 +2975,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. @@ -2990,7 +3007,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. @@ -3018,7 +3035,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. @@ -3048,7 +3065,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. @@ -3082,7 +3099,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. @@ -3120,7 +3137,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. @@ -3141,6 +3158,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, @@ -3159,7 +3210,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. @@ -3185,7 +3236,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. @@ -3207,7 +3258,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. @@ -3233,7 +3284,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, *, precise: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the softmax along the given axis. @@ -3248,6 +3299,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. @@ -3268,7 +3323,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. @@ -3296,7 +3351,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"); @@ -3316,7 +3371,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. @@ -3337,14 +3392,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: 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) -> tuple[array, ...]"), R"pbdoc( Generate multidimensional coordinate grids from 1-D coordinate arrays @@ -3357,7 +3412,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", @@ -3377,7 +3432,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. @@ -3414,7 +3469,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. @@ -3442,18 +3497,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); } } }, @@ -3464,7 +3524,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', 'reflect', 'symmetric'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Pad an array with a constant value @@ -3479,7 +3539,9 @@ 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 + "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. Returns: @@ -3511,7 +3573,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. @@ -3548,7 +3610,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. @@ -3580,7 +3642,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. @@ -3618,7 +3680,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. @@ -3655,7 +3717,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. @@ -3691,7 +3753,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. @@ -3716,7 +3778,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. @@ -3737,8 +3799,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`. @@ -3758,7 +3819,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`. @@ -3832,7 +3893,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. @@ -3859,7 +3920,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 @@ -3917,7 +3978,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 @@ -3987,7 +4048,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 @@ -4023,7 +4084,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 @@ -4098,7 +4159,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 @@ -4184,7 +4245,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 @@ -4286,7 +4347,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 @@ -4319,8 +4380,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. @@ -4336,8 +4396,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. @@ -4371,7 +4430,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. @@ -4390,7 +4449,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. @@ -4425,7 +4484,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. @@ -4447,7 +4506,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. @@ -4478,7 +4537,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``. @@ -4510,7 +4569,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. @@ -4537,7 +4596,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. @@ -4570,7 +4629,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 @@ -4608,7 +4667,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``. @@ -4648,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. @@ -4709,7 +4768,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. @@ -4724,13 +4783,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`` @@ -4764,7 +4823,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. @@ -4801,6 +4860,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: 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. + + 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, @@ -4810,7 +4921,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. @@ -4846,7 +4957,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. @@ -4870,7 +4981,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. @@ -4890,7 +5001,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. @@ -4910,7 +5021,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. @@ -4937,7 +5048,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``. @@ -4959,7 +5070,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. @@ -4989,7 +5100,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. @@ -5028,7 +5139,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. @@ -5070,7 +5181,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. @@ -5102,7 +5213,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 @@ -5138,7 +5249,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. @@ -5168,13 +5279,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. @@ -5191,13 +5302,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. @@ -5214,13 +5325,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. @@ -5250,7 +5361,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. @@ -5322,8 +5433,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. @@ -5350,7 +5460,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. @@ -5411,7 +5521,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. @@ -5439,7 +5549,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. @@ -5466,7 +5576,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. @@ -5493,7 +5603,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. @@ -5521,7 +5631,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. @@ -5549,7 +5659,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. @@ -5574,7 +5684,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. @@ -5596,7 +5706,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. @@ -5622,7 +5732,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. @@ -5686,7 +5796,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. @@ -5721,7 +5831,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. @@ -5749,7 +5859,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. @@ -5768,7 +5878,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. @@ -5795,7 +5905,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. @@ -5834,7 +5944,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. @@ -5863,7 +5973,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. @@ -5892,7 +6002,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. @@ -5943,7 +6053,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``. @@ -5970,7 +6080,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 @@ -6020,7 +6130,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. @@ -6038,7 +6148,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/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) diff --git a/python/src/random.cpp b/python/src/random.cpp index ddddb56eed..8485faea41 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, @@ -109,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. @@ -144,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. @@ -187,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. @@ -227,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. @@ -274,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. @@ -282,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. @@ -311,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. @@ -355,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. @@ -389,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. @@ -435,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. @@ -478,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. @@ -511,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/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/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``. 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/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/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 00d3182bc0..cbb9663046 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,51 @@ 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 + ) + ) + + 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") 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__": diff --git a/python/tests/test_array.py b/python/tests/test_array.py index dfab616716..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 @@ -1205,7 +1214,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]]) @@ -1548,6 +1557,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)) @@ -1868,11 +1908,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) @@ -2129,16 +2168,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 +2263,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 +2693,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) @@ -2663,16 +2713,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) diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index fa6aed6f09..c85fcf09a6 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -444,6 +444,98 @@ 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_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) @@ -1416,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() diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 0db6257f93..7dbbaea049 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -538,6 +538,125 @@ 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") + + # 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 +852,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 +1367,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) @@ -1239,6 +1474,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] 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): 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)) 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() diff --git a/python/tests/test_double.py b/python/tests/test_double.py index f957e8fc96..3186e7e573 100644 --- a/python/tests/test_double.py +++ b/python/tests/test_double.py @@ -296,11 +296,55 @@ 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) + 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_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() 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() diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index c45dd4a606..23f0645b3a 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 @@ -313,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") @@ -704,6 +723,86 @@ 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)) + + 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/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): diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 7bd867084e..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 @@ -360,6 +387,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 diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 0bb190bd99..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) @@ -185,6 +208,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 diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index afdf75d7c8..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,10 +65,11 @@ 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)) - # 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 +79,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) @@ -154,6 +157,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) @@ -188,6 +201,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( @@ -293,6 +323,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) @@ -435,6 +472,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 @@ -493,6 +543,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 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() 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 diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 8ad72a323c..c5e6db94a7 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): @@ -192,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()) @@ -672,7 +696,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 +707,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 +764,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 +804,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 @@ -989,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) @@ -1127,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) @@ -2193,6 +2287,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() diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 028d9a4573..2b6f9324f5 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 @@ -454,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) @@ -474,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]: @@ -724,6 +768,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( [ @@ -783,6 +840,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) @@ -813,6 +873,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( [ @@ -1091,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] @@ -1328,6 +1436,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) @@ -1488,6 +1622,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) @@ -1994,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) @@ -2124,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)) @@ -2134,6 +2277,60 @@ 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 + # 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) + + # 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) @@ -2152,6 +2349,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 @@ -2293,6 +2493,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]) @@ -2321,6 +2557,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) @@ -2489,6 +2743,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", @@ -2500,8 +2858,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 @@ -2598,7 +2960,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) @@ -2625,6 +2987,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]]]) @@ -2833,6 +3246,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)) @@ -3185,6 +3619,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) @@ -3373,8 +3834,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]) @@ -3515,6 +3981,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/python/tests/test_optimizers.py b/python/tests/test_optimizers.py index 68a88d1820..3d51fedf32 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): @@ -610,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_quantized.py b/python/tests/test_quantized.py index f031d2aa92..5165ce09a4 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,5 +1,8 @@ -# Copyright © 2023 Apple Inc. +# Copyright © 2023-2026 Apple Inc. +import os +import platform +import subprocess import unittest from itertools import product @@ -7,11 +10,19 @@ 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)) 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) @@ -22,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)) @@ -154,6 +165,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") @@ -161,11 +181,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( @@ -173,12 +189,171 @@ 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_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_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) tests = product( [256, 512, 67], # M - [64, 256], # N + [64, 256, 512], # N ["nvfp4", "mxfp8"], # mode ) for M, N, mode in tests: @@ -186,12 +361,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 @@ -224,31 +395,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) @@ -256,7 +450,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 @@ -319,6 +513,92 @@ 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 + # 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) + + 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) + + 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) k1, k2 = mx.random.split(key) @@ -409,7 +689,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 @@ -431,6 +711,42 @@ 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_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) @@ -482,6 +798,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 @@ -498,7 +849,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): @@ -532,7 +883,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( @@ -1072,6 +1428,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) @@ -1157,6 +1607,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": @@ -1189,6 +1643,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"), @@ -1271,6 +1727,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: diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 551c32993c..96a789b608 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]: @@ -230,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 ) @@ -325,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/python/tests/test_reduce.py b/python/tests/test_reduce.py index 24d5a688db..6ac8fc1504 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) @@ -218,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) diff --git a/python/tests/test_tree.py b/python/tests/test_tree.py index 171c15b12b..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 @@ -121,6 +155,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() 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() 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) 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() 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/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 { 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 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; 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); +} 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 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") { 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; diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 987ee4df9b..3da0a2950b 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") { @@ -2960,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") { @@ -3140,6 +3201,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") { @@ -3282,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") { @@ -4337,6 +4419,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); 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") { 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); +}