diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index 2fc81d589..1262ded04 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.github/scripts/build-linux-cuda-cmake.sh @@ -17,13 +17,6 @@ mkdir -p build pushd build cmake -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=ON -DMLX_C_BUILD_EXAMPLES=OFF .. -G Ninja -# Verify the local mlx patch is actually in place: without it the examples below -# crash on exit, which is easy to misread as an unrelated CUDA failure. -git -C _deps/mlx-src apply --reverse --check "$PWD/../cmake/mlx.patch" || { - echo "error: cmake/mlx.patch is not applied to _deps/mlx-src" - exit 1 -} - ninja ./example1 --device gpu ./tutorial --device gpu diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c558a90c..3ddeeb7c5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -37,22 +37,10 @@ endif() set(apply_patch ${CMAKE_COMMAND} -DREPO=) set(apply_patch_script -P ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake) -# two local patches for CUDA: -# -# * https://github.com/ml-explore/mlx/pull/4480: leak the global command encoder -# map rather than synchronizing on process shutdown -# * make ~CudaHandle non-throwing: cu::Worker is detached and holds the last -# reference to itself, so ~Worker -> ~CudaEvent -> cudaEventDestroy runs on -# the worker thread at an arbitrary point during exit, racing the CUDA -# runtime's own teardown. A failing destroy there terminated the process. -set(mlx_patch - ${apply_patch} -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/mlx.patch - ${apply_patch_script}) - # note: this must be declared (and named `mlx`) before mlx-c is made available -# so that our pinned/patched version wins over the one mlx-c declares. Because -# we now pull mlx in first, we are also responsible for the option defaults that -# mlx-c would normally set for us. +# so that our pinned version wins over the one mlx-c declares. Because we now +# pull mlx in first, we are also responsible for the option defaults that mlx-c +# would normally set for us. set(MLX_BUILD_TESTS OFF) set(MLX_BUILD_EXAMPLES OFF) set(MLX_BUILD_BENCHMARKS OFF) @@ -61,8 +49,7 @@ set(MLX_BUILD_PYTHON_BINDINGS OFF) FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG "v0.32.2" - PATCH_COMMAND ${mlx_patch}) + GIT_TAG "v0.32.4") FetchContent_MakeAvailable(mlx) # mlx-c @@ -70,7 +57,7 @@ set(MLX_C_BUILD_EXAMPLES OFF) FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" - GIT_TAG "ebc88f10caa1b625e6b581437a8dea6df8a70085") + GIT_TAG "v0.8.0") FetchContent_MakeAvailable(mlx-c) # swift-numerics diff --git a/Package.swift b/Package.swift index fc70166fb..80f9d2829 100644 --- a/Package.swift +++ b/Package.swift @@ -21,6 +21,7 @@ let noMetalCmlxExcludes = [ "mlx/mlx/backend/metal/compiled.cpp", "mlx/mlx/backend/metal/conv.cpp", "mlx/mlx/backend/metal/copy.cpp", + "mlx/mlx/backend/metal/cross_entropy.cpp", "mlx/mlx/backend/metal/custom_kernel.cpp", "mlx/mlx/backend/metal/device.cpp", "mlx/mlx/backend/metal/device_info.cpp", @@ -29,6 +30,7 @@ let noMetalCmlxExcludes = [ "mlx/mlx/backend/metal/event.cpp", "mlx/mlx/backend/metal/fence.cpp", "mlx/mlx/backend/metal/fft.cpp", + "mlx/mlx/backend/metal/gated_delta_update.cpp", "mlx/mlx/backend/metal/hadamard.cpp", "mlx/mlx/backend/metal/indexing.cpp", "mlx/mlx/backend/metal/jit_kernels.cpp", @@ -245,6 +247,7 @@ let cmlx = Target.target( .headerSearchPath("json/single_include/nlohmann"), .headerSearchPath("fmt/include"), .define("MLX_VERSION", to: "\"0.32.0\""), + .unsafeFlags(["-Wno-shorten-64-to-32"]), ], linkerSettings: linkerSettings, plugins: cudaBuildPlugins, diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift index 15ecd2e88..f28c552c7 100644 --- a/Plugins/CudaBuild/plugin.swift +++ b/Plugins/CudaBuild/plugin.swift @@ -16,6 +16,7 @@ struct CudaBuild: BuildToolPlugin { var cppLanguageStandard: String? = nil var verbose: Bool = false var codeGeneration: [CodeGeneration] = [] + var hostCompiler: String? = nil init() {} @@ -28,9 +29,29 @@ struct CudaBuild: BuildToolPlugin { verbose = try c.decodeIfPresent(Bool.self, forKey: .verbose) ?? false codeGeneration = try c.decodeIfPresent([CodeGeneration].self, forKey: .codeGeneration) ?? [] + hostCompiler = try c.decodeIfPresent(String.self, forKey: .hostCompiler) } } + /// Host compiler nvcc should use, in priority order: the + /// `MLX_CUDA_HOST_COMPILER` environment variable, the `hostCompiler` key in + /// CudaBuild.json, then the toolchain's own clang++. + /// + /// It must still be a clang: nvcc preprocesses with it and SwiftPM compiles + /// the generated `.cpp`, so a GCC-preprocessed translation unit will not build. + func resolveHostCompiler(context: PluginContext, settings: Settings) -> String { + if let env = ProcessInfo.processInfo.environment["MLX_CUDA_HOST_COMPILER"], !env.isEmpty { + return env + } + if let configured = settings.hostCompiler, !configured.isEmpty { + return configured + } + guard let clangUrl = try? context.tool(named: "clang++") else { + fatalError("clang++ not found") + } + return clangUrl.url.path + } + func createBuildCommands(context: PluginContext, target: Target) async throws -> [Command] { print("CUDA Build Plugin") @@ -40,12 +61,6 @@ struct CudaBuild: BuildToolPlugin { return [] } - guard let clangUrl = try? context.tool(named: "clang++") else { - fatalError("clang++ not found") - } - - print("Use clang++ at: \(clangUrl.url.path)") - let sourceDir = target.directoryURL print("Source directory: \(sourceDir.path)") @@ -61,6 +76,10 @@ struct CudaBuild: BuildToolPlugin { settings = Settings() } + let hostCompiler = resolveHostCompiler(context: context, settings: settings) + + print("Use clang++ at: \(hostCompiler)") + // Scan source directory for .cu files let sourceDirPath = sourceDir.path.hasSuffix("/") ? sourceDir.path : sourceDir.path + "/" @@ -132,6 +151,22 @@ struct CudaBuild: BuildToolPlugin { } let stdArgs = settings.cppLanguageStandard.map { ["--std", $0] } ?? [] + // Include roots that live outside the package: the CCCL copy to compile + // device code against, plus anything named in MLX_CUDA_INCLUDE_PATHS + // (cudnn-frontend, CUTLASS). Kept in step with Package.swift, which + // applies the same roots to the C++ sources. + let environment = ProcessInfo.processInfo.environment + var externalIncludeRoots: [String] = [] + if let ccclDir = environment["MLX_CCCL_DIR"] { + externalIncludeRoots.append(ccclDir) + } + let extraRoots = environment["MLX_CUDA_INCLUDE_PATHS"] ?? "" + externalIncludeRoots += extraRoots.split(separator: ":").map(String.init) + let externalIncludeArgs: [String] = + externalIncludeRoots + .filter { !$0.isEmpty } + .flatMap { ["-I", $0] } + for inputFile in sourceCuFiles + generatedCuFiles { let outputCpp = URL(string: inputFile.relativePath, relativeTo: outputDir)! .deletingPathExtension().appendingPathExtension("cpp") @@ -141,12 +176,13 @@ struct CudaBuild: BuildToolPlugin { "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)", executable: encuda.url, arguments: compileArguments( - clangpp: clangUrl.url.path, + clangpp: hostCompiler, sourceDir: sourceDir.path, verboseFlag: verboseFlag, stdArgs: stdArgs, incrementalFlag: incrementalFlag, headerSearchPathArgs: headerSearchPathArgs, + externalIncludeArgs: externalIncludeArgs, input: inputFile.path, output: outputCpp.path ), @@ -171,7 +207,7 @@ struct CudaBuild: BuildToolPlugin { displayName: "Linking CUDA objects", executable: encuda.url, arguments: linkArguments( - clangpp: clangUrl.url.path, + clangpp: hostCompiler, verboseFlag: verboseFlag, stdArgs: stdArgs, incrementalFlag: incrementalFlag, @@ -195,6 +231,7 @@ struct CudaBuild: BuildToolPlugin { stdArgs: [String], incrementalFlag: [String], headerSearchPathArgs: [String], + externalIncludeArgs: [String], input: String, output: String ) -> [String] { @@ -205,6 +242,7 @@ struct CudaBuild: BuildToolPlugin { arguments += ["--clangpp", clangpp] arguments += ["-I", sourceDir] arguments += headerSearchPathArgs + arguments += externalIncludeArgs arguments += [input] arguments += ["-o", output] return arguments diff --git a/Source/Cmlx/include-framework/Metal.hpp b/Source/Cmlx/include-framework/Metal.hpp index a30f8636c..066c673b4 100644 --- a/Source/Cmlx/include-framework/Metal.hpp +++ b/Source/Cmlx/include-framework/Metal.hpp @@ -2,7 +2,7 @@ // // Metal.hpp.in // -// Autogenerated on October 05, 2026. +// Autogenerated from commit d5839e5e317d5df0f18951e89f0c9c36c8304a73. // // Copyright 2020-2024 Apple Inc. // diff --git a/Source/Cmlx/include-framework/mlx-allocator.h b/Source/Cmlx/include-framework/mlx-allocator.h index b014aee0d..9823a63fb 100644 --- a/Source/Cmlx/include-framework/mlx-allocator.h +++ b/Source/Cmlx/include-framework/mlx-allocator.h @@ -4,9 +4,20 @@ #pragma once #include +#include +#include #include +namespace mlx::core::allocator { +class Data; +} // namespace mlx::core::allocator + +namespace mlx::core::cu { +class CommandEncoder; +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); +} // namespace mlx::core::cu + namespace mlx::core::allocator { // Simple wrapper around buffer pointers @@ -52,12 +63,47 @@ class MLX_API Allocator { MLX_API Allocator& allocator(); -inline Buffer malloc(size_t size) { - return allocator().malloc(size); -} +using Deleter = std::function; + +// The allocator owns data from malloc +class Data { + public: + Data(Buffer buffer, Deleter d) : buffer_(buffer), deleter_(std::move(d)) {} + Data(const Data& other) = delete; + Data& operator=(const Data& other) = delete; + Data(Data&& other) noexcept + : buffer_(std::exchange(other.buffer_, Buffer(nullptr))), + deleter_(std::exchange(other.deleter_, nullptr)), + owned_(std::exchange(other.owned_, false)) {} + ~Data() { + if (owned_) { + allocator().free(buffer_); + } else if (deleter_) { + deleter_(buffer_); + } + } + + Buffer buffer() const& { + return buffer_; + } + Buffer buffer() const&& = delete; + + bool is_owned() const { + return owned_; + } + + private: + Buffer buffer_; + Deleter deleter_; + bool owned_{false}; + + explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {} + friend Data malloc(size_t size); + friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder); +}; -inline void free(Buffer buffer) { - allocator().free(buffer); +inline Data malloc(size_t size) { + return Data(allocator().malloc(size)); } // Make a Buffer from a raw pointer of the given size without a copy. If a diff --git a/Source/Cmlx/include-framework/mlx-array.h b/Source/Cmlx/include-framework/mlx-array.h index 4af752e76..5fde4ee6e 100644 --- a/Source/Cmlx/include-framework/mlx-array.h +++ b/Source/Cmlx/include-framework/mlx-array.h @@ -20,7 +20,6 @@ namespace mlx::core { // Forward declaration class Primitive; -using Deleter = std::function; using ShapeElem = int32_t; using Shape = SmallVector; using Strides = SmallVector; @@ -70,12 +69,8 @@ class MLX_API array { Dtype dtype, const std::function& deleter); - /* Build an array from a buffer */ - explicit array( - allocator::Buffer data, - Shape shape, - Dtype dtype, - Deleter deleter = allocator::free); + /* Build an array from data */ + explicit array(allocator::Data data, Shape shape, Dtype dtype); /** Assignment to rvalue does not compile. */ array& operator=(const array& other) && = delete; @@ -239,22 +234,7 @@ class MLX_API array { return reinterpret_cast(array_desc_->primitive.get()); } - struct Data { - allocator::Buffer buffer; - Deleter d; - Data(allocator::Buffer buffer, Deleter d = allocator::free) - : buffer(buffer), d(d) {} - // Not copyable - Data(const Data& d) = delete; - Data& operator=(const Data& d) = delete; - Data(Data&& o) : buffer(o.buffer), d(o.d) { - o.buffer = allocator::Buffer(nullptr); - o.d = [](allocator::Buffer) {}; - } - ~Data() { - d(buffer); - } - }; + using Data = allocator::Data; struct Flags { // True iff there are no gaps in the underlying data. Each item @@ -303,7 +283,8 @@ class MLX_API array { /** True indicates the arrays buffer is safe to reuse */ bool is_donatable() const { - return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1); + return array_desc_.use_count() == 1 && + (array_desc_->data.use_count() == 1) && array_desc_->data->is_owned(); } /** The array's siblings. */ @@ -360,11 +341,8 @@ class MLX_API array { return array_desc_->data_size; } - allocator::Buffer& buffer() { - return array_desc_->data->buffer; - } - const allocator::Buffer& buffer() const { - return array_desc_->data->buffer; + allocator::Buffer buffer() const { + return array_desc_->data->buffer(); } size_t buffer_size() const { @@ -448,24 +426,14 @@ class MLX_API array { // Check if the array is a tracer array bool is_tracer() const; - void set_data(allocator::Buffer buffer, Deleter d = allocator::free); + void set_data(allocator::Data data); void set_data( - allocator::Buffer buffer, + allocator::Data data, size_t data_size, Strides strides, Flags flags, - Deleter d) { - set_data(buffer, data_size, std::move(strides), flags, 0, std::move(d)); - } - - void set_data( - allocator::Buffer buffer, - size_t data_size, - Strides strides, - Flags flags, - int64_t offset = 0, - Deleter d = allocator::free); + int64_t offset = 0); void copy_shared_buffer( const array& other, diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-device.h b/Source/Cmlx/include-framework/mlx-backend-metal-device.h index b7c06c1a4..3521bbee0 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-device.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-device.h @@ -63,8 +63,13 @@ class MLX_API CommandEncoder { template >> void set_vector_bytes(const Vec& vec, size_t nelems, int idx) { - get_command_encoder()->setBytes( - vec.data(), nelems * sizeof(typename Vec::value_type), idx); + using T = typename Vec::value_type; + if (nelems > 0) { + get_command_encoder()->setBytes(vec.data(), nelems * sizeof(T), idx); + } else { + T val{}; + get_command_encoder()->setBytes(&val, sizeof(T), idx); + } } template >> void set_vector_bytes(const Vec& vec, int idx) { @@ -90,7 +95,7 @@ class MLX_API CommandEncoder { } void barrier(); - void end_encoding(); + void end_encoding(bool wait_in_flight = false); void wait_event(Event event, uint64_t value); void signal_event(Event event, uint64_t value); bool needs_commit() const; diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-resident.h b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h index a5c6b58f4..a72a1cfe2 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-resident.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h @@ -1,11 +1,14 @@ #ifdef __cplusplus -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #pragma once #include +#include #include #include +#include +#include #include #include @@ -97,7 +100,7 @@ class ResidencySets { // 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); + bool add_set_locked(std::string* 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); @@ -119,6 +122,9 @@ class ResidencySets { // taking the lock. std::atomic num_sets_{0}; mutable std::mutex mtx_; + std::condition_variable refresh_cv_; + bool stop_refresh_{false}; + std::thread refresh_thread_; }; } // namespace mlx::core::metal diff --git a/Source/Cmlx/include-framework/mlx-c-distributed_group.h b/Source/Cmlx/include-framework/mlx-c-distributed_group.h index c55ee98c0..d79b6b228 100644 --- a/Source/Cmlx/include-framework/mlx-c-distributed_group.h +++ b/Source/Cmlx/include-framework/mlx-c-distributed_group.h @@ -4,6 +4,7 @@ #define MLX_DISTRIBUTED_GROUP_H #include +#include #include @@ -41,15 +42,67 @@ int mlx_distributed_init( bool strict, const char* bk /* may be null */); +/** + * Byte-level all-gather callback used by the JACCL side channel. + * + * \param rank The rank of the process. + * \param size The size of the group. + * \param src Pointer to this rank's data of size n_bytes. + * \param 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. + * \param n_bytes The number of bytes contributed by each rank. + * \param payload Opaque user data passed to the callback. + * \return Zero on success, non-zero on failure. + */ +typedef int (*mlx_distributed_all_gather_fn)( + int rank, + int size, + const char* src, + char* dst, + size_t n_bytes, + void* payload); + +/** + * Initialize JACCL with a custom metadata all-gather callback. + * + * `bk` must be `jaccl`. The synchronous callback reads `n_bytes` from `src` + * and fills `size * n_bytes` in `dst`, ordered by `rank`. Buffers are borrowed. + * The callback returns zero on success, nonzero on failure. + * + * If a `jaccl` group is already cached, it is returned and `all_gather` is + * never called: the callback is dropped and, with `dtor`, `payload` is + * destroyed before this function returns. Call + * `mlx_distributed_clear_backends()` first to initialize with a new callback. + * + * With `dtor`, ownership of `payload` transfers on entry, including failure. + * Destruction occurs once after all references are released. Cached or + * live groups may delay it. Without `dtor`, `payload` is borrowed and must + * remain valid while the callback is retained. + * + * Returns zero on success, nonzero on error. `res` is unchanged on failure. + */ +int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)); + +/** + * Clear the distributed backend cache. + */ +int mlx_distributed_clear_backends(void); + /** * Get the rank. */ -int mlx_distributed_group_rank(mlx_distributed_group group); +int mlx_distributed_group_get_rank(int* res, mlx_distributed_group group); /** * Get the group size. */ -int mlx_distributed_group_size(mlx_distributed_group group); +int mlx_distributed_group_get_size(int* res, mlx_distributed_group group); /** * Split the group. diff --git a/Source/Cmlx/include-framework/mlx-c-fast.h b/Source/Cmlx/include-framework/mlx-c-fast.h index a69b5b2ea..514ff9a48 100644 --- a/Source/Cmlx/include-framework/mlx-c-fast.h +++ b/Source/Cmlx/include-framework/mlx-c-fast.h @@ -96,6 +96,16 @@ int mlx_fast_cuda_kernel_apply( const mlx_fast_cuda_kernel_config config, const mlx_stream stream); +int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s); int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, diff --git a/Source/Cmlx/include-framework/mlx-c-memory.h b/Source/Cmlx/include-framework/mlx-c-memory.h index 8ae92ae49..839f19069 100644 --- a/Source/Cmlx/include-framework/mlx-c-memory.h +++ b/Source/Cmlx/include-framework/mlx-c-memory.h @@ -30,6 +30,7 @@ extern "C" { int mlx_clear_cache(void); int mlx_get_active_memory(size_t* res); +int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays); int mlx_get_cache_memory(size_t* res); int mlx_get_memory_limit(size_t* res); int mlx_get_peak_memory(size_t* res); diff --git a/Source/Cmlx/include-framework/mlx-c-ops.h b/Source/Cmlx/include-framework/mlx-c-ops.h index 7531fe08f..0e32d381a 100644 --- a/Source/Cmlx/include-framework/mlx-c-ops.h +++ b/Source/Cmlx/include-framework/mlx-c-ops.h @@ -548,6 +548,7 @@ int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s); int mlx_gather_qqmm( @@ -703,6 +704,7 @@ int mlx_matmul( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -1360,6 +1362,15 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s); int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/include-framework/mlx-fft.h b/Source/Cmlx/include-framework/mlx-fft.h index 628113909..5b7cc2fe4 100644 --- a/Source/Cmlx/include-framework/mlx-fft.h +++ b/Source/Cmlx/include-framework/mlx-fft.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include @@ -213,6 +214,7 @@ inline array irfft2( StreamOrDevice s = {}) { return irfftn(a, axes, norm, s); } + /** Compute the discrete Fourier Transform sample frequencies. */ MLX_API array fftfreq(int n, double d = 1.0, StreamOrDevice s = {}); @@ -235,5 +237,31 @@ MLX_API array ifftshift(const array& a, StreamOrDevice s = {}); MLX_API array ifftshift(const array& a, const std::vector& axes, StreamOrDevice s = {}); +/** Compute the Short-Time Fourier Transform of the last axis of `x`. */ +MLX_API array stft( + const array& x, + int n_fft = 2048, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + const std::string& pad_mode = "reflect", + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + StreamOrDevice s = {}); + +/** Compute the inverse Short-Time Fourier Transform. */ +MLX_API array istft( + const array& stft_matrix, + const std::optional& n_fft = std::nullopt, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + const std::optional& length = std::nullopt, + StreamOrDevice s = {}); + } // namespace mlx::core::fft #endif diff --git a/Source/Cmlx/include-framework/mlx-ops.h b/Source/Cmlx/include-framework/mlx-ops.h index d3ff5a0d5..5e3d6cc1e 100644 --- a/Source/Cmlx/include-framework/mlx-ops.h +++ b/Source/Cmlx/include-framework/mlx-ops.h @@ -787,6 +787,23 @@ min(const array& a, MLX_API array min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); +MLX_API std::tuple +minmax(const array& a, bool keepdims, StreamOrDevice s = {}); +inline std::tuple minmax(const array& a, StreamOrDevice s = {}) { + return minmax(a, false, to_stream(s)); +} + +/** The minimum and maximum of the elements of an array along the given axes. */ +MLX_API std::tuple minmax( + const array& a, + const std::vector& axes, + bool keepdims = false, + StreamOrDevice s = {}); + +/** The minimum and maximum of the elements of an array along the given axis. */ +MLX_API std::tuple +minmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); + /** Returns the Hanning window of size M. */ MLX_API array hanning(int M, StreamOrDevice s = {}); @@ -882,6 +899,27 @@ MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); /** Returns topk elements of the array along a given axis. */ MLX_API array topk(const array& a, int k, int axis, StreamOrDevice s = {}); +/** + * Returns the sorted unique elements of the flattened array, and optionally + * the index of the first occurrence of each, the inverse indices and the + * counts. The output has the given size, and + * is truncated if ``size`` is smaller than the number of unique values. + * If ``size`` is larger, the result is padded with ``fill_value``, or the + * first of the sorted unique values if no fill value is provided. An empty + * input has none, so it throws unless ``size`` is zero or a fill value is + * given. + * A truncated output loses the counts of the values it drops, and its inverse + * indices are clamped to the last entry. + */ +MLX_API std::vector unique( + const array& a, + int size, + bool return_index = false, + bool return_inverse = false, + bool return_counts = false, + const std::optional& fill_value = std::nullopt, + StreamOrDevice s = {}); + /** Cumulative logsumexp of an array. */ MLX_API array logcumsumexp( const array& a, diff --git a/Source/Cmlx/include-framework/mlx-primitives.h b/Source/Cmlx/include-framework/mlx-primitives.h index 692f821ab..549b40961 100644 --- a/Source/Cmlx/include-framework/mlx-primitives.h +++ b/Source/Cmlx/include-framework/mlx-primitives.h @@ -1080,8 +1080,13 @@ class FFT : public UnaryPrimitive { Stream stream, const std::vector& axes, bool inverse, - bool real) - : UnaryPrimitive(stream), axes_(axes), inverse_(inverse), real_(real) {} + bool real, + bool odd_out = false) + : UnaryPrimitive(stream), + axes_(axes), + inverse_(inverse), + real_(real), + odd_out_(odd_out) {} void eval_cpu(const std::vector& inputs, array& out) override; void eval_gpu(const std::vector& inputs, array& out) override; @@ -1092,13 +1097,16 @@ class FFT : public UnaryPrimitive { bool is_equivalent(const Primitive& other) const override; auto state() const { - return std::make_tuple(axes_, inverse_, real_); + return std::make_tuple(axes_, inverse_, real_, odd_out_); } private: std::vector axes_; bool inverse_; bool real_; + // The output length of an inverse real transform is not recoverable from + // the input shape. + bool odd_out_; }; class Flatten : public UnaryPrimitive { @@ -1140,6 +1148,20 @@ class Floor : public UnaryPrimitive { DEFINE_INPUT_OUTPUT_SHAPE() }; +class FloorDivide : public UnaryPrimitive { + public: + explicit FloorDivide(Stream stream) : UnaryPrimitive(stream) {} + + 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(FloorDivide) + DEFINE_DEFAULT_IS_EQUIVALENT() + DEFINE_INPUT_OUTPUT_SHAPE() +}; + class Full : public UnaryPrimitive { public: explicit Full(Stream stream) : UnaryPrimitive(stream) {} diff --git a/Source/Cmlx/include-framework/mlx-types-complex.h b/Source/Cmlx/include-framework/mlx-types-complex.h index f9dda2b71..764667911 100644 --- a/Source/Cmlx/include-framework/mlx-types-complex.h +++ b/Source/Cmlx/include-framework/mlx-types-complex.h @@ -46,6 +46,11 @@ struct complex64_t : public std::complex { operator float() const { return real(); }; + + // Explicit, so int casts still go through operator float + explicit operator bool() const { + return real() != 0 || imag() != 0; + }; }; inline bool operator>=(const complex64_t& a, const complex64_t& b) { diff --git a/Source/Cmlx/include-framework/mlx-version.h b/Source/Cmlx/include-framework/mlx-version.h index f672cd22f..f02758282 100644 --- a/Source/Cmlx/include-framework/mlx-version.h +++ b/Source/Cmlx/include-framework/mlx-version.h @@ -7,7 +7,7 @@ #define MLX_VERSION_MAJOR 0 #define MLX_VERSION_MINOR 32 -#define MLX_VERSION_PATCH 3 +#define MLX_VERSION_PATCH 4 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) diff --git a/Source/Cmlx/include/mlx/c/distributed_group.h b/Source/Cmlx/include/mlx/c/distributed_group.h index bfbaa80eb..a784cf639 100644 --- a/Source/Cmlx/include/mlx/c/distributed_group.h +++ b/Source/Cmlx/include/mlx/c/distributed_group.h @@ -4,6 +4,7 @@ #define MLX_DISTRIBUTED_GROUP_H #include +#include #include "mlx/c/stream.h" @@ -41,15 +42,67 @@ int mlx_distributed_init( bool strict, const char* bk /* may be null */); +/** + * Byte-level all-gather callback used by the JACCL side channel. + * + * \param rank The rank of the process. + * \param size The size of the group. + * \param src Pointer to this rank's data of size n_bytes. + * \param 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. + * \param n_bytes The number of bytes contributed by each rank. + * \param payload Opaque user data passed to the callback. + * \return Zero on success, non-zero on failure. + */ +typedef int (*mlx_distributed_all_gather_fn)( + int rank, + int size, + const char* src, + char* dst, + size_t n_bytes, + void* payload); + +/** + * Initialize JACCL with a custom metadata all-gather callback. + * + * `bk` must be `jaccl`. The synchronous callback reads `n_bytes` from `src` + * and fills `size * n_bytes` in `dst`, ordered by `rank`. Buffers are borrowed. + * The callback returns zero on success, nonzero on failure. + * + * If a `jaccl` group is already cached, it is returned and `all_gather` is + * never called: the callback is dropped and, with `dtor`, `payload` is + * destroyed before this function returns. Call + * `mlx_distributed_clear_backends()` first to initialize with a new callback. + * + * With `dtor`, ownership of `payload` transfers on entry, including failure. + * Destruction occurs once after all references are released. Cached or + * live groups may delay it. Without `dtor`, `payload` is borrowed and must + * remain valid while the callback is retained. + * + * Returns zero on success, nonzero on error. `res` is unchanged on failure. + */ +int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)); + +/** + * Clear the distributed backend cache. + */ +int mlx_distributed_clear_backends(void); + /** * Get the rank. */ -int mlx_distributed_group_rank(mlx_distributed_group group); +int mlx_distributed_group_get_rank(int* res, mlx_distributed_group group); /** * Get the group size. */ -int mlx_distributed_group_size(mlx_distributed_group group); +int mlx_distributed_group_get_size(int* res, mlx_distributed_group group); /** * Split the group. diff --git a/Source/Cmlx/include/mlx/c/fast.h b/Source/Cmlx/include/mlx/c/fast.h index 2786c1ce8..affa08dee 100644 --- a/Source/Cmlx/include/mlx/c/fast.h +++ b/Source/Cmlx/include/mlx/c/fast.h @@ -96,6 +96,16 @@ int mlx_fast_cuda_kernel_apply( const mlx_fast_cuda_kernel_config config, const mlx_stream stream); +int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s); int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, diff --git a/Source/Cmlx/include/mlx/c/memory.h b/Source/Cmlx/include/mlx/c/memory.h index bae9e08ec..972142c75 100644 --- a/Source/Cmlx/include/mlx/c/memory.h +++ b/Source/Cmlx/include/mlx/c/memory.h @@ -30,6 +30,7 @@ extern "C" { int mlx_clear_cache(void); int mlx_get_active_memory(size_t* res); +int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays); int mlx_get_cache_memory(size_t* res); int mlx_get_memory_limit(size_t* res); int mlx_get_peak_memory(size_t* res); diff --git a/Source/Cmlx/include/mlx/c/ops.h b/Source/Cmlx/include/mlx/c/ops.h index 741778570..23359f63f 100644 --- a/Source/Cmlx/include/mlx/c/ops.h +++ b/Source/Cmlx/include/mlx/c/ops.h @@ -548,6 +548,7 @@ int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s); int mlx_gather_qqmm( @@ -703,6 +704,7 @@ int mlx_matmul( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -1360,6 +1362,15 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s); int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/mlx-c/CMakeLists.txt b/Source/Cmlx/mlx-c/CMakeLists.txt index f96c72a79..a9a3ed1c6 100644 --- a/Source/Cmlx/mlx-c/CMakeLists.txt +++ b/Source/Cmlx/mlx-c/CMakeLists.txt @@ -8,7 +8,7 @@ set(CMAKE_CXX_STANDARD 20) # ----------------------------- Configuration ----------------------------- if(NOT MLX_C_VERSION) - set(MLX_C_VERSION 0.6.0) + set(MLX_C_VERSION 0.8.0) endif() option(BUILD_SHARED_LIBS "Build mlx C as a shared library" OFF) @@ -35,7 +35,7 @@ else() FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG v0.32.2) + GIT_TAG v0.32.4) FetchContent_MakeAvailable(mlx) endif() diff --git a/Source/Cmlx/mlx-c/docs/src/conf.py b/Source/Cmlx/mlx-c/docs/src/conf.py index 660fbb7df..026ee180c 100644 --- a/Source/Cmlx/mlx-c/docs/src/conf.py +++ b/Source/Cmlx/mlx-c/docs/src/conf.py @@ -10,7 +10,7 @@ project = "MLX C" copyright = "2023-2026, MLX Contributors" author = "MLX Contributors" -version = "0.6.0" +version = "0.8.0" release = version # -- General configuration --------------------------------------------------- diff --git a/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp b/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp index d0db5d5ae..ff483c8b8 100644 --- a/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp @@ -41,21 +41,90 @@ mlx_distributed_init(mlx_distributed_group* res, bool strict, const char* bk) { } } -extern "C" int mlx_distributed_group_rank(mlx_distributed_group group) { +extern "C" int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)) { try { - return mlx_distributed_group_get_(group).rank(); + // Take ownership before validation so failures release the payload. + std::shared_ptr owner; + if (dtor) { + owner = std::shared_ptr(payload, dtor); + } else { + owner = std::shared_ptr(payload, [](void*) {}); + } + + if (!res) { + throw std::invalid_argument("group output must not be null"); + } + if (!bk) { + throw std::invalid_argument("backend must not be null"); + } + if (!all_gather) { + throw std::invalid_argument("all-gather callback must not be null"); + } + + mlx::core::distributed::AllGatherFactory factory = + [all_gather, owner]( + int rank, int size) -> mlx::core::distributed::AllGatherFn { + return [all_gather, owner, rank, size]( + const char* src, char* dst, size_t n_bytes) { + if (all_gather(rank, size, src, dst, n_bytes, owner.get()) != 0) { + throw std::runtime_error("all-gather callback returned an error"); + } + }; + }; + + mlx_distributed_group_set_( + *res, mlx::core::distributed::init(strict, bk, std::move(factory))); + + return 0; } catch (std::exception& e) { mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_distributed_clear_backends(void) { + try { + mlx::core::distributed::clear_backends(); return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; } } -extern "C" int mlx_distributed_group_size(mlx_distributed_group group) { +extern "C" int mlx_distributed_group_get_rank( + int* res, + mlx_distributed_group group) { try { - return mlx_distributed_group_get_(group).size(); + if (!res) { + throw std::invalid_argument("rank output must not be null"); + } + *res = mlx_distributed_group_get_(group).rank(); + return 0; } catch (std::exception& e) { mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_distributed_group_get_size( + int* res, + mlx_distributed_group group) { + try { + if (!res) { + throw std::invalid_argument("size output must not be null"); + } + *res = mlx_distributed_group_get_(group).size(); return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; } } diff --git a/Source/Cmlx/mlx-c/mlx/c/distributed_group.h b/Source/Cmlx/mlx-c/mlx/c/distributed_group.h index bfbaa80eb..a784cf639 100644 --- a/Source/Cmlx/mlx-c/mlx/c/distributed_group.h +++ b/Source/Cmlx/mlx-c/mlx/c/distributed_group.h @@ -4,6 +4,7 @@ #define MLX_DISTRIBUTED_GROUP_H #include +#include #include "mlx/c/stream.h" @@ -41,15 +42,67 @@ int mlx_distributed_init( bool strict, const char* bk /* may be null */); +/** + * Byte-level all-gather callback used by the JACCL side channel. + * + * \param rank The rank of the process. + * \param size The size of the group. + * \param src Pointer to this rank's data of size n_bytes. + * \param 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. + * \param n_bytes The number of bytes contributed by each rank. + * \param payload Opaque user data passed to the callback. + * \return Zero on success, non-zero on failure. + */ +typedef int (*mlx_distributed_all_gather_fn)( + int rank, + int size, + const char* src, + char* dst, + size_t n_bytes, + void* payload); + +/** + * Initialize JACCL with a custom metadata all-gather callback. + * + * `bk` must be `jaccl`. The synchronous callback reads `n_bytes` from `src` + * and fills `size * n_bytes` in `dst`, ordered by `rank`. Buffers are borrowed. + * The callback returns zero on success, nonzero on failure. + * + * If a `jaccl` group is already cached, it is returned and `all_gather` is + * never called: the callback is dropped and, with `dtor`, `payload` is + * destroyed before this function returns. Call + * `mlx_distributed_clear_backends()` first to initialize with a new callback. + * + * With `dtor`, ownership of `payload` transfers on entry, including failure. + * Destruction occurs once after all references are released. Cached or + * live groups may delay it. Without `dtor`, `payload` is borrowed and must + * remain valid while the callback is retained. + * + * Returns zero on success, nonzero on error. `res` is unchanged on failure. + */ +int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)); + +/** + * Clear the distributed backend cache. + */ +int mlx_distributed_clear_backends(void); + /** * Get the rank. */ -int mlx_distributed_group_rank(mlx_distributed_group group); +int mlx_distributed_group_get_rank(int* res, mlx_distributed_group group); /** * Get the group size. */ -int mlx_distributed_group_size(mlx_distributed_group group); +int mlx_distributed_group_get_size(int* res, mlx_distributed_group group); /** * Split the group. diff --git a/Source/Cmlx/mlx-c/mlx/c/fast.cpp b/Source/Cmlx/mlx-c/mlx/c/fast.cpp index 79a788d4c..e89488d2d 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fast.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/fast.cpp @@ -14,7 +14,8 @@ namespace mlx::core { void prefault(const array& a) { if (!a.data_shared_ptr()) return; - const uint8_t* ptr = static_cast(const_cast(a.buffer()).raw_ptr()); + allocator::Buffer buf = a.buffer(); + const uint8_t* ptr = static_cast(buf.raw_ptr()); if (!ptr) return; // volatile to prevent the compiler from optimizing the read away @@ -291,6 +292,37 @@ extern "C" int mlx_fast_cuda_kernel_apply( return 0; } +extern "C" int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s) { + try { + mlx_vector_array_set_( + *res, + mlx::core::fast::gated_delta_update( + mlx_array_get_(queries), + mlx_array_get_(keys), + mlx_array_get_(values), + mlx_array_get_(gates), + mlx_array_get_(beta_), + (initial_state.ctx + ? std::make_optional(mlx_array_get_(initial_state)) + : std::nullopt), + (mask.ctx ? std::make_optional(mlx_array_get_(mask)) + : std::nullopt), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, @@ -641,38 +673,6 @@ extern "C" int mlx_fast_rope_dynamic( } return 0; } -extern "C" int mlx_fast_scaled_dot_product_attention( - mlx_array* res, - const mlx_array queries, - const mlx_array keys, - const mlx_array values, - float scale, - const char* mask_mode, - const mlx_array mask_arr /* may be null */, - const mlx_array sinks /* may be null */, - bool force_fused, - const mlx_stream s) { - try { - mlx_array_set_( - *res, - mlx::core::fast::scaled_dot_product_attention( - mlx_array_get_(queries), - mlx_array_get_(keys), - mlx_array_get_(values), - scale, - std::string(mask_mode), - (mask_arr.ctx ? std::make_optional(mlx_array_get_(mask_arr)) - : std::nullopt), - (sinks.ctx ? std::make_optional(mlx_array_get_(sinks)) - : std::nullopt), - force_fused, - mlx_stream_get_(s))); - } catch (std::exception& e) { - mlx_error(e.what()); - return 1; - } - return 0; -} #include #include @@ -1221,3 +1221,35 @@ extern "C" int mlx_fast_pread_into_offset( } return 0; } +extern "C" int mlx_fast_scaled_dot_product_attention( + mlx_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + float scale, + const char* mask_mode, + const mlx_array mask_arr /* may be null */, + const mlx_array sinks /* may be null */, + bool force_fused, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::fast::scaled_dot_product_attention( + mlx_array_get_(queries), + mlx_array_get_(keys), + mlx_array_get_(values), + scale, + std::string(mask_mode), + (mask_arr.ctx ? std::make_optional(mlx_array_get_(mask_arr)) + : std::nullopt), + (sinks.ctx ? std::make_optional(mlx_array_get_(sinks)) + : std::nullopt), + force_fused, + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} diff --git a/Source/Cmlx/mlx-c/mlx/c/fast.h b/Source/Cmlx/mlx-c/mlx/c/fast.h index 2786c1ce8..affa08dee 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fast.h +++ b/Source/Cmlx/mlx-c/mlx/c/fast.h @@ -96,6 +96,16 @@ int mlx_fast_cuda_kernel_apply( const mlx_fast_cuda_kernel_config config, const mlx_stream stream); +int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s); int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, diff --git a/Source/Cmlx/mlx-c/mlx/c/memory.cpp b/Source/Cmlx/mlx-c/mlx/c/memory.cpp index fb38921d5..c72b9415b 100644 --- a/Source/Cmlx/mlx-c/mlx/c/memory.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/memory.cpp @@ -26,6 +26,17 @@ extern "C" int mlx_get_active_memory(size_t* res) { } return 0; } +extern "C" int mlx_get_array_buffer_size( + size_t* res, + const mlx_vector_array arrays) { + try { + *res = mlx::core::get_array_buffer_size(mlx_vector_array_get_(arrays)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_get_cache_memory(size_t* res) { try { *res = mlx::core::get_cache_memory(); diff --git a/Source/Cmlx/mlx-c/mlx/c/memory.h b/Source/Cmlx/mlx-c/mlx/c/memory.h index bae9e08ec..972142c75 100644 --- a/Source/Cmlx/mlx-c/mlx/c/memory.h +++ b/Source/Cmlx/mlx-c/mlx/c/memory.h @@ -30,6 +30,7 @@ extern "C" { int mlx_clear_cache(void); int mlx_get_active_memory(size_t* res); +int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays); int mlx_get_cache_memory(size_t* res); int mlx_get_memory_limit(size_t* res); int mlx_get_peak_memory(size_t* res); diff --git a/Source/Cmlx/mlx-c/mlx/c/ops.cpp b/Source/Cmlx/mlx-c/mlx/c/ops.cpp index 9e1b06bd4..50ba9b846 100644 --- a/Source/Cmlx/mlx-c/mlx/c/ops.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/ops.cpp @@ -1687,6 +1687,7 @@ extern "C" int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s) { try { @@ -1708,7 +1709,8 @@ extern "C" int mlx_gather_qmm( (bits.has_value ? std::make_optional(bits.value) : std::nullopt), std::string(mode), - std::nullopt, // global_scale (added in mlx 0.32.3; not exposed in the C API) + (global_scale.ctx ? std::make_optional(mlx_array_get_(global_scale)) + : std::nullopt), sorted_indices, mlx_stream_get_(s))); } catch (std::exception& e) { @@ -2294,6 +2296,18 @@ extern "C" int mlx_matmul( } return 0; } +extern "C" int +mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::matrix_transpose(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -4380,6 +4394,33 @@ extern "C" int mlx_unflatten( } return 0; } +extern "C" int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s) { + try { + mlx_vector_array_set_( + *res, + mlx::core::unique( + mlx_array_get_(a), + size, + return_index, + return_inverse, + return_counts, + (fill_value.ctx ? std::make_optional(mlx_array_get_(fill_value)) + : std::nullopt), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/mlx-c/mlx/c/ops.h b/Source/Cmlx/mlx-c/mlx/c/ops.h index 741778570..23359f63f 100644 --- a/Source/Cmlx/mlx-c/mlx/c/ops.h +++ b/Source/Cmlx/mlx-c/mlx/c/ops.h @@ -548,6 +548,7 @@ int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s); int mlx_gather_qqmm( @@ -703,6 +704,7 @@ int mlx_matmul( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -1360,6 +1362,15 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s); int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/mlx-generated/compiled_preamble.cpp b/Source/Cmlx/mlx-generated/compiled_preamble.cpp index c23168e74..48006a428 100644 --- a/Source/Cmlx/mlx-generated/compiled_preamble.cpp +++ b/Source/Cmlx/mlx-generated/compiled_preamble.cpp @@ -115,6 +115,9 @@ struct complex64_t : public std::complex { operator float() const { return real(); }; + explicit operator bool() const { + return real() != 0 || imag() != 0; + }; }; inline bool operator>=(const complex64_t& a, const complex64_t& b) { return (a.real() > b.real()) || @@ -305,6 +308,27 @@ Simd divide(Simd a_, Simd b_) { return a / b; } template +Simd floor_divide(Simd a_, Simd b_) { + T a = a_.value; + T b = b_.value; + if constexpr (std::is_integral_v) { + if (b == 0) { + return T(0); + } + T q = a / b; + if constexpr (is_signed_v) { + if (a % b != 0 && (a < 0 != b < 0)) { + q -= 1; + } + } + return q; + } else if constexpr (is_complex) { + return a / b; + } else { + return floor(divide(a_, b_)); + } +} +template Simd remainder(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; @@ -686,6 +710,7 @@ using namespace mlx::core::simd; struct Add { template Simd operator()(Simd x, Simd y) { return operator+(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct ArcTan2 { template Simd operator()(Simd x, Simd y) { return atan2(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct Divide { template Simd operator()(Simd x, Simd y) { return divide(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; +struct FloorDivide { template Simd operator()(Simd x, Simd y) { return floor_divide(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct Multiply { template Simd operator()(Simd x, Simd y) { return operator*(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct Subtract { template Simd operator()(Simd x, Simd y) { return operator-(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct LogicalAnd { template Simd operator()(Simd x, Simd y) { return operator&&(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; diff --git a/Source/Cmlx/mlx-generated/conv.cpp b/Source/Cmlx/mlx-generated/conv.cpp index 25c6d5433..3cdec8672 100644 --- a/Source/Cmlx/mlx-generated/conv.cpp +++ b/Source/Cmlx/mlx-generated/conv.cpp @@ -224,8 +224,8 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -374,8 +374,8 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -561,6 +561,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -582,6 +583,7 @@ struct Conv2DWeightBlockLoader { weight_hw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -596,7 +598,7 @@ struct Conv2DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -658,9 +660,9 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; short kdil_d; short kdil_h; @@ -842,9 +844,9 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -1064,6 +1066,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -1085,6 +1088,7 @@ struct Conv3DWeightBlockLoader { weight_dhw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -1099,7 +1103,7 @@ struct Conv3DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -1206,6 +1210,7 @@ struct Conv2DInputBlockLoaderSmallChannels { // Rows / strided reads within the block STEEL_CONST short n_rows = BROWS / TROWS; + static_assert(n_rows > 0, "BM must be at least tgp_size * vec_size / BK"); // Thread location indices const short thread_idx; @@ -1368,6 +1373,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -1388,6 +1394,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { params(params_), weight_hw(thread_idx % TCOLS), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -1395,7 +1402,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { if (bi >= BROWS || bj >= BCOLS) return; - if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -1424,7 +1431,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } } else { for (short i = 0; i < BROWS; i += TROWS) { - if (((read_n + i) < params->O)) { + if (((read_n + i) < N)) { STEEL_PRAGMA_UNROLL for (short j = 0; j < n_channels; j++) { dst[i * dst_ld + j] = curr_src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h index b7089f3a2..0dab87176 100644 --- a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -755,7 +755,7 @@ constexpr char jit_source_cast_op[] = { 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, - 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, 0x69, + 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x7c, 0x7c, 0x20, 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, @@ -772,7 +772,7 @@ constexpr char jit_source_cast_op[] = { 0x72, 0x28, 0x29, 0x28, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, - 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x31, 0x2c, 0x20, 0x31, 0x7d, 0x20, + 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x31, 0x2c, 0x20, 0x30, 0x7d, 0x20, 0x3a, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x30, 0x2c, 0x20, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6e, @@ -2043,6 +2043,140 @@ constexpr char jit_source_fp16_math[] = { 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; +constexpr char jit_source_gather[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, + 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x4e, 0x49, 0x44, 0x58, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, 0x44, + 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, + 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, + 0x69, 0x64, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, + 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, + 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, + 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, + 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, + 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, + 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, + 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, + 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, + 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, + 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, + 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, + 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, + 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x25, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x65, + 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, + 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, + 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, + 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, + 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, + 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, + 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, + 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x73, 0x72, + 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, + 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, + 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + constexpr char jit_source_gather_axis[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, @@ -2185,140 +2319,6 @@ constexpr char jit_source_gather_axis[] = { 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; -constexpr char jit_source_gather[] = { - 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, - 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, - 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, - 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, - 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, - 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, - 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, - 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, - 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, - 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, - 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, - 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x4e, 0x49, 0x44, 0x58, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, 0x44, - 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, - 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, - 0x69, 0x64, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x28, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, - 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, - 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, - 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, - 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, - 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, - 0x65, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, - 0x32, 0x5f, 0x74, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, - 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, - 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, - 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, - 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, - 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, - 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, - 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, - 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, - 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, - 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, - 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, - 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, - 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, - 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, - 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, - 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, - 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, - 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, - 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, - 0x20, 0x25, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, - 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, - 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, - 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, - 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x65, - 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, - 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, - 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, - 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, - 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, - 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, - 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, - 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, - 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, - 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, - 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, - 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, - 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, - 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, - 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, - 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, - 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, - 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, - 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, - 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x73, 0x72, - 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, - 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, - 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, - 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, - 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, - 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, - 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, -}; - constexpr char jit_source_gemm_sm70[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, @@ -5494,47 +5494,288 @@ constexpr char jit_source_qmm_sm80[] = { 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x20, 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x28, 0x74, 0x43, 0x72, 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x28, 0x69, 0x29, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, - 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2c, 0x20, - 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, 0x2c, 0x20, 0x72, 0x32, - 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x5f, - 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, - 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, - 0x66, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, - 0x2c, 0x20, 0x74, 0x43, 0x70, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, - 0x74, 0x43, 0x73, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, - 0x67, 0x43, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, - 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, - 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, - 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, - 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, - 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3d, 0x20, 0x49, 0x6e, 0x74, - 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x7b, - 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6d, 0x61, 0x6b, - 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, - 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, - 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, - 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, - 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x30, - 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, + 0x65, 0x20, 0x6c, 0x61, 0x73, 0x74, 0x20, 0x70, 0x72, 0x65, 0x66, 0x65, + 0x74, 0x63, 0x68, 0x65, 0x73, 0x20, 0x6d, 0x61, 0x79, 0x20, 0x73, 0x74, + 0x69, 0x6c, 0x6c, 0x20, 0x62, 0x65, 0x20, 0x6c, 0x61, 0x6e, 0x64, 0x69, + 0x6e, 0x67, 0x20, 0x69, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x62, 0x75, + 0x66, 0x66, 0x65, 0x72, 0x73, 0x20, 0x43, 0x20, 0x72, 0x65, 0x75, 0x73, + 0x65, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x63, 0x70, 0x5f, 0x61, 0x73, 0x79, + 0x6e, 0x63, 0x5f, 0x77, 0x61, 0x69, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, + 0x6f, 0x70, 0x79, 0x28, 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x63, 0x2c, 0x20, 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, + 0x2c, 0x20, 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x63, 0x2c, 0x20, 0x74, 0x43, 0x70, 0x43, 0x2c, 0x20, + 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x2c, 0x20, 0x73, 0x32, + 0x67, 0x5f, 0x74, 0x43, 0x67, 0x43, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, + 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3d, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, + 0x7a, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, - 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, - 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, + 0x7a, 0x65, 0x29, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x28, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x3c, 0x5f, 0x30, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, + 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, + 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, + 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, 0x61, + 0x75, 0x6e, 0x63, 0x68, 0x5f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, 0x5f, + 0x5f, 0x28, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x29, 0x29, 0x29, 0x3a, 0x3a, + 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6b, 0x65, 0x72, + 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, + 0x6c, 0x65, 0x2a, 0x20, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2a, 0x20, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, + 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, + 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, + 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, 0x61, + 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, + 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, + 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x7a, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, + 0x69, 0x6e, 0x65, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, + 0x28, 0x6d, 0x69, 0x78, 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, + 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, + 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, + 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, + 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, + 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4e, 0x2c, + 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x49, 0x6e, + 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, + 0x6e, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, + 0x4e, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, + 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, + 0x6e, 0x64, 0x6c, 0x65, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, + 0x73, 0x74, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, + 0x32, 0x3e, 0x28, 0x64, 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x29, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, + 0x65, 0x73, 0x65, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, + 0x6c, 0x6c, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, + 0x6d, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, + 0x70, 0x74, 0x72, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, + 0x29, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, + 0x64, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, + 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, + 0x43, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, + 0x20, 0x6e, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, + 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, + 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, + 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x46, 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, + 0x72, 0x2c, 0x20, 0x75, 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x6c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x69, 0x6e, 0x70, 0x75, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, + 0x73, 0x6c, 0x69, 0x63, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, 0x73, 0x5f, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, + 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x20, 0x3f, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, + 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x61, + 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, 0x5f, 0x6e, + 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, + 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, + 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, + 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, + 0x20, 0x74, 0x68, 0x65, 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, + 0x69, 0x61, 0x74, 0x65, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x20, 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, + 0x65, 0x72, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, + 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x6d, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, + 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, + 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, + 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, + 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, + 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, + 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, + 0x6d, 0x53, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, + 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, + 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, + 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, + 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x20, 0x3d, 0x20, 0x6d, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x30, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, + 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, + 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, + 0x70, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x53, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, - 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, - 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, - 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x74, 0x79, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, 0x61, 0x75, 0x6e, 0x63, 0x68, @@ -5543,223 +5784,324 @@ constexpr char jit_source_qmm_sm80[] = { 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x29, 0x29, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, - 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, - 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x41, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x51, 0x75, 0x61, - 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2a, 0x20, - 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x5a, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, - 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x6c, 0x68, 0x73, - 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, - 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, - 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x43, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x2c, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, - 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, - 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, - 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, - 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, - 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, - 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, - 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, - 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, - 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x79, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x5f, 0x63, 0x6f, - 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, - 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x7a, 0x29, 0x3b, 0x0a, 0x0a, - 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, - 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x6d, 0x69, 0x78, - 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, - 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, - 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, - 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, - 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, - 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, - 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, - 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, - 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, - 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, - 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, - 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, - 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, - 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, - 0x75, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, - 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x3b, 0x0a, - 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, 0x6e, 0x64, 0x6c, 0x65, - 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x69, 0x6e, - 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x62, 0x72, 0x6f, - 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x64, - 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, - 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x29, - 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, 0x65, 0x73, 0x65, 0x6e, - 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, 0x6c, 0x6c, 0x20, 0x74, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, - 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x20, + 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x72, 0x68, 0x73, 0x5f, 0x6b, 0x65, 0x72, + 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, + 0x6c, 0x65, 0x2a, 0x20, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2a, 0x20, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, + 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2a, 0x20, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6e, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6b, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, + 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, + 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x20, 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, + 0x72, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x20, 0x3d, 0x20, 0x6d, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x2a, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, + 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x6d, 0x20, 0x2d, 0x20, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, + 0x20, 0x3d, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3c, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x20, + 0x3f, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3a, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, + 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x6d, 0x69, + 0x78, 0x65, 0x64, 0x29, 0x2e, 0x20, 0x41, 0x6c, 0x6c, 0x20, 0x67, 0x61, + 0x74, 0x68, 0x65, 0x72, 0x65, 0x64, 0x20, 0x72, 0x6f, 0x77, 0x73, 0x20, + 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x6f, 0x6e, 0x65, 0x20, 0x6d, 0x61, 0x74, + 0x72, 0x69, 0x78, 0x3b, 0x20, 0x65, 0x61, 0x63, 0x68, 0x20, 0x72, 0x75, + 0x6e, 0x20, 0x6f, 0x66, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x65, 0x71, + 0x75, 0x61, 0x6c, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x20, 0x69, 0x6e, 0x73, 0x69, 0x64, 0x65, 0x20, 0x74, + 0x68, 0x65, 0x20, 0x72, 0x6f, 0x77, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, + 0x72, 0x65, 0x75, 0x73, 0x65, 0x73, 0x20, 0x6f, 0x6e, 0x65, 0x20, 0x65, + 0x78, 0x70, 0x65, 0x72, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, + 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, + 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, + 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, + 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, + 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, + 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x47, 0x72, + 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x28, 0x6e, 0x2c, 0x20, + 0x6b, 0x2c, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, + 0x6e, 0x74, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, + 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x31, 0x29, 0x2c, 0x20, 0x64, 0x41, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x3c, 0x51, 0x75, + 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, 0x29, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6b, + 0x2c, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, 0x6e, + 0x74, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, - 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, - 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, - 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x3c, - 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, 0x29, 0x2c, 0x20, 0x6d, - 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, - 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, - 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, - 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, - 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, - 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, - 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, - 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, - 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, - 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, - 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, - 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, - 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, - 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, - 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, - 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, - 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x46, - 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x2c, 0x20, 0x75, - 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x6c, 0x6f, 0x6f, - 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x69, 0x6e, 0x70, 0x75, - 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, - 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, - 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, - 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, - 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, - 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, - 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, - 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, - 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, - 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, - 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, - 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, - 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, - 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, - 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, - 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, - 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, - 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, - 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, - 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, - 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, - 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, - 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, - 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, - 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, - 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, - 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, - 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, - 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, + 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, 0x31, 0x29, 0x2c, 0x20, + 0x64, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, + 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, + 0x28, 0x53, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, + 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x5a, + 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, - 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, - 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, - 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, - 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, - 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, - 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, - 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, - 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, - 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, - 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, - 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, - 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, - 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, - 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, - 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, - 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, - 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, - 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, - 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, + 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, + 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, + 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, + 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x77, 0x68, 0x69, 0x6c, 0x65, 0x20, 0x28, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3c, 0x20, + 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x3d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, + 0x78, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x67, 0x70, + 0x5f, 0x62, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x6f, 0x77, 0x20, 0x3d, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x20, 0x31, 0x3b, 0x20, + 0x72, 0x6f, 0x77, 0x20, 0x3c, 0x20, 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, + 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x6f, 0x77, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, + 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x72, 0x68, + 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x79, 0x5f, + 0x72, 0x6f, 0x77, 0x20, 0x2b, 0x20, 0x72, 0x6f, 0x77, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x21, 0x3d, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x20, 0x3d, 0x20, 0x72, 0x6f, 0x77, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, + 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x72, 0x65, + 0x61, 0x6b, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, + 0x3d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x20, 0x2d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x41, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, + 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x79, 0x5f, 0x72, 0x6f, + 0x77, 0x20, 0x2b, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, + 0x30, 0x29, 0x2c, 0x20, 0x6d, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x20, 0x2b, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x30, 0x29, 0x2c, + 0x20, 0x6d, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, + 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, + 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, 0x6e, + 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, + 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, + 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x75, + 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x20, 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, - 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, - 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x41, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x72, 0x75, + 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, + 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, 0x20, 0x72, + 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, + 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, - 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, - 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, - 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, - 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, - 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, - 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, - 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, - 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, - 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, - 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, - 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, - 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, - 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, - 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x5a, 0x20, - 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, - 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, - 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, - 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, - 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, - 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, - 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, - 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, 0x6f, 0x72, - 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, - 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x5f, 0x6d, - 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, - 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x67, - 0x41, 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, - 0x5f, 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x38, 0x30, - 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x28, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, - 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x42, 0x2c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, - 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, - 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, - 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, - 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, + 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, + 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, + 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, + 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, + 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x5a, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, + 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, + 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, + 0x38, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x42, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, + 0x77, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x72, 0x6f, + 0x77, 0x73, 0x20, 0x3d, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, + 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x3a, 0x3a, 0x76, + 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3e, 0x20, + 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, + 0x20, 0x3c, 0x3d, 0x20, 0x31, 0x36, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, + 0x6d, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x2c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x28, + 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3c, 0x3d, 0x20, + 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x6e, + 0x74, 0x3c, 0x33, 0x32, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3e, 0x20, 0x31, 0x36, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, + 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3c, 0x3d, 0x20, 0x31, + 0x36, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x6e, 0x74, + 0x3c, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, 0x65, + 0x20, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x72, 0x75, 0x6e, 0x27, 0x73, 0x20, + 0x6c, 0x6f, 0x61, 0x64, 0x73, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x6e, + 0x6f, 0x74, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x77, 0x72, 0x69, 0x74, 0x65, + 0x20, 0x43, 0x20, 0x77, 0x68, 0x69, 0x6c, 0x65, 0x20, 0x69, 0x74, 0x20, + 0x69, 0x73, 0x20, 0x73, 0x74, 0x69, 0x6c, 0x6c, 0x20, 0x62, 0x65, 0x69, + 0x6e, 0x67, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x2e, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_qmm_sm90[] = { @@ -5954,262 +6296,44 @@ constexpr char jit_source_qmm_sm90[] = { 0x3a, 0x3a, 0x53, 0x74, 0x61, 0x67, 0x65, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x41, 0x75, 0x74, 0x6f, 0x43, 0x61, 0x72, 0x76, 0x65, 0x6f, 0x75, 0x74, 0x3c, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, - 0x3c, 0x69, 0x6e, 0x74, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x28, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x70, 0x69, - 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x53, 0x68, 0x61, 0x72, 0x65, - 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x29, 0x29, 0x3e, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, - 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, 0x6d, 0x3a, 0x3a, 0x4b, 0x65, - 0x72, 0x6e, 0x65, 0x6c, 0x54, 0x6d, 0x61, 0x57, 0x61, 0x72, 0x70, 0x53, - 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x69, 0x7a, 0x65, 0x64, 0x43, 0x6f, - 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x3e, 0x3a, 0x3a, - 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x4f, 0x70, - 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, - 0x6d, 0x3a, 0x3a, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3a, 0x3a, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x47, 0x65, 0x6d, 0x6d, 0x55, 0x6e, - 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x3c, 0x53, 0x68, 0x61, 0x70, - 0x65, 0x3c, 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x2c, 0x20, - 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x3e, 0x2c, 0x20, 0x4d, - 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x2c, 0x20, 0x45, 0x70, 0x69, - 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x7d, 0x0a, - 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, - 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, - 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, - 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, - 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, - 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, 0x6b, 0x65, - 0x72, 0x6e, 0x65, 0x6c, 0x5f, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, - 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, - 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, - 0x74, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x43, 0x74, - 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x28, 0x29, 0x29, 0x3b, 0x0a, - 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, - 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, - 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, -}; - -constexpr char jit_source_scatter_axis[] = { - 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, - 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, - 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, - 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, - 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, - 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, - 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, - 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, - 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, - 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, - 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, - 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, - 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, - 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, - 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, - 0x55, 0x70, 0x64, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, - 0x6f, 0x6c, 0x20, 0x49, 0x64, 0x78, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, - 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, - 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, - 0x74, 0x65, 0x72, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x28, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, - 0x70, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x20, 0x69, 0x6e, 0x64, 0x69, - 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, - 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, - 0x72, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, - 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, - 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, - 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, - 0x73, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, - 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, - 0x49, 0x4d, 0x3e, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, - 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, - 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, - 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, - 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, 0x75, - 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, - 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, - 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, - 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, - 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, - 0x20, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, - 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, - 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, - 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x69, - 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, - 0x69, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, - 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, - 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, - 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, - 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, - 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x20, 0x2a, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, - 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, - 0x6f, 0x73, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, - 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x78, 0x2c, 0x20, 0x79, - 0x2c, 0x20, 0x7a, 0x5d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x5f, 0x74, 0x6f, 0x5f, 0x64, 0x69, 0x6d, 0x73, 0x28, 0x69, 0x6e, 0x64, - 0x65, 0x78, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x29, 0x3b, 0x0a, 0x0a, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, - 0x64, 0x78, 0x20, 0x3d, 0x20, 0x7a, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, - 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x3b, 0x0a, - 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, - 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x49, 0x64, 0x78, 0x43, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, - 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, - 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, 0x78, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, - 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, - 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, - 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x5f, - 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, - 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x69, - 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, - 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x69, 0x64, 0x78, 0x5f, - 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, - 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, - 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, - 0x5d, 0x2c, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, - 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, - 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, - 0x61, 0x78, 0x69, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x55, 0x70, - 0x64, 0x43, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, - 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, - 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, - 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, - 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, - 0x63, 0x20, 0x2b, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, - 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, - 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, - 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, - 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, - 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, - 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, - 0x0a, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, - 0x20, 0x2b, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, - 0x75, 0x70, 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, - 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, - 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, - 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, -}; - -constexpr char jit_source_scatter_ops[] = { - 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, - 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, - 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, - 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, - 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, - 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, - 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x61, 0x74, 0x6f, - 0x6d, 0x69, 0x63, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, - 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, - 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, - 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, - 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x41, 0x73, 0x73, 0x69, 0x67, - 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, - 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, - 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x2a, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, - 0x63, 0x74, 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x53, 0x75, - 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, - 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, - 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, - 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x64, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, - 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, - 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x61, 0x78, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, - 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, - 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, - 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, - 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x6f, 0x75, 0x74, - 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, - 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x69, 0x6e, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, - 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, - 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, - 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, - 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, - 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, - 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, - 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, - 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x3c, 0x69, 0x6e, 0x74, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x28, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x70, 0x69, + 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x53, 0x68, 0x61, 0x72, 0x65, + 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x29, 0x29, 0x3e, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, + 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, 0x6d, 0x3a, 0x3a, 0x4b, 0x65, + 0x72, 0x6e, 0x65, 0x6c, 0x54, 0x6d, 0x61, 0x57, 0x61, 0x72, 0x70, 0x53, + 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x69, 0x7a, 0x65, 0x64, 0x43, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x3e, 0x3a, 0x3a, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x4f, 0x70, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, + 0x6d, 0x3a, 0x3a, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3a, 0x3a, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x47, 0x65, 0x6d, 0x6d, 0x55, 0x6e, + 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x3c, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x3c, 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x2c, 0x20, + 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x3e, 0x2c, 0x20, 0x4d, + 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x2c, 0x20, 0x45, 0x70, 0x69, + 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, + 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, 0x6b, 0x65, + 0x72, 0x6e, 0x65, 0x6c, 0x5f, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, + 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, + 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x43, 0x74, + 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x28, 0x29, 0x29, 0x3b, 0x0a, + 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_scatter[] = { @@ -6577,6 +6701,224 @@ constexpr char jit_source_scatter[] = { 0x63, 0x75, 0x0a, 0x00, }; +constexpr char jit_source_scatter_axis[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, + 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, + 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, + 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, + 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, + 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, + 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x55, 0x70, 0x64, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x49, 0x64, 0x78, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, + 0x74, 0x65, 0x72, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, + 0x70, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x20, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, + 0x72, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, + 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, + 0x73, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, + 0x49, 0x4d, 0x3e, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, + 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, + 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, + 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, + 0x20, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, + 0x69, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, + 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, + 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x20, 0x2a, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, + 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, + 0x6f, 0x73, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x78, 0x2c, 0x20, 0x79, + 0x2c, 0x20, 0x7a, 0x5d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x5f, 0x74, 0x6f, 0x5f, 0x64, 0x69, 0x6d, 0x73, 0x28, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x29, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x3d, 0x20, 0x7a, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, + 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x49, 0x64, 0x78, 0x43, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, 0x78, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, + 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, + 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, + 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, + 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, + 0x5d, 0x2c, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, + 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, + 0x61, 0x78, 0x69, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x55, 0x70, + 0x64, 0x43, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, + 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, + 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, + 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, + 0x20, 0x2b, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, + 0x75, 0x70, 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, + 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, + 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x61, 0x74, 0x6f, + 0x6d, 0x69, 0x63, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, + 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, + 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x41, 0x73, 0x73, 0x69, 0x67, + 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x2a, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x53, 0x75, + 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x64, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x61, 0x78, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x6f, 0x75, 0x74, + 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, + 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x69, 0x6e, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, + 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, + 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, + 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + constexpr char jit_source_slice_update[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, diff --git a/Source/Cmlx/mlx-generated/fft.cpp b/Source/Cmlx/mlx-generated/fft.cpp index d60d565b8..36c59094c 100644 --- a/Source/Cmlx/mlx-generated/fft.cpp +++ b/Source/Cmlx/mlx-generated/fft.cpp @@ -664,10 +664,20 @@ constexpr complex_t operator*(complex_t a, complex_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}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -675,11 +685,7 @@ template < 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 = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx-generated/fp_quantized.cpp b/Source/Cmlx/mlx-generated/fp_quantized.cpp index c9fade97e..ad0c884be 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized.cpp @@ -497,6 +497,7 @@ template < typename T, int group_size, int bits, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( @@ -530,6 +531,15 @@ METAL_FUNC void fp_qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // Reuse the last valid row for reads beyond the output tile. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; @@ -539,8 +549,12 @@ METAL_FUNC void fp_qmv_fast_impl( load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device auto* sl = scales + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device auto* sl = scales + src * in_vec_size_g; U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); @@ -558,7 +572,7 @@ METAL_FUNC void fp_qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { if constexpr (has_global_scale) { y[row] = static_cast(result[row] * inv_scale_enc); } else { @@ -959,6 +973,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const bool has_global_scale = false, const int BM = 32, const int BK = 32, @@ -968,7 +983,7 @@ METAL_FUNC void fp_qmm_t_impl( const device uint8_t* scales, const device float* global_scale, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -993,7 +1008,7 @@ METAL_FUNC void fp_qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1406,6 +1421,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( @@ -1447,6 +1463,7 @@ template < T, group_size, bits, + partial_rows, has_global_scale, results_per_simdgroup>( w, @@ -1466,6 +1483,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv( @@ -1674,14 +1692,14 @@ template template < typename T, - const int group_size, - const int bits, - const bool aligned_N, - const bool batched, - const bool has_global_scale = false, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool aligned_N, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, @@ -1725,7 +1743,16 @@ template < s_strides, tid); } - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -1847,7 +1874,7 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( + fp_qmv_fast_impl( w, scales, global_scale, @@ -2044,7 +2071,16 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -2074,7 +2110,7 @@ template < const device uint32_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], const device T* x [[buffer(2)]], - device T* y [[buffer(3)]], + device float* y [[buffer(3)]], const constant int& K [[buffer(4)]], const constant int& N [[buffer(5)]], const constant int& M [[buffer(6)]], @@ -2099,7 +2135,7 @@ template < scales += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - fp_qmm_t_impl( + fp_qmm_t_impl( (const device uint32_t*)wl, scales, nullptr, diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp index 2dff85266..94a90fe8d 100644 --- a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp @@ -1115,10 +1115,11 @@ METAL_FUNC void tile_matmad_nax( } // namespace mlx /////////////////////////////////////////////////////////////////////////////// -// Contents from "mlx/backend/metal/kernels/gated_delta_update_nax.h" +// Contents from "mlx/backend/metal/kernels/gated_delta_nax_ops.h" /////////////////////////////////////////////////////////////////////////////// -#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax.h" +#line 1 "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +// Copyright © 2026 Apple Inc. #include @@ -1130,15 +1131,146 @@ using namespace metal; using namespace mpp; using namespace mpp::tensor_ops; -// NAX MACROS I can probably do a nice template instead of doing this -// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for -// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS #define AT_NAX(TILE, IDX) TILE.elems()[IDX] +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + #define SUB_NAX(TILE0, TILE1, TILE2) \ { \ STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ } \ } @@ -1146,7 +1278,7 @@ using namespace mpp::tensor_ops; #define ADD_NAX(TILE0, TILE1, TILE2) \ { \ STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ } \ } @@ -1177,6 +1309,15 @@ using namespace mpp::tensor_ops; } \ } +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + #define SCALE_BETA_NAX(TILE0, BETA2) \ { \ STEEL_PRAGMA_UNROLL \ @@ -1352,7 +1493,6 @@ METAL_FUNC static constexpr void mman( // Create matmul op mpp::tensor_ops::matmul2d gemm_op; - // Create matmul operands in registers using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< AType, BType, @@ -1362,6 +1502,7 @@ METAL_FUNC static constexpr void mman( BType, CType>()); + // Create matmul operands in registers auto ct_a = gemm_op.template get_left_input_cooperative_tensor(); auto ct_b = @@ -1399,9 +1540,9 @@ METAL_FUNC static constexpr void mman( #define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ mlx::steel::mma< \ - float, \ - float, \ - float, \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ TA, \ TB, \ mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ @@ -1413,9 +1554,9 @@ METAL_FUNC static constexpr void mman( #define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ mlx::steel::mma< \ - float, \ - float, \ - float, \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ TA, \ TB, \ mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ @@ -1427,9 +1568,9 @@ METAL_FUNC static constexpr void mman( #define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ mlx::steel::mma< \ - float, \ - float, \ - float, \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ TA, \ TB, \ mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ @@ -1443,9 +1584,9 @@ METAL_FUNC static constexpr void mman( #define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ mlx::steel::mman< \ - float, \ - float, \ - float, \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ TA, \ TB, \ mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ @@ -1459,9 +1600,9 @@ METAL_FUNC static constexpr void mman( #define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ mlx::steel::mman< \ - float, \ - float, \ - float, \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ TA, \ TB, \ mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ @@ -1473,7 +1614,117 @@ METAL_FUNC static constexpr void mman( B.frag_at(0, (BO) + 1), \ metal::bool_constant{}); -template +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_update_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax.h" +// Copyright © 2026 Apple Inc. + + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +/////////////////////////////////////////////////////////////////////////////// +// Function constants +/////////////////////////////////////////////////////////////////////////////// + +constant bool save_state_cache [[function_constant(200)]]; + +template [[kernel]] void gated_delta_fused_nax( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], @@ -1484,6 +1735,9 @@ template device InT* y [[buffer(6)]], device float* state_out [[buffer(7)]], constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9)]], // [B, Hv, n_ckpt, Dv, Dk] + device float* chunk_mats [[buffer(10)]], // [B, Hv, n_chunks, 3, 16, 16] + device float* chunk_delta [[buffer(11)]], // [B, Hv, n_chunks, C, Dv] uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { @@ -1516,6 +1770,13 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; + const int n_chunks = (T + C - 1) / C; + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + auto c_state = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + + auto o_chunk_mats = chunk_mats + n * n_chunks * 3 * 256; + auto o_chunk_delta = chunk_delta + n * n_chunks * C * Dv + dv_idx; + threadgroup float gamma_all[C * 4]; threadgroup float* gamma = gamma_all + sg_id * C; @@ -1548,6 +1809,10 @@ template } mlx::steel::NAXTile TMP_tile; + // Which chunk process_chunk is on, so it knows whether this is a segment + // boundary. + int chunk_idx = 0; + auto process_chunk = [&](const short valid_rows, auto bounded_tag) __attribute__((always_inline)) { constexpr bool B = decltype(bounded_tag)::value; @@ -1560,19 +1825,23 @@ template } }; + // Checkpoint the state entering this chunk before anything mutates it, on + // segment boundaries only. + if (save_state_cache && (chunk_idx % Ckpt) == 0) { + S_tile.store(c_state, Dk); + c_state += Dv * Dk; + } + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? metal::fast::log( metal::max( - static_cast( - g_[thread_index_in_simdgroup * Hv + hv_idx]), - 1e-6f)) + float(g_[thread_index_in_simdgroup * Hv + hv_idx]), 1e-6)) : 0.0f; auto gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = static_cast(gamma_val); } - simdgroup_barrier(mem_flags::mem_threadgroup); beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; @@ -1584,6 +1853,12 @@ template MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); } + // KKt_raw is identical for every Dv tile of this chunk; only the first + // tile needs to write it out for the vjp kernel to reuse. + if (save_state_cache && dv_idx == 0) { + KKt_tile.store(o_chunk_mats, 16); + } + KKtK_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); @@ -1594,6 +1869,12 @@ template SUB_NAX(Tinv_tile, I_tile, TMP_tile); } + // Same as KKt_raw: cache the pre-decay inverse before it is overwritten + // below by the gamma-scaled (TUinv) version. + if (save_state_cache && dv_idx == 0) { + Tinv_tile.store(o_chunk_mats + 256, 16); + } + STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { load_seq(K_tile, k_ + nn * 16, Dk * Hk); @@ -1615,20 +1896,29 @@ template SUB_NAX(delta_tile, U_tile, WS_tile) - tmp_tile.clear(); QKt_tile.clear(); + tmp_tile.clear(); for (int kk = 0; kk < Dk; kk += 32) { load_seq(Q_tile, q_ + kk, Hk * Dk); load_seq(K_tile, k_ + kk, Hk * Dk); - MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); - SCALE_ROW_NAX(Q_tile, gamma); MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); } - SCALE_TRI_NAX(QKt_tile, gamma) + if (save_state_cache) { + if (dv_idx == 0) { + QKt_tile.store(o_chunk_mats + 512, 16); + } + // Every Dv tile stores its own slice. + delta_tile.store(o_chunk_delta, Dv); + o_chunk_mats += 3 * 256; + o_chunk_delta += C * Dv; + } + // Output path runs unconditionally: the forward serves both inference and + // the training save, so y and state_out are always needed. + SCALE_TRI_NAX(QKt_tile, gamma) out_tile = tmp_tile; MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); @@ -1650,6 +1940,8 @@ template SCALE2_NAX(K_tile, gamma); MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); } + + chunk_idx++; }; int t = 0; diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp new file mode 100644 index 000000000..bf46ed675 --- /dev/null +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp @@ -0,0 +1,3816 @@ +namespace mlx::core::metal { + +const char* gated_delta_update_nax_vjp() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/defines.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/defines.h" +// Copyright © 2024 Apple Inc. + + +#define STEEL_CONST static constant constexpr const +#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") +#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/type_traits.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/type_traits.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace metal { + +template +struct is_empty : metal::bool_constant<__is_empty(T)> {}; + +#ifdef __cpp_variable_templates +template +constexpr constant bool is_empty_v = is_empty::value; +#endif + +template +struct make_void { + typedef void type; +}; + +template +using void_t = typename make_void::type; + +template +struct is_static : metal::bool_constant>::value> {}; + +template +struct pointer_element {}; + +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; + +template +using pointer_element_t = typename pointer_element>::type; + +} // namespace metal + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// Integral constant with casting +/////////////////////////////////////////////////////////////////////////////// + +template +struct integral_constant { + static constexpr constant T value = v; + using value_type = T; + using type = integral_constant; + + METAL_FUNC constexpr operator value_type() const thread noexcept { + return value; + } +}; + +template +using bool_constant = integral_constant; +using true_type = bool_constant; +using false_type = bool_constant; + +template +struct is_integral : bool_constant::value> {}; + +template +struct is_integral> + : bool_constant::value> {}; + +template +constexpr constant bool is_integral_v = is_integral::value; + +template +using Int = integral_constant; + +/////////////////////////////////////////////////////////////////////////////// +// Binary Operators on Integral constants +/////////////////////////////////////////////////////////////////////////////// + +#define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ + } + +integral_const_binop(+, operator+); +integral_const_binop(-, operator-); +integral_const_binop(*, operator*); +integral_const_binop(/, operator/); + +integral_const_binop(==, operator==); +integral_const_binop(!=, operator!=); +integral_const_binop(<, operator<); +integral_const_binop(>, operator>); +integral_const_binop(<=, operator<=); +integral_const_binop(>=, operator>=); + +integral_const_binop(&&, operator&&); +integral_const_binop(||, operator||); + +template >> +METAL_FUNC constexpr auto operator||(true_type, T) { + return true_type{}; +} +template >> +METAL_FUNC constexpr auto operator||(T, true_type) { + return true_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(false_type, T) { + return false_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(T, false_type) { + return false_type{}; +} + +// Dispatch utilities +template +void dispatch_bool(bool v, F f) { + if (v) { + f(true_type{}); + } else { + f(false_type{}); + } +} + +template +constexpr void const_for_loop(F f) { + if constexpr (start < stop) { + constexpr auto idx = Int{}; + f(idx); + const_for_loop(f); + } +} + +#undef integral_const_binop + +/////////////////////////////////////////////////////////////////////////////// +// Reduction operators +/////////////////////////////////////////////////////////////////////////////// + +template +METAL_FUNC constexpr T sum(T x) { + return x; +} + +template +METAL_FUNC constexpr auto sum(T x, Us... us) { + return x + sum(us...); +} + +} // namespace steel +} // namespace mlx + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/nax.h" +// Copyright © 2025 Apple Inc. + + +#include +#include +#include + + +#include + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// MMA helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// NAX Steel with new tiles +/////////////////////////////////////////////////////////////////////////////// + +struct BaseNAXFrag { + STEEL_CONST short kFragRows = 16; + STEEL_CONST short kFragCols = 16; + + STEEL_CONST short kElemsPerFrag = (kFragRows * kFragCols) / 32; + + STEEL_CONST short kElemRows = 2; + STEEL_CONST short kElemCols = 4; + + STEEL_CONST short kElemRowsJump = 8; + + static_assert( + kElemRows * kElemCols == kElemsPerFrag, + "MMAFrag shape is not consistent with MMAFrag size"); + + template + using dtype_frag_t = typename metal::vec; + + METAL_FUNC static short2 get_coord() { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4; + return short2{fn, fm}; + } + + METAL_FUNC static short2 get_coord(short idx) { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)) + (idx >> 2) * 8; + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4 + idx % 4; + return short2{fn, fm}; + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + c + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_rows( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j)]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_safe( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((r < lx) && ((c + j) < ly)) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } else { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_rows( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_safe( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if (r < lx && (c + j) < ly) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename StartX, + typename StopX, + typename StartY, + typename StopY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_slice( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + StartX start_x, + StopX stop_x, + StartY start_y, + StopY stop_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + + const_for_loop<0, kElemRows, 1>([&](auto idx_row) { + const auto r = off_x + idx_row * Int{}; + if (r >= stop_x - sc.y || r < start_x - sc.y) { + return; + } + + const_for_loop<0, kElemCols, 1>([&](auto idx_col) { + const auto c = off_y + idx_col; + if (c >= stop_y - sc.x || c < start_y - sc.x) { + return; + } + + const auto src_idx = idx_row * Int{} + idx_col; + dst[(r + sc.y) * str_x + (c + sc.x) * str_y] = + static_cast(src[src_idx]); + }); + }); + } + + template + METAL_FUNC static constexpr void row_reduce( + thread const dtype_frag_t& inp_vals, + thread T* reduced_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + T thr_reduce = Op::apply( + Op::apply(inp_vals[i * kElemCols + 0], inp_vals[i * kElemCols + 1]), + Op::apply(inp_vals[i * kElemCols + 2], inp_vals[i * kElemCols + 3])); + + T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1)); + qgr_reduce = Op::apply(thr_reduce, qgr_reduce); + + T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8)); + sgr_reduce = Op::apply(qgr_reduce, sgr_reduce); + + reduced_vals[i] = Op::apply(reduced_vals[i], sgr_reduce); + } + } + + template + METAL_FUNC static constexpr void row_bin_op( + thread dtype_frag_t& inp_vals, + thread T* row_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + inp_vals[i * kElemCols + j] = + Op::apply(inp_vals[i * kElemCols + j], row_vals[i]); + } + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; + } + } +}; + +template < + typename T, + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> +struct NAXTile { + using NAXFrag_t = NAXFrag_; + using elem_type = T; + + STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; + STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; + STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; + + STEEL_CONST short kTileRows = kTileRows_; + STEEL_CONST short kTileCols = kTileCols_; + + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; + + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; + + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; + + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; + + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; + + METAL_FUNC NAXTile() thread {} + + 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) + thread { + return val_frags[i * kTileCols + j]; + } + + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j) const thread { + return val_frags[i * kTileCols + j]; + } + + template + 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 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) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + 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 { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + 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 thread { + auto vptr = (thread T*)(&vals); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_reduce( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + 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) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_bin_op( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + 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( + frag_at(), + src, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + 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( + frag_at(), + dst, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + 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( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + 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( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void + 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( + frag_at(), + src, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + 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( + frag_at(), + src, + ld, + Int<1>{}, + src_tile_dims.y, + src_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), + dst, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + 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( + frag_at(), + dst, + ld, + Int<1>{}, + dst_tile_dims.y, + dst_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_slice( + device U* dst, + const int ld, + const short2 start, + 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( + frag_at(), + dst, + ld, + Int<1>{}, + start.y, + stop.y, + start.x, + stop.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } +}; + +template < + class CTile, + class ATile, + class BTile, + bool transpose_a, + bool transpose_b> +METAL_FUNC void tile_matmad_nax( + thread CTile& C, + thread ATile& A, + metal::bool_constant, + thread BTile& B, + metal::bool_constant) { + // Static checks + constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); + + constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); + + constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); + + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; + + if constexpr (TN == 1 && TM % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm += 2) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } + } + } + } +} + +} // namespace steel +} // namespace mlx + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +// Copyright © 2026 Apple Inc. + +#include + +#include +#include + + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/atomic.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/atomic.h" +// Copyright © 2023 Apple Inc. + + +#include +#include + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// Atomic utils +/////////////////////////////////////////////////////////////////////////////// + +#pragma METAL internals : enable +template +constexpr constant bool is_metal_atomic = _disjunction< + is_same, + is_same, + is_same, + is_same>::value; + +#pragma METAL internals : disable + +template +struct mlx_atomic { + atomic val; +}; + +template +struct mlx_atomic>> { + atomic val; +}; + +/////////////////////////////////////////////////////////////////////////////// +// Native metal atomics +/////////////////////////////////////////////////////////////////////////////// + +template , bool> = true> +METAL_FUNC T +mlx_atomic_load_explicit(device mlx_atomic* object, size_t offset) { + return atomic_load_explicit(&(object[offset].val), memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void +mlx_atomic_store_explicit(device mlx_atomic* object, T val, size_t offset) { + atomic_store_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_and_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_and_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_or_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_or_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_min_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_max_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_add_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_add_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_mul_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + T expected = mlx_atomic_load_explicit(object, offset); + while (!mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val * expected, offset)) { + // Workaround: Metal's atomic_compare_exchange_weak_explicit does + // not perform bitwise comparison as required by the C++ atomics spec. + // The compiler lowers the success check to `fcmp fast ueq` under + // no-nans-fp-math, which evaluates to false when either operand is NaN - + // even when the bit patterns are identical. With NaN in memory the CAS + // can never succeed, so the loop spins. Bail out instead: memory is + // already NaN and that is the correct reduction result regardless of + // this thread's update. + if constexpr (metal::is_floating_point_v) { + if (isnan(expected)) { + break; + } + } + } +} + +template , bool> = true> +METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit( + device mlx_atomic* object, + thread T* expected, + T val, + size_t offset) { + return atomic_compare_exchange_weak_explicit( + &(object[offset].val), + expected, + val, + memory_order_relaxed, + memory_order_relaxed); +} + +// Specialization for float since it does not atomic_fetch_min_explicit +template <> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + float val, + size_t offset) { + float expected = mlx_atomic_load_explicit(object, offset); + while (val < expected) { + if (mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val, offset)) { + return; + } + } +} + +// Specialization for float since it does not atomic_fetch_max_explicit +template <> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + float val, + size_t offset) { + float expected = mlx_atomic_load_explicit(object, offset); + while (val > expected) { + if (mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val, offset)) { + return; + } + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Custom atomics +/////////////////////////////////////////////////////////////////////////////// + +namespace { + +template +constexpr constant uint packing_size = sizeof(uint) / sizeof(T); + +template +union uint_or_packed { + T val[packing_size]; + uint bits; +}; + +template +struct mlx_atomic_update_helper { + 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; + } +}; + +template +METAL_FUNC void mlx_atomic_update_and_store( + device mlx_atomic* object, + T update, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + + mlx_atomic_update_helper helper; + uint_or_packed expected; + expected.bits = + atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed); + + while (Op::condition(update, expected.val[elem_offset]) && + !mlx_atomic_compare_exchange_weak_explicit( + object, + &(expected.bits), + helper(expected, update, elem_offset), + pack_offset)) { + } +} + +template +struct __None { + static bool condition(T a, T b) { +#pragma unused(a) +#pragma unused(b) + return true; + } + + T operator()(T a, T b) thread { +#pragma unused(b) + return a; + } +}; + +template +struct __Add { + static bool condition(T a, T b) { +#pragma unused(a) +#pragma unused(b) + return true; + } + + T operator()(T a, T b) thread { + return a + b; + } +}; + +template +struct __Mul { + static bool condition(T a, T b) { +#pragma unused(a) + return b != 0; + } + + T operator()(T a, T b) thread { + return a * b; + } +}; + +template +struct __Max { + static bool condition(T a, T b) { + return a > b; + } + + T operator()(T a, T b) thread { + return max(a, b); + } +}; + +template +struct __Min { + static bool condition(T a, T b) { + return a < b; + } + + T operator()(T a, T b) thread { + return min(a, b); + } +}; + +} // namespace + +template , bool> = true> +METAL_FUNC T +mlx_atomic_load_explicit(device mlx_atomic* object, size_t offset) { + size_t pack_offset = offset / sizeof(T); + size_t elem_offset = offset % sizeof(T); + uint_or_packed packed_val; + packed_val.bits = + atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed); + return packed_val.val[elem_offset]; +} + +template , bool> = true> +METAL_FUNC void +mlx_atomic_store_explicit(device mlx_atomic* object, T val, size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_and_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + uint_or_packed identity; + identity.bits = __UINT32_MAX__; + identity.val[elem_offset] = val; + + atomic_fetch_and_explicit( + &(object[pack_offset].val), identity.bits, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_or_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + uint_or_packed identity; + identity.bits = 0; + identity.val[elem_offset] = val; + + atomic_fetch_or_explicit( + &(object[pack_offset].val), identity.bits, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_add_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_mul_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit( + device mlx_atomic* object, + thread uint* expected, + uint val, + size_t offset) { + return atomic_compare_exchange_weak_explicit( + &(object[offset].val), + expected, + val, + memory_order_relaxed, + memory_order_relaxed); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h" +// Copyright © 2026 Apple Inc. + + +#include + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +template +METAL_FUNC void reduce_tile_tg( + thread TileT& acc, + threadgroup float* scratch, + device mlx_atomic* out, + int row_stride, + int col_off, + short valid_rows, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(sg_id * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _s = 1; _s < kNSG; _s++) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(_s * 32 + simd_lane_id) * kE + _i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Simdgroup 0 alone holds the total. + if (sg_id != 0) { + return; + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * row_stride + col_off + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(out, AT_NAX(acc, _i), idx); + } else { + mlx_atomic_store_explicit(out, AT_NAX(acc, _i), idx); + } + } + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + device float* seg_states [[buffer(17)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + const int n_chunks = (T + C - 1) / C; + // Ceiling: the floor form gives 0 when n_chunks < Ckpt, which would leave the + // allocation empty and fault on the first read. + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + const short tail = short(T - (n_chunks - 1) * C); + + // Bases, indexed by chunk rather than walked. With two nested loops the + // incremental decrements are no longer expressible. + auto q_base = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_base = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dq_base = dq + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dk_base = dk + b_idx * T * Hk * Dk + hk_idx * Dk; + + auto v_base = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto dv_base = dv + b_idx * T * Hv * Dv + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto g_base = g + b_idx * T * Hv; + auto beta_base = beta + b_idx * T * Hv; + auto dg_base = dg + b_idx * T * Hv; + auto db_base = db + b_idx * T * Hv; + + auto cm_base = chunk_mats + n * n_chunks * 3 * 256; + auto cd_base = chunk_delta + n * n_chunks * C * Dv + dv_idx; + auto ck_base = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + // Each simdgroup owns a distinct dv slice, so these regions are disjoint and + // the replay needs no cross-simdgroup synchronisation. + auto seg_base = seg_states + (n * Ckpt * Dv + dv_idx) * Dk; + + // Per-chunk pointers, assigned from the bases before each process_chunk call. + auto q_ = q_base; + auto k_ = k_base; + auto dq_ = dq_base; + auto dk_ = dk_base; + auto v_ = v_base; + auto dv_ = dv_base; + auto co_ = co_base; + auto g_ = g_base; + auto beta_ = beta_base; + auto dg_ = dg_base; + auto db_ = db_base; + auto i_chunk_mats = cm_base; + auto i_chunk_delta = cd_base; + auto i_seg = seg_base; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, taken from the replayed segment. + _M16xDk S_tile; + + _M16x32 K_tile, Q_tile; + _M16x16 V_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // From the replayed segment rather than a per-chunk checkpoint. + S_tile.load(i_seg, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + if (fm < valid_rows) { + db_st[fm] = _b0; + } + if (fm1 < valid_rows) { + db_st[fm1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (fm == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (fm < valid_rows) { + dg_st[fm] = _d0; + } + if (fm1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (fm1 < valid_rows) { + dg_st[fm1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk segments in reverse. Each is replayed forward from its checkpoint to + // recover the entry states, then walked backwards. + // Backward over one segment: replay its entry states from the checkpoint, + // then walk it in reverse. HasTail is a compile-time tag so valid_rows is the + // constant C on every chunk of an interior segment, which folds away the + // row guards in the dq/dk stores and the load_rows switch in the replay. + auto process_segment = [&](const int seg, + auto tail_tag) __attribute__((always_inline)) { + constexpr bool HasTail = decltype(tail_tag)::value; + + const int seg_start = seg * Ckpt; + const int rem = n_chunks - seg_start; + const int n_in_seg = (rem < Ckpt) ? rem : Ckpt; + const int n_full = HasTail ? n_in_seg - 1 : n_in_seg; + + // delta is cached per chunk, so the replay is the state recurrence alone: + // S <- exp(gamma_C) S + K_dec^T delta + // No T^-1, no W, no U -- roughly a tenth of a full chunk's work, which is + // what makes trading Ckpt-fold memory for it worthwhile. + // + // Scoped so S is dead before process_chunk runs: it is a _M16xDk, and + // having it live alongside S_tile and dS_tile would put three of them in + // flight. + { + _M16xDk S; + S.load(ck_base + seg * Dv * Dk, Dk); + + auto replay_one = [&](const int j, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + S.store(seg_base + j * Dv * Dk, Dk); + + const int c = seg_start + j; + auto k_c = k_base + c * C * Hk * Dk; + auto g_c = g_base + c * C * Hv; + auto cd_c = cd_base + c * C * Dv; + + float gv = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + g_c[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + auto gs = simd_prefix_inclusive_sum(gv); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gs); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float g_last = gamma[C - 1]; + const float dec[2] = { + metal::fast::exp(g_last - gamma[fm]), + metal::fast::exp(g_last - gamma[fm1])}; + + _M16x16 delta_tile; + delta_tile.load(cd_c, Dv); + + S = S * metal::fast::exp(g_last); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 Kd; + if constexpr (B) { + Kd.load_rows(k_c + kk, Dk * Hk, valid_rows); + } else { + Kd.load(k_c + kk, Dk * Hk); + } + Kd = scale_rows(Kd, dec); + MMA16x32x16(S, kk / 16, delta_tile, true, 0, Kd, false, 0); + } + }; + + for (int j = 0; j < n_full; j++) { + replay_one(j, C, metal::false_type{}); + } + if constexpr (HasTail) { + replay_one(n_in_seg - 1, tail, metal::true_type{}); + } + } + + auto set_chunk_ptrs = [&](const int j) __attribute__((always_inline)) { + const int c = seg_start + j; + q_ = q_base + c * C * Hk * Dk; + k_ = k_base + c * C * Hk * Dk; + dq_ = dq_base + c * C * Hk * Dk; + dk_ = dk_base + c * C * Hk * Dk; + v_ = v_base + c * C * Hv * Dv; + dv_ = dv_base + c * C * Hv * Dv; + co_ = co_base + c * C * Hv * Dv; + g_ = g_base + c * C * Hv; + beta_ = beta_base + c * C * Hv; + dg_ = dg_base + c * C * Hv; + db_ = db_base + c * C * Hv; + i_chunk_mats = cm_base + c * 3 * 256; + i_chunk_delta = cd_base + c * C * Dv; + i_seg = seg_base + j * Dv * Dk; + }; + + if constexpr (HasTail) { + set_chunk_ptrs(n_in_seg - 1); + process_chunk(tail, metal::true_type{}); + } + for (int j = n_full - 1; j >= 0; --j) { + set_chunk_ptrs(j); + process_chunk(C, metal::false_type{}); + } + }; + + // Only the last segment can hold a short chunk + int seg = n_ckpt - 1; + if (tail != C) { + process_segment(seg, metal::true_type{}); + --seg; + } + for (; seg >= 0; --seg) { + process_segment(seg, metal::false_type{}); + } + + dS_tile.store(o_dh, Dk); +} + +template +[[kernel]] void gated_delta_vjp_fused_nax1( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; // KdKb_tmp folded away + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_r0 < valid_rows) { + db_st[_r0] = _b0; + } + if (_r1 < valid_rows) { + db_st[_r1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (_r0 == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (_r0 < valid_rows) { + dg_st[_r0] = _d0; + } + if (_r1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (_r1 < valid_rows) { + dg_st[_r1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +template +METAL_FUNC void reduce_tile_tg0( + thread TileT& acc, + threadgroup float* scratch, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short lo = kNSG / 2; lo > 0; lo >>= 1) { + if (sg_id >= lo && sg_id < 2 * lo) { + const short dst = sg_id - lo; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(dst * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id < lo) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(sg_id * 32 + simd_lane_id) * kE + _i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax0( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 2) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 NI_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(NI_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(NI_tile, _i) = (_c.x == _c.y) ? -1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw; + MUL_NAX(QKt_tile, QKt_tile, tri_decay); + + TUinv_tile = TWinv_tile; + MUL_NAX(TUinv_tile, TUinv_tile, tri_decay); + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + MUL_NAX(dODT, dODT, tri_decay); + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + SCALE2_P(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + _M16x32 Kb_tile; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + + SCALE_BETA_NAX_O(Kb_tile, K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg(dq_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dq_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + const short _fn = _c.x + _f * 16; + const short _fm = _c.y; + if (_fm < valid_rows) { + const int idx = _fm * Hk * Dk + kk + _fn; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } else { + // Dv reduction done in threadgroup memory: single writer. + mlx_atomic_store_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } + } + } + } + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + SCALE_BETA_NAX(dVb_tile, beta_fm); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dVb_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.y < valid_rows) { + dv_[_c.y * Hv * Dv + dv_idx + _c.x] = AT_NAX(dVb_tile, _i); + } + } + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + MUL_NAX(dTU_tile, dTU_tile, tri_decay); // dTinv = dT_W + (dT_U * D). + + ADD_NAX(dTinv_tile, dTW_tile, dTU_tile); + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + SCALE_NAX(dA_tile, -1.0f); + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = dA_tile; + SCALE_TRIEQ_NAX1(G_tile, beta_fm); + _M16x16 GGt_tile; + GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile; + _M16x16 R_tile; + + MUL_NAX(P_tile, QKt_raw, dODT); + MUL_NAX(R_tile, TWinv_tile, dTU_tile); + ADD_NAX(dgam_pair, P_tile, R_tile); + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, dKb_tile, t1, KdKb_tmp; + _M16x32 QgdOS, TbK, KdKC_blk; + + load_seq(Q_tile, q_ + kk, Hk * Dk); // [16 x 32] + load_seq(K_tile, k_ + kk, Dk * Hk); // [16 x 32] + + // gamma: + rowsum(Q * gamma*(dO @ S)) + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(QgdOS, row_exp); + MULA_NAX(dgam_row32, QgdOS, Q_tile); + + // dODT.T @ Q and (G + G.T) @ K + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + // B (T.T @ dW_raw), dW_raw = -gamma * (ddelta @ S) + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + + // dbeta term uses dK_b before the beta scaling, with raw K + MUL_NAX(KdKb_tmp, dKb_tile, K_tile); + ADD_NAX(KdKb, KdKb, KdKb_tmp); + + SCALE_BETA_NAX(dKb_tile, beta_fm); + ADD_NAX(dk_acc, dk_acc, dKb_tile); + + // D_C * (delta @ dS) + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + SCALE2_P(t1, dec_exp); + ADD_NAX(dk_acc, dk_acc, t1); + + // gamma stuff + MUL_NAX(KdKC_blk, K_tile, t1); + SUB_NAX(dgam_row32, dgam_row32, KdKC_blk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC_blk)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC_blk, _i); + } + + // gamma stuff: gamma_C * sum(S * dS) + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16)[_i] * + dS_tile.frag_at(0, kk / 16)[_i]; + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]; + } + + // gamma stuff + SCALE_BETA_NAX(K_tile, beta_fm); + SCALE_NAX(dW_raw, -1.0f); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + MULS_NAX(dgam_row32, TbK, dW_raw); + + // dk reduces over Dv the same way dq does. + reduce_tile_tg(dk_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dk_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * Hk * Dk + kk + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } else { + mlx_atomic_store_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } + } + } + } + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + _M16x16 VdVb; + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + MUL_NAX(VdVb, V_tile, TUddelta); + + _M16x16 AKKt; + TRIL_NAX(AKKt, dA_tile); + MUL_NAX(AKKt, AKKt, KKt_tile); + + ADD_NAX(VdVb, VdVb, AKKt); + ROWSUM1_NAX(dbeta_acc, VdVb); + ROWSUM1_NAX(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _r = 0; _r < 2; _r++) { + float _b = dbeta_acc[_r]; + _b += simd_shuffle_xor(_b, ushort(1)); + _b += simd_shuffle_xor(_b, ushort(8)); + + const short _row = fm + _r * mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_row < valid_rows) { + db_st[_row] = _b; + } + } + } + + // Fold the two fragments: only the row sum of dgam_row is used downstream, + // and row sums are additive across fragments. + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + AT_NAX(dgam_row, _i) = AT_NAX(dgam_row32, _i) + + AT_NAX(dgam_row32, mlx::steel::BaseNAXFrag::kElemsPerFrag + _i); + } + + // reduce and store dgamma: rowsum(row + pair) - colsum(pair) + ADD_NAX(dgam_row, dgam_row, dgam_pair); + _M16x16 dgam_tile, cs_tile; + MM16x16x16(dgam_tile, 0, dgam_row, false, 0, Ones_tile, false, 0); + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + SUB_NAX(dgam_tile, dgam_tile, cs_tile); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dgam_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.x == 0) { + dg_st[_c.y] = AT_NAX(dgam_tile, _i); + } + } + // The gamma_{C-1}-only term lands on the same slot the loop just wrote, + // so the lanes must be ordered before lane 0 accumulates into it. + simdgroup_barrier(mem_flags::mem_threadgroup); + const float dgam_last_red = simd_sum(dgam_last); + if (simd_lane_id == 0 && valid_rows > 0) { + dg_st[valid_rows - 1] += dgam_last_red; + } + } + + // One barrier serves both stages; simdgroup 0 commits. + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + SCALE_ROW_P(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + NSCALE_ROW_P(W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +// Postprocessing dg. Similar idea to what is done in the triton kernel. +// The kernel computes the gradient with respect to the log. +template +[[kernel]] void gated_delta_dgamma_to_dg( + const device InT* g [[buffer(0)]], + device float* dg [[buffer(1)]], // in: dL/dgamma, out: dL/dg + constant int& T [[buffer(2)]], + constant int& Hv [[buffer(3)]], + constant int& n_total [[buffer(4)]], // B * Hv + uint2 pos [[thread_position_in_grid]]) { + const int n = int(pos.x); // b * Hv + hv + const int c = int(pos.y); // chunk index + + if (n >= n_total) { + return; + } + + const int t0 = c * C; + if (t0 >= T) { + return; + } + const int len = min(C, T - t0); + + const int b_idx = n / Hv; + const int hv_idx = n % Hv; + const int base = b_idx * T * Hv + hv_idx; + + float acc = 0.0f; + for (int j = len - 1; j >= 0; --j) { + const int idx = base + (t0 + j) * Hv; + acc += static_cast(dg[idx]); + float gv = metal::max(static_cast(g[idx]), 1e-6f); + dg[idx] = acc / gv; + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/gemv.cpp b/Source/Cmlx/mlx-generated/gemv.cpp index 0a4a5c78a..71a0757cd 100644 --- a/Source/Cmlx/mlx-generated/gemv.cpp +++ b/Source/Cmlx/mlx-generated/gemv.cpp @@ -681,10 +681,20 @@ constexpr complex_t operator*(complex_t a, complex_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}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -692,11 +702,7 @@ template < 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 = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx-generated/gemv_masked.cpp b/Source/Cmlx/mlx-generated/gemv_masked.cpp index 8e2faebe0..e5269ed6d 100644 --- a/Source/Cmlx/mlx-generated/gemv_masked.cpp +++ b/Source/Cmlx/mlx-generated/gemv_masked.cpp @@ -511,7 +511,7 @@ struct GEMVTKernel { const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x); const int out_mask_offset = - !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0]; + !has_output_mask ? 0 : n_block_idx * out_mask_strides[0]; int mat_mask_offset = !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0]; diff --git a/Source/Cmlx/mlx-generated/masked_scatter.cpp b/Source/Cmlx/mlx-generated/masked_scatter.cpp index d660afad9..25e4ffb0c 100644 --- a/Source/Cmlx/mlx-generated/masked_scatter.cpp +++ b/Source/Cmlx/mlx-generated/masked_scatter.cpp @@ -44,7 +44,7 @@ template if (src_contiguous) { out[idx] = src[batch_idx * src_batch_size + src_index]; } else { - out[idx] = src[elem_to_loc( + out[idx] = src[elem_to_loc( batch_idx * src_batch_size + src_index, src_shapes, src_strides, diff --git a/Source/Cmlx/mlx-generated/metal/complex.h b/Source/Cmlx/mlx-generated/metal/complex.h index 629b6289b..874e837ca 100644 --- a/Source/Cmlx/mlx-generated/metal/complex.h +++ b/Source/Cmlx/mlx-generated/metal/complex.h @@ -264,10 +264,20 @@ constexpr complex_t operator*(complex_t a, complex_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}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -275,11 +285,7 @@ template < 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 = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx-generated/metal/conv.metal b/Source/Cmlx/mlx-generated/metal/conv.metal index 74027df73..0f49414eb 100644 --- a/Source/Cmlx/mlx-generated/metal/conv.metal +++ b/Source/Cmlx/mlx-generated/metal/conv.metal @@ -418,8 +418,8 @@ winograd_conv_2d_weight_transform( // wt_out is stored transposed (A x A x C x O) short ohw_0 = sm * 8 + sn; short ohw_1 = sm * 8 + sn + 1; - device T* wt_out_0 = wt_out + ohw_0 * C * O + ko; - device T* wt_out_1 = wt_out + ohw_1 * C * O + ko; + device T* wt_out_0 = wt_out + size_t(ohw_0) * C * O + ko; + device T* wt_out_1 = wt_out + size_t(ohw_1) * C * O + ko; // Prepare shared memory threadgroup T Ws[BO][R][R][BC]; diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized.h b/Source/Cmlx/mlx-generated/metal/fp_quantized.h index 880ad589f..57d66a7c7 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized.h @@ -334,6 +334,7 @@ template < typename T, int group_size, int bits, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( @@ -367,6 +368,15 @@ METAL_FUNC void fp_qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // Reuse the last valid row for reads beyond the output tile. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; @@ -376,8 +386,12 @@ METAL_FUNC void fp_qmv_fast_impl( load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device auto* sl = scales + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device auto* sl = scales + src * in_vec_size_g; U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); @@ -395,7 +409,7 @@ METAL_FUNC void fp_qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { if constexpr (has_global_scale) { y[row] = static_cast(result[row] * inv_scale_enc); } else { @@ -796,6 +810,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const bool has_global_scale = false, const int BM = 32, const int BK = 32, @@ -805,7 +820,7 @@ METAL_FUNC void fp_qmm_t_impl( const device uint8_t* scales, const device float* global_scale, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -830,7 +845,7 @@ METAL_FUNC void fp_qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1243,6 +1258,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( @@ -1284,6 +1300,7 @@ template < T, group_size, bits, + partial_rows, has_global_scale, results_per_simdgroup>( w, @@ -1303,6 +1320,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv( @@ -1511,14 +1529,14 @@ template template < typename T, - const int group_size, - const int bits, - const bool aligned_N, - const bool batched, - const bool has_global_scale = false, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool aligned_N, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, @@ -1562,7 +1580,16 @@ template < s_strides, tid); } - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -1684,7 +1711,7 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( + fp_qmv_fast_impl( w, scales, global_scale, @@ -1881,7 +1908,16 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -1911,7 +1947,7 @@ template < const device uint32_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], const device T* x [[buffer(2)]], - device T* y [[buffer(3)]], + device float* y [[buffer(3)]], const constant int& K [[buffer(4)]], const constant int& N [[buffer(5)]], const constant int& M [[buffer(6)]], @@ -1936,7 +1972,7 @@ template < scales += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - fp_qmm_t_impl( + fp_qmm_t_impl( (const device uint32_t*)wl, scales, nullptr, diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h b/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h new file mode 100644 index 000000000..154a6e4ba --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h @@ -0,0 +1,588 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include + +#include +#include + +#include "steel/gemm/nax.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h index af29ff161..841be0393 100644 --- a/Source/Cmlx/mlx-generated/metal/gated_delta_update.h +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h @@ -3,6 +3,8 @@ #include #include "utils.h" +constant bool save_state [[function_constant(200)]]; + #define AT(TILE, IDX) TILE.thread_elements()[IDX] #define SUB(TILE0, TILE1, TILE2) \ { \ @@ -263,25 +265,24 @@ template } /* - auto grid = MTL::Size(32, Dv, B * Hv); + auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); */ -template +template [[kernel]] void gated_delta_seq( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], const device float* state_in [[buffer(3)]], - const device InT* g [[buffer(4)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* g [[buffer(4)]], // [B, T, Hv] const device InT* beta [[buffer(5)]], // [B, T, Hv] - // [B, Hv, Dv, Dk] device InT* y [[buffer(6)]], // [B, T, Hv, Dv] device float* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9), function_constant(save_state)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - // kernel implementation auto n = thread_position_in_grid.z; auto b_idx = n / Hv; auto hv_idx = n % Hv; @@ -313,7 +314,23 @@ template auto g_ = g + b_idx * T * Hv; auto beta_ = beta + b_idx * T * Hv; + // Checkpoint every Ckpt-th step rather than every step: at T = 4096 with + // B = 4, Hv = 32 a dense cache is 34 GB, which the backward's replay makes + // unnecessary. + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + auto c_state_base = state_cache + n * n_ckpt * Dv * Dk + dv_idx * Dk; + for (int t = 0; t < T; ++t) { + // Stored before the decay, so slot t/Ckpt holds the state at *entry* to + // step t. The backward replays forward from here, so it must be the entry + // state and not the exit state. + if (save_state && (t % Ckpt) == 0) { + auto c_state = c_state_base + (t / Ckpt) * Dv * Dk; + for (int i = 0; i < n_per_t; ++i) { + c_state[n_per_t * dk_idx + i] = state[i]; + } + } + float kv_mem = 0.0f; for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal b/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal new file mode 100644 index 000000000..230ef4caf --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal @@ -0,0 +1,50 @@ +#include "gated_delta_update.h" +#include "utils.h" + +using namespace metal; + +#define instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) + +#define instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "gated_delta_fused_chunk_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #c, \ + gated_delta_fused_chunk, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +// Ckpt is a template parameter of the sequential kernel, so the inference path +// has to name one even though it never stores. These must cover every value +// gated_delta_ckpt() can return, since the host puts it in the kernel name. +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 16) \ + instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) + +#define instantiate_gated_delta(in_type) \ + instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 64) + +instantiate_gated_delta(float); +instantiate_gated_delta(bfloat16_t); +instantiate_gated_delta(float16_t); diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h index 1a576f4aa..d9c85f432 100644 --- a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h @@ -1,360 +1,19 @@ +// Copyright © 2026 Apple Inc. #pragma once -#include - -#include -#include - -#include "steel/gemm/nax.h" +#include "gated_delta_nax_ops.h" using namespace metal; using namespace mpp; using namespace mpp::tensor_ops; -// NAX MACROS I can probably do a nice template instead of doing this -// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for -// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns -#define AT_NAX(TILE, IDX) TILE.elems()[IDX] - -#define SUB_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ - } \ - } - -#define ADD_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ - } \ - } - -#define FMA_NAX(TILE0, S, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ - (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ - } \ - } - -#define SCALE_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - AT_NAX(TILE0, _i) *= (S); \ - } \ - } - -#define SCALE_ROW_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= \ - metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ - } \ - } - -#define SCALE_BETA_NAX(TILE0, BETA2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ - } \ - } - -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ - AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ - } \ - } - -#define SCALE_TRI_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ - ? 0.f \ - : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ - } \ - } - -#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ - } \ - } - -namespace mlx { -namespace steel { -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A0, - const thread BaseNAXFrag::dtype_frag_t& A1, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B0, - const thread BaseNAXFrag::dtype_frag_t& B1, - metal::bool_constant) { - // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 16, 32, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A0[i]; - ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; - ct_b[i] = B0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; - ct_c[i] = C[i]; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B, - metal::bool_constant) { - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = B[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - ct_c[i] = C[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - } +/////////////////////////////////////////////////////////////////////////////// +// Function constants +/////////////////////////////////////////////////////////////////////////////// - gemm_op.run(ct_a, ct_b, ct_c); +constant bool save_state_cache [[function_constant(200)]]; - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mman( - thread BaseNAXFrag::dtype_frag_t& Cn0, - thread BaseNAXFrag::dtype_frag_t& Cn1, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& Bn0, - const thread BaseNAXFrag::dtype_frag_t& Bn1, - metal::bool_constant) { - // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - // Create matmul op - mpp::tensor_ops::matmul2d gemm_op; - - // Create matmul operands in registers - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - - // Create matmul output in register - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - // Load A in to left operand registers - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = Bn0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; - ct_c[i] = Cn0[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; - } - - // Do matmul - gemm_op.run(ct_a, ct_b, ct_c); - - // Copy out results - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - Cn0[i] = ct_c[i]; - Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; - } -} - -} // namespace steel -} // namespace mlx - -#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - A.frag_at(0, (AO) + 1), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -template +template [[kernel]] void gated_delta_fused_nax( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], @@ -365,6 +24,9 @@ template device InT* y [[buffer(6)]], device float* state_out [[buffer(7)]], constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9)]], // [B, Hv, n_ckpt, Dv, Dk] + device float* chunk_mats [[buffer(10)]], // [B, Hv, n_chunks, 3, 16, 16] + device float* chunk_delta [[buffer(11)]], // [B, Hv, n_chunks, C, Dv] uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { @@ -397,6 +59,13 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; + const int n_chunks = (T + C - 1) / C; + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + auto c_state = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + + auto o_chunk_mats = chunk_mats + n * n_chunks * 3 * 256; + auto o_chunk_delta = chunk_delta + n * n_chunks * C * Dv + dv_idx; + threadgroup float gamma_all[C * 4]; threadgroup float* gamma = gamma_all + sg_id * C; @@ -429,6 +98,10 @@ template } mlx::steel::NAXTile TMP_tile; + // Which chunk process_chunk is on, so it knows whether this is a segment + // boundary. + int chunk_idx = 0; + auto process_chunk = [&](const short valid_rows, auto bounded_tag) __attribute__((always_inline)) { constexpr bool B = decltype(bounded_tag)::value; @@ -441,19 +114,23 @@ template } }; + // Checkpoint the state entering this chunk before anything mutates it, on + // segment boundaries only. + if (save_state_cache && (chunk_idx % Ckpt) == 0) { + S_tile.store(c_state, Dk); + c_state += Dv * Dk; + } + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? metal::fast::log( metal::max( - static_cast( - g_[thread_index_in_simdgroup * Hv + hv_idx]), - 1e-6f)) + float(g_[thread_index_in_simdgroup * Hv + hv_idx]), 1e-6)) : 0.0f; auto gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = static_cast(gamma_val); } - simdgroup_barrier(mem_flags::mem_threadgroup); beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; @@ -465,6 +142,12 @@ template MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); } + // KKt_raw is identical for every Dv tile of this chunk; only the first + // tile needs to write it out for the vjp kernel to reuse. + if (save_state_cache && dv_idx == 0) { + KKt_tile.store(o_chunk_mats, 16); + } + KKtK_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); @@ -475,6 +158,12 @@ template SUB_NAX(Tinv_tile, I_tile, TMP_tile); } + // Same as KKt_raw: cache the pre-decay inverse before it is overwritten + // below by the gamma-scaled (TUinv) version. + if (save_state_cache && dv_idx == 0) { + Tinv_tile.store(o_chunk_mats + 256, 16); + } + STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { load_seq(K_tile, k_ + nn * 16, Dk * Hk); @@ -496,20 +185,29 @@ template SUB_NAX(delta_tile, U_tile, WS_tile) - tmp_tile.clear(); QKt_tile.clear(); + tmp_tile.clear(); for (int kk = 0; kk < Dk; kk += 32) { load_seq(Q_tile, q_ + kk, Hk * Dk); load_seq(K_tile, k_ + kk, Hk * Dk); - MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); - SCALE_ROW_NAX(Q_tile, gamma); MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); } - SCALE_TRI_NAX(QKt_tile, gamma) + if (save_state_cache) { + if (dv_idx == 0) { + QKt_tile.store(o_chunk_mats + 512, 16); + } + // Every Dv tile stores its own slice. + delta_tile.store(o_chunk_delta, Dv); + o_chunk_mats += 3 * 256; + o_chunk_delta += C * Dv; + } + // Output path runs unconditionally: the forward serves both inference and + // the training save, so y and state_out are always needed. + SCALE_TRI_NAX(QKt_tile, gamma) out_tile = tmp_tile; MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); @@ -531,6 +229,8 @@ template SCALE2_NAX(K_tile, gamma); MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); } + + chunk_idx++; }; int t = 0; @@ -548,4 +248,4 @@ template } S_tile.store(o_state, Dk); -} \ No newline at end of file +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h new file mode 100644 index 000000000..0e6ea0939 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h @@ -0,0 +1,1737 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "gated_delta_nax_ops.h" + +#include +#include "atomic.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +template +METAL_FUNC void reduce_tile_tg( + thread TileT& acc, + threadgroup float* scratch, + device mlx_atomic* out, + int row_stride, + int col_off, + short valid_rows, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(sg_id * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _s = 1; _s < kNSG; _s++) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(_s * 32 + simd_lane_id) * kE + _i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Simdgroup 0 alone holds the total. + if (sg_id != 0) { + return; + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * row_stride + col_off + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(out, AT_NAX(acc, _i), idx); + } else { + mlx_atomic_store_explicit(out, AT_NAX(acc, _i), idx); + } + } + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + device float* seg_states [[buffer(17)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + const int n_chunks = (T + C - 1) / C; + // Ceiling: the floor form gives 0 when n_chunks < Ckpt, which would leave the + // allocation empty and fault on the first read. + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + const short tail = short(T - (n_chunks - 1) * C); + + // Bases, indexed by chunk rather than walked. With two nested loops the + // incremental decrements are no longer expressible. + auto q_base = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_base = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dq_base = dq + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dk_base = dk + b_idx * T * Hk * Dk + hk_idx * Dk; + + auto v_base = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto dv_base = dv + b_idx * T * Hv * Dv + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto g_base = g + b_idx * T * Hv; + auto beta_base = beta + b_idx * T * Hv; + auto dg_base = dg + b_idx * T * Hv; + auto db_base = db + b_idx * T * Hv; + + auto cm_base = chunk_mats + n * n_chunks * 3 * 256; + auto cd_base = chunk_delta + n * n_chunks * C * Dv + dv_idx; + auto ck_base = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + // Each simdgroup owns a distinct dv slice, so these regions are disjoint and + // the replay needs no cross-simdgroup synchronisation. + auto seg_base = seg_states + (n * Ckpt * Dv + dv_idx) * Dk; + + // Per-chunk pointers, assigned from the bases before each process_chunk call. + auto q_ = q_base; + auto k_ = k_base; + auto dq_ = dq_base; + auto dk_ = dk_base; + auto v_ = v_base; + auto dv_ = dv_base; + auto co_ = co_base; + auto g_ = g_base; + auto beta_ = beta_base; + auto dg_ = dg_base; + auto db_ = db_base; + auto i_chunk_mats = cm_base; + auto i_chunk_delta = cd_base; + auto i_seg = seg_base; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, taken from the replayed segment. + _M16xDk S_tile; + + _M16x32 K_tile, Q_tile; + _M16x16 V_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // From the replayed segment rather than a per-chunk checkpoint. + S_tile.load(i_seg, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + if (fm < valid_rows) { + db_st[fm] = _b0; + } + if (fm1 < valid_rows) { + db_st[fm1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (fm == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (fm < valid_rows) { + dg_st[fm] = _d0; + } + if (fm1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (fm1 < valid_rows) { + dg_st[fm1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk segments in reverse. Each is replayed forward from its checkpoint to + // recover the entry states, then walked backwards. + // Backward over one segment: replay its entry states from the checkpoint, + // then walk it in reverse. HasTail is a compile-time tag so valid_rows is the + // constant C on every chunk of an interior segment, which folds away the + // row guards in the dq/dk stores and the load_rows switch in the replay. + auto process_segment = [&](const int seg, + auto tail_tag) __attribute__((always_inline)) { + constexpr bool HasTail = decltype(tail_tag)::value; + + const int seg_start = seg * Ckpt; + const int rem = n_chunks - seg_start; + const int n_in_seg = (rem < Ckpt) ? rem : Ckpt; + const int n_full = HasTail ? n_in_seg - 1 : n_in_seg; + + // delta is cached per chunk, so the replay is the state recurrence alone: + // S <- exp(gamma_C) S + K_dec^T delta + // No T^-1, no W, no U -- roughly a tenth of a full chunk's work, which is + // what makes trading Ckpt-fold memory for it worthwhile. + // + // Scoped so S is dead before process_chunk runs: it is a _M16xDk, and + // having it live alongside S_tile and dS_tile would put three of them in + // flight. + { + _M16xDk S; + S.load(ck_base + seg * Dv * Dk, Dk); + + auto replay_one = [&](const int j, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + S.store(seg_base + j * Dv * Dk, Dk); + + const int c = seg_start + j; + auto k_c = k_base + c * C * Hk * Dk; + auto g_c = g_base + c * C * Hv; + auto cd_c = cd_base + c * C * Dv; + + float gv = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + g_c[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + auto gs = simd_prefix_inclusive_sum(gv); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gs); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float g_last = gamma[C - 1]; + const float dec[2] = { + metal::fast::exp(g_last - gamma[fm]), + metal::fast::exp(g_last - gamma[fm1])}; + + _M16x16 delta_tile; + delta_tile.load(cd_c, Dv); + + S = S * metal::fast::exp(g_last); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 Kd; + if constexpr (B) { + Kd.load_rows(k_c + kk, Dk * Hk, valid_rows); + } else { + Kd.load(k_c + kk, Dk * Hk); + } + Kd = scale_rows(Kd, dec); + MMA16x32x16(S, kk / 16, delta_tile, true, 0, Kd, false, 0); + } + }; + + for (int j = 0; j < n_full; j++) { + replay_one(j, C, metal::false_type{}); + } + if constexpr (HasTail) { + replay_one(n_in_seg - 1, tail, metal::true_type{}); + } + } + + auto set_chunk_ptrs = [&](const int j) __attribute__((always_inline)) { + const int c = seg_start + j; + q_ = q_base + c * C * Hk * Dk; + k_ = k_base + c * C * Hk * Dk; + dq_ = dq_base + c * C * Hk * Dk; + dk_ = dk_base + c * C * Hk * Dk; + v_ = v_base + c * C * Hv * Dv; + dv_ = dv_base + c * C * Hv * Dv; + co_ = co_base + c * C * Hv * Dv; + g_ = g_base + c * C * Hv; + beta_ = beta_base + c * C * Hv; + dg_ = dg_base + c * C * Hv; + db_ = db_base + c * C * Hv; + i_chunk_mats = cm_base + c * 3 * 256; + i_chunk_delta = cd_base + c * C * Dv; + i_seg = seg_base + j * Dv * Dk; + }; + + if constexpr (HasTail) { + set_chunk_ptrs(n_in_seg - 1); + process_chunk(tail, metal::true_type{}); + } + for (int j = n_full - 1; j >= 0; --j) { + set_chunk_ptrs(j); + process_chunk(C, metal::false_type{}); + } + }; + + // Only the last segment can hold a short chunk + int seg = n_ckpt - 1; + if (tail != C) { + process_segment(seg, metal::true_type{}); + --seg; + } + for (; seg >= 0; --seg) { + process_segment(seg, metal::false_type{}); + } + + dS_tile.store(o_dh, Dk); +} + +template +[[kernel]] void gated_delta_vjp_fused_nax1( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; // KdKb_tmp folded away + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_r0 < valid_rows) { + db_st[_r0] = _b0; + } + if (_r1 < valid_rows) { + db_st[_r1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (_r0 == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (_r0 < valid_rows) { + dg_st[_r0] = _d0; + } + if (_r1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (_r1 < valid_rows) { + dg_st[_r1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +template +METAL_FUNC void reduce_tile_tg0( + thread TileT& acc, + threadgroup float* scratch, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short lo = kNSG / 2; lo > 0; lo >>= 1) { + if (sg_id >= lo && sg_id < 2 * lo) { + const short dst = sg_id - lo; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(dst * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id < lo) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(sg_id * 32 + simd_lane_id) * kE + _i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax0( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 2) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 NI_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(NI_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(NI_tile, _i) = (_c.x == _c.y) ? -1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw; + MUL_NAX(QKt_tile, QKt_tile, tri_decay); + + TUinv_tile = TWinv_tile; + MUL_NAX(TUinv_tile, TUinv_tile, tri_decay); + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + MUL_NAX(dODT, dODT, tri_decay); + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + SCALE2_P(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + _M16x32 Kb_tile; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + + SCALE_BETA_NAX_O(Kb_tile, K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg(dq_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dq_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + const short _fn = _c.x + _f * 16; + const short _fm = _c.y; + if (_fm < valid_rows) { + const int idx = _fm * Hk * Dk + kk + _fn; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } else { + // Dv reduction done in threadgroup memory: single writer. + mlx_atomic_store_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } + } + } + } + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + SCALE_BETA_NAX(dVb_tile, beta_fm); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dVb_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.y < valid_rows) { + dv_[_c.y * Hv * Dv + dv_idx + _c.x] = AT_NAX(dVb_tile, _i); + } + } + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + MUL_NAX(dTU_tile, dTU_tile, tri_decay); // dTinv = dT_W + (dT_U * D). + + ADD_NAX(dTinv_tile, dTW_tile, dTU_tile); + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + SCALE_NAX(dA_tile, -1.0f); + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = dA_tile; + SCALE_TRIEQ_NAX1(G_tile, beta_fm); + _M16x16 GGt_tile; + GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile; + _M16x16 R_tile; + + MUL_NAX(P_tile, QKt_raw, dODT); + MUL_NAX(R_tile, TWinv_tile, dTU_tile); + ADD_NAX(dgam_pair, P_tile, R_tile); + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, dKb_tile, t1, KdKb_tmp; + _M16x32 QgdOS, TbK, KdKC_blk; + + load_seq(Q_tile, q_ + kk, Hk * Dk); // [16 x 32] + load_seq(K_tile, k_ + kk, Dk * Hk); // [16 x 32] + + // gamma: + rowsum(Q * gamma*(dO @ S)) + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(QgdOS, row_exp); + MULA_NAX(dgam_row32, QgdOS, Q_tile); + + // dODT.T @ Q and (G + G.T) @ K + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + // B (T.T @ dW_raw), dW_raw = -gamma * (ddelta @ S) + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + + // dbeta term uses dK_b before the beta scaling, with raw K + MUL_NAX(KdKb_tmp, dKb_tile, K_tile); + ADD_NAX(KdKb, KdKb, KdKb_tmp); + + SCALE_BETA_NAX(dKb_tile, beta_fm); + ADD_NAX(dk_acc, dk_acc, dKb_tile); + + // D_C * (delta @ dS) + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + SCALE2_P(t1, dec_exp); + ADD_NAX(dk_acc, dk_acc, t1); + + // gamma stuff + MUL_NAX(KdKC_blk, K_tile, t1); + SUB_NAX(dgam_row32, dgam_row32, KdKC_blk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC_blk)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC_blk, _i); + } + + // gamma stuff: gamma_C * sum(S * dS) + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16)[_i] * + dS_tile.frag_at(0, kk / 16)[_i]; + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]; + } + + // gamma stuff + SCALE_BETA_NAX(K_tile, beta_fm); + SCALE_NAX(dW_raw, -1.0f); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + MULS_NAX(dgam_row32, TbK, dW_raw); + + // dk reduces over Dv the same way dq does. + reduce_tile_tg(dk_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dk_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * Hk * Dk + kk + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } else { + mlx_atomic_store_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } + } + } + } + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + _M16x16 VdVb; + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + MUL_NAX(VdVb, V_tile, TUddelta); + + _M16x16 AKKt; + TRIL_NAX(AKKt, dA_tile); + MUL_NAX(AKKt, AKKt, KKt_tile); + + ADD_NAX(VdVb, VdVb, AKKt); + ROWSUM1_NAX(dbeta_acc, VdVb); + ROWSUM1_NAX(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _r = 0; _r < 2; _r++) { + float _b = dbeta_acc[_r]; + _b += simd_shuffle_xor(_b, ushort(1)); + _b += simd_shuffle_xor(_b, ushort(8)); + + const short _row = fm + _r * mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_row < valid_rows) { + db_st[_row] = _b; + } + } + } + + // Fold the two fragments: only the row sum of dgam_row is used downstream, + // and row sums are additive across fragments. + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + AT_NAX(dgam_row, _i) = AT_NAX(dgam_row32, _i) + + AT_NAX(dgam_row32, mlx::steel::BaseNAXFrag::kElemsPerFrag + _i); + } + + // reduce and store dgamma: rowsum(row + pair) - colsum(pair) + ADD_NAX(dgam_row, dgam_row, dgam_pair); + _M16x16 dgam_tile, cs_tile; + MM16x16x16(dgam_tile, 0, dgam_row, false, 0, Ones_tile, false, 0); + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + SUB_NAX(dgam_tile, dgam_tile, cs_tile); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dgam_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.x == 0) { + dg_st[_c.y] = AT_NAX(dgam_tile, _i); + } + } + // The gamma_{C-1}-only term lands on the same slot the loop just wrote, + // so the lanes must be ordered before lane 0 accumulates into it. + simdgroup_barrier(mem_flags::mem_threadgroup); + const float dgam_last_red = simd_sum(dgam_last); + if (simd_lane_id == 0 && valid_rows > 0) { + dg_st[valid_rows - 1] += dgam_last_red; + } + } + + // One barrier serves both stages; simdgroup 0 commits. + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + SCALE_ROW_P(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + NSCALE_ROW_P(W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +// Postprocessing dg. Similar idea to what is done in the triton kernel. +// The kernel computes the gradient with respect to the log. +template +[[kernel]] void gated_delta_dgamma_to_dg( + const device InT* g [[buffer(0)]], + device float* dg [[buffer(1)]], // in: dL/dgamma, out: dL/dg + constant int& T [[buffer(2)]], + constant int& Hv [[buffer(3)]], + constant int& n_total [[buffer(4)]], // B * Hv + uint2 pos [[thread_position_in_grid]]) { + const int n = int(pos.x); // b * Hv + hv + const int c = int(pos.y); // chunk index + + if (n >= n_total) { + return; + } + + const int t0 = c * C; + if (t0 >= T) { + return; + } + const int len = min(C, T - t0); + + const int b_idx = n / Hv; + const int hv_idx = n % Hv; + const int base = b_idx * T * Hv + hv_idx; + + float acc = 0.0f; + for (int j = len - 1; j >= 0; --j) { + const int idx = base + (t0 + j) * Hv; + acc += static_cast(dg[idx]); + float gv = metal::max(static_cast(g[idx]), 1e-6f); + dg[idx] = acc / gv; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h new file mode 100644 index 000000000..86082c272 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h @@ -0,0 +1,184 @@ +// Copyright © 2024 Apple Inc. +#pragma once + +#include +#include + +#include "atomic.h" +#include "utils.h" + +using namespace metal; + +template +[[kernel]] void gated_delta_vjp_seq( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* b [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B*Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device mlx_atomic* dg [[buffer(12)]], + device mlx_atomic* db [[buffer(13)]], + device float* dh [[buffer(14)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + const int qk_stride = Hk * Dk; + const int v_stride = Hv * Dv; + const int g_stride = Hv; + + auto q_base = q + b_idx * T * qk_stride + hk_idx * Dk; + auto k_base = k + b_idx * T * qk_stride + hk_idx * Dk; + auto dq_base = dq + b_idx * T * qk_stride + hk_idx * Dk; + auto dk_base = dk + b_idx * T * qk_stride + hk_idx * Dk; + + auto v_base = v + b_idx * T * v_stride + hv_idx * Dv; + auto dv_base = dv + b_idx * T * v_stride + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * v_stride + hv_idx * Dv; + + auto g_base = g + b_idx * T * g_stride + hv_idx; + auto b_base = b + b_idx * T * g_stride + hv_idx; + auto dg_base = dg + b_idx * T * g_stride + hv_idx; + auto db_base = db + b_idx * T * g_stride + hv_idx; + + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + + float s_hat[n_per_t]; + auto base_state = cot_h + (n * Dv + dv_idx) * Dk; + for (int i = 0; i < n_per_t; i++) { + s_hat[i] = base_state[n_per_t * dk_idx + i]; + } + + // Only this thread's n_per_t slice of the state is ever needed, so a whole + // segment of entry states costs Ckpt * n_per_t registers and the replay needs + // no device buffer. + float seg[Ckpt][n_per_t]; + float s_prev[n_per_t]; + float s_dec[n_per_t]; + + for (int seg_idx = n_ckpt - 1; seg_idx >= 0; --seg_idx) { + const int t0 = seg_idx * Ckpt; + const int seg_len = metal::min(Ckpt, T - t0); + + // Replay forward from the checkpoint, recording the entry state of each + // step in the segment. + auto c_state = + state_cache + n * n_ckpt * Dv * Dk + seg_idx * Dv * Dk + dv_idx * Dk; + + float s[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + s[i] = c_state[n_per_t * dk_idx + i]; + } + + for (int j = 0; j < seg_len; ++j) { + const int t = t0 + j; + + for (int i = 0; i < n_per_t; ++i) { + seg[j][i] = s[i]; + } + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto k_t = k_base + t * qk_stride; + + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + const int s_idx = n_per_t * dk_idx + i; + s[i] *= gamma; + kv_mem += s[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + + float delta = + beta * (static_cast(v_base[t * v_stride + dv_idx]) - kv_mem); + + for (int i = 0; i < n_per_t; ++i) { + s[i] += delta * static_cast(k_t[n_per_t * dk_idx + i]); + } + } + + // Backward over the same steps, newest first. + for (int j = seg_len - 1; j >= 0; --j) { + const int t = t0 + j; + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto q_t = q_base + t * qk_stride; + auto k_t = k_base + t * qk_stride; + auto dq_t = dq_base + t * qk_stride; + auto dk_t = dk_base + t * qk_stride; + + float kv_mem = 0.0f; + float co = static_cast(co_base[t * v_stride + dv_idx]); + float w = 0.0f; + for (int i = 0; i < n_per_t; i++) { + const int s_idx = n_per_t * dk_idx + i; + s_prev[i] = seg[j][i]; + s_dec[i] = s_prev[i] * gamma; + kv_mem += s_dec[i] * static_cast(k_t[s_idx]); + + s_hat[i] += co * static_cast(q_t[s_idx]); + + w += s_hat[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + w = simd_sum(w); + + if (thread_index_in_simdgroup == 0) { + dv_base[t * v_stride + dv_idx] = beta * w; + } + + float u = static_cast(v_base[t * v_stride + dv_idx]) - kv_mem; + float delta = beta * u; + + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(db_base + t * g_stride, w * u, 0); + } + + float dgamma = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + + float s_t = s_dec[i] + delta * static_cast(k_t[s_idx]); + mlx_atomic_fetch_add_explicit(dq_t, co * s_t, s_idx); + + float contrib = beta * (u * s_hat[i] - w * s_dec[i]); + mlx_atomic_fetch_add_explicit(dk_t, contrib, s_idx); + + s_hat[i] -= beta * w * static_cast(k_t[s_idx]); + + dgamma += s_hat[i] * s_prev[i]; + } + dgamma = simd_sum(dgamma); + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(dg_base + t * g_stride, dgamma, 0); + } + + for (int i = 0; i < n_per_t; ++i) { + s_hat[i] *= gamma; + } + } + } + + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + dh[(n * Dv + dv_idx) * Dk + s_idx] = s_hat[i]; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal new file mode 100644 index 000000000..172cef32a --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal @@ -0,0 +1,35 @@ +#include "gated_delta_update_vjp.h" +#include "gated_delta_update.h" +#include "utils.h" + +using namespace metal; + +#define instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_vjp_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_vjp_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) + +#define instantiate_gated_delta_vjp_dims(in_type, dk, dv, hk, hv) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 16) + +#define instantiate_gated_delta_vjp(in_type) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 64) + +instantiate_gated_delta_vjp(float); +instantiate_gated_delta_vjp(bfloat16_t); +instantiate_gated_delta_vjp(float16_t); \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gemv_masked.h b/Source/Cmlx/mlx-generated/metal/gemv_masked.h index 30371eb06..930ba3b94 100644 --- a/Source/Cmlx/mlx-generated/metal/gemv_masked.h +++ b/Source/Cmlx/mlx-generated/metal/gemv_masked.h @@ -452,7 +452,7 @@ struct GEMVTKernel { const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x); const int out_mask_offset = - !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0]; + !has_output_mask ? 0 : n_block_idx * out_mask_strides[0]; int mat_mask_offset = !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0]; diff --git a/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h b/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h index 2ba547400..fe2e2b2fe 100644 --- a/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h +++ b/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h @@ -32,7 +32,7 @@ template if (src_contiguous) { out[idx] = src[batch_idx * src_batch_size + src_index]; } else { - out[idx] = src[elem_to_loc( + out[idx] = src[elem_to_loc( batch_idx * src_batch_size + src_index, src_shapes, src_strides, diff --git a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h index 3a97de9bb..e361e7414 100644 --- a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h +++ b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h @@ -38,8 +38,8 @@ METAL_FUNC void scatter_impl( for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) { LocT out_idx = out_offset; for (int i = 0; i < NIDX; ++i) { - auto idx_loc = indices.row_contiguous[i] - ? ind_idx + LocT idx_loc = indices.row_contiguous[i] + ? LocT(ind_idx) : elem_to_loc( ind_idx, &indices.shapes[indices.ndim * i], @@ -50,7 +50,7 @@ METAL_FUNC void scatter_impl( out_idx += static_cast(idx_val) * static_cast(out_strides[ax]); } - auto upd_idx = ind_idx * static_cast(upd_size) + gid.x; + LocT upd_idx = ind_idx * static_cast(upd_size) + gid.x; if constexpr (!UPD_ROW_CONTIG) { upd_idx = elem_to_loc(upd_idx, upd_shape, upd_strides, upd_ndim); } diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index be1ebff59..48fe55cd9 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -28,15 +28,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -46,8 +61,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -61,7 +76,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -71,8 +86,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -86,7 +101,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -107,15 +122,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -125,8 +155,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -141,7 +171,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -151,8 +181,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -166,7 +196,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -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,14 +557,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -585,9 +679,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(); @@ -753,7 +847,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -765,7 +859,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(); @@ -787,19 +882,37 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; y += tid.x * out_vec_size + out_row; - for (int k = 0; k < in_vec_size; k += block_size) { + const int aligned_end = (in_vec_size / block_size) * block_size; + + for (int k = 0; k < aligned_end; k += block_size) { U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -812,9 +925,37 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } + // Tail block: handles K that is a multiple of 512 (dispatch gate) but not + // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for + // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; + // their result[row] stays at 0 and contributes 0 to the simd_sum below. + if (aligned_end < in_vec_size) { + bool in_bounds = + (aligned_end + static_cast(simd_lid) * values_per_thread) < + in_vec_size; + if (in_bounds) { + U sum = load_vector(x, x_thread); + + for (int row = 0; row < results_per_simdgroup; row++) { + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; + + U s = sl[0]; + U b = bl[0]; + result[row] += + qdot(wl, x_thread, s, b, sum); + } + } + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1193,6 +1334,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1201,7 +1343,7 @@ METAL_FUNC void qmm_t_impl( const device T* scales, const device T* biases, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -1226,7 +1368,7 @@ METAL_FUNC void qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1602,6 +1744,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( @@ -1642,7 +1785,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1660,6 +1803,7 @@ template < int group_size, const int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv( @@ -1946,7 +2090,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -1977,7 +2121,7 @@ template < const device T* scales [[buffer(1)]], const device T* biases [[buffer(2)]], const device T* x [[buffer(3)]], - device T* y [[buffer(4)]], + device float* y [[buffer(4)]], const constant int& K [[buffer(5)]], const constant int& N [[buffer(6)]], const constant int& M [[buffer(7)]], @@ -2005,7 +2149,7 @@ template < biases += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2336,7 +2480,7 @@ template < s_strides, b_strides, tid); - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2617,14 +2761,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; @@ -2728,7 +2881,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/Source/Cmlx/mlx-generated/metal/quantized_nax.h b/Source/Cmlx/mlx-generated/metal/quantized_nax.h index a34774992..95f26868f 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/quantized_nax.h @@ -31,15 +31,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -49,8 +64,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -64,7 +79,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -74,8 +89,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -89,7 +104,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -110,15 +125,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -128,8 +158,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -144,7 +174,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -154,8 +184,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -169,7 +199,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -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 == 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); + } + } - 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); @@ -487,14 +567,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -587,9 +689,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -715,9 +817,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h index b1546adb5..fce36a4cd 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h @@ -117,7 +117,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); @@ -172,6 +172,7 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], + const constant int64_t& column_offset [[buffer(11)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -193,7 +194,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT out_idx = gid.y + gsize.y * IdxT(gid.z); @@ -243,7 +244,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += out_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -311,7 +312,8 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], - const constant size_t& out_size [[buffer(11)]], + const constant int64_t& column_offset [[buffer(11)]], + const constant size_t& out_size [[buffer(12)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -336,7 +338,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT full_idx = gid.y + gsize.y * IdxT(gid.z); @@ -383,7 +385,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += full_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h index 604efa785..9a1e363e1 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h @@ -3,6 +3,7 @@ template [[kernel]] void init_reduce( device T* out [[buffer(0)]], - uint tid [[thread_position_in_grid]]) { - out[tid] = Op::init; + uint2 index [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + out[index.x + grid_dim.x * int64_t(index.y)] = Op::init; } diff --git a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal index 6797bdf3e..a27ad455b 100644 --- a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal @@ -2,10 +2,28 @@ // clang-format off #include "utils.h" +#include "sdpa_blocked.h" #include "sdpa_vector.h" using namespace metal; +instantiate_kernel( + "sdpa_blocked_scale_copy_bfloat16", + sdpa_blocked_scale_copy, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_scale_copy_float16", + sdpa_blocked_scale_copy, + float16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_bfloat16", + sdpa_blocked_causal_mask, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_float16", + sdpa_blocked_causal_mask, + float16_t) + // SDPA vector instantiations #define instantiate_sdpa_vector_aggregation(type, value_dim) \ instantiate_kernel( \ @@ -61,4 +79,17 @@ using namespace metal; instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) instantiate_sdpa_vector_heads(float16_t) - // clang-format on + +#define instantiate_sdpa_vector_tk1(type, dim) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_" #type "_" #dim "_" #dim "_tk1", \ + sdpa_vector_2pass_1, \ + type, \ + dim, \ + dim, \ + 1) + +instantiate_sdpa_vector_tk1(bfloat16_t, 256) +instantiate_sdpa_vector_tk1(float, 192) +instantiate_sdpa_vector_tk1(float, 256) +// clang-format on diff --git a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal new file mode 100644 index 000000000..b96809e52 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal @@ -0,0 +1,154 @@ +// Copyright © 2026 Apple Inc. + +#include "utils.h" + +using namespace metal; + +template +[[kernel]] void sdpa_vjp_odo( + const device InT* o [[buffer(0)]], // [B, H, qL, Dv] + const device InT* cot_o [[buffer(1)]], // [B, H, qL, Dv] + device float* odo [[buffer(2)]], // [B, H, qL] + constant int& qL [[buffer(3)]], + uint3 tpg [[thread_position_in_grid]], + uint simd_lane_id [[thread_index_in_simdgroup]]) { + const uint bh = tpg.z; + const int i = int(tpg.y); + + if (i >= qL) { + return; + } + + const size_t row = (size_t)bh * qL + i; + auto o_i = o + row * Dv; + auto co_i = cot_o + row * Dv; + + float acc = 0.0f; + for (int d = int(simd_lane_id); d < Dv; d += 32) { + acc += float(o_i[d]) * float(co_i[d]); + } + acc = simd_sum(acc); + + if (simd_lane_id == 0) { + odo[row] = acc; + } +} + +struct SDPAVJPTileParams { + int bq; + int bk; + int qL; + int i0; + int j0; + float scale; + int diag_off; + int causal; +}; + +template +[[kernel]] void sdpa_vjp_ds( + const device T* S [[buffer(0)]], // [BH, bq, bk] + const device T* dP [[buffer(1)]], // [BH, bq, bk] + const device float* lse [[buffer(2)]], // [B, H, qL] + const device float* odo [[buffer(3)]], // [B, H, qL] + device T* dS [[buffer(4)]], // [BH, bq, bk] + device T* P [[buffer(5)]], // [BH, bk, bq] + const constant SDPAVJPTileParams& p [[buffer(6)]], + uint3 gid [[thread_position_in_grid]]) { + int col = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (row >= p.bq || col >= p.bk) { + return; + } + + size_t sbase = (size_t(bh) * size_t(p.bq) + size_t(row)) * size_t(p.bk); + size_t qi = size_t(bh) * size_t(p.qL) + size_t(p.i0 + row); + + float l = lse[qi]; + float dlt = odo[qi]; + // Fully masked rows carry lse == -inf; exp(-inf - -inf) would be NaN. + bool dead = (l < 0.0f) && metal::isinf(l); + int lim = p.i0 + row + p.diag_off; + + float pv; + if (dead || (p.causal != 0 && (p.j0 + col) > lim)) { + pv = 0.0f; + } else { + pv = metal::fast::exp(static_cast(S[sbase + col]) * p.scale - l); + } + float dsv = pv * (static_cast(dP[sbase + col]) - dlt) * p.scale; + + dS[sbase + col] = static_cast(dsv); + P[sbase + col] = static_cast(pv); +} + +template +[[kernel]] void sdpa_vjp_reduce( + const device T* src [[buffer(0)]], + device float* acc [[buffer(1)]], + device T* out [[buffer(2)]], + const constant int& rows [[buffer(3)]], + const constant int& dim [[buffer(4)]], + const constant int& group [[buffer(5)]], + const constant int& acc_rows [[buffer(6)]], + const constant int& row_off [[buffer(7)]], + uint3 gid [[thread_position_in_grid]]) { + int c = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (c >= dim || row >= rows) { + return; + } + + size_t sbase = size_t(bh) * size_t(group) * size_t(rows) * size_t(dim); + float sum = 0.0f; + for (int g = 0; g < group; ++g) { + size_t si = sbase + + (size_t(g) * size_t(rows) + size_t(row)) * size_t(dim) + size_t(c); + sum += static_cast(src[si]); + } + + size_t ai = + (size_t(bh) * size_t(acc_rows) + size_t(row_off + row)) * size_t(dim) + + size_t(c); + float total = Accum ? acc[ai] + sum : sum; + acc[ai] = total; + out[ai] = static_cast(total); +} + +// Instantiations +#define instantiate_odo(in_type, dv) \ + instantiate_kernel( \ + "sdpa_vjp_odo_" #in_type "_" #dv, sdpa_vjp_odo, in_type, dv) + +#define instantiate_odo_shapes(in_type) \ + instantiate_odo(in_type, 64); \ + instantiate_odo(in_type, 72); \ + instantiate_odo(in_type, 80); \ + instantiate_odo(in_type, 96); \ + instantiate_odo(in_type, 128); \ + instantiate_odo(in_type, 192); \ + instantiate_odo(in_type, 256); \ + instantiate_odo(in_type, 512); \ + +instantiate_odo_shapes(bfloat16_t); +instantiate_odo_shapes(float16_t); +instantiate_odo_shapes(float); + +#define instantiate_sdpa_vjp_ds(tname, type) \ + instantiate_kernel("sdpa_vjp_ds_" #tname, sdpa_vjp_ds, type) + +instantiate_sdpa_vjp_ds(float, float); +instantiate_sdpa_vjp_ds(float16_t, float16_t); +instantiate_sdpa_vjp_ds(bfloat16_t, bfloat16_t); + +#define instantiate_sdpa_vjp_reduce(tname, type) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_add_" #tname, sdpa_vjp_reduce, type, true) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_set_" #tname, sdpa_vjp_reduce, type, false) + +instantiate_sdpa_vjp_reduce(float, float); +instantiate_sdpa_vjp_reduce(float16_t, float16_t); +instantiate_sdpa_vjp_reduce(bfloat16_t, bfloat16_t); diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h b/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h new file mode 100644 index 000000000..444d5e9ed --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h @@ -0,0 +1,45 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "defines.h" + +using namespace metal; + +template +[[kernel]] void sdpa_blocked_scale_copy( + const device T* q [[buffer(0)]], + device T* q_block [[buffer(1)]], + constant float& scale [[buffer(2)]], + constant int64_t* q_strides [[buffer(3)]], + constant int& heads [[buffer(4)]], + constant int& rows [[buffer(5)]], + constant int& head_dim [[buffer(6)]], + constant int& q_offset [[buffer(7)]], + uint3 elem [[thread_position_in_grid]]) { + int d = elem.x; + int row = elem.y; + int head = elem.z % heads; + int batch = elem.z / heads; + int64_t q_index = int64_t(batch) * q_strides[0] + + int64_t(head) * q_strides[1] + int64_t(q_offset + row) * q_strides[2] + d; + size_t q_block_index = (size_t(elem.z) * rows + row) * head_dim + d; + q_block[q_block_index] = q[q_index] * static_cast(scale); +} + +template +[[kernel]] void sdpa_blocked_causal_mask( + device T* scores [[buffer(0)]], + constant int& tile_rows [[buffer(1)]], + constant int& columns [[buffer(2)]], + uint2 elem [[thread_position_in_grid]]) { + int column = elem.x; + int score_row = elem.y; + int row = score_row % tile_rows; + if (column > row) { + size_t index = size_t(score_row) * columns + columns - tile_rows + column; + scores[index] = Limits::finite_min; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h index 4594f0389..a38265f55 100644 --- a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h +++ b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h @@ -216,6 +216,7 @@ template const int q_batch_head_idx = tid.x; const int q_seq_idx = tid.y; const int kv_head_idx = q_batch_head_idx / gqa_factor; + const int kv_seq_idx = simd_gid; const int o_offset = q_batch_head_idx * tpg.y + q_seq_idx; const int q_offset = query_transposed ? tpg.x * q_seq_idx + q_batch_head_idx : o_offset; @@ -226,11 +227,11 @@ template simd_lid * v_per_thread; if (bool_mask) { bmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } if (float_mask) { fmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } out += o_offset * V + simd_gid * v_per_thread; @@ -331,7 +332,7 @@ template } } -template +template [[kernel]] void sdpa_vector_2pass_1( const device T* queries [[buffer(0)]], const device T* keys [[buffer(1)]], @@ -414,8 +415,73 @@ template sum_exp_score = 1; } + const int stride_k_tok = blocks * int(k_seq_stride); + const int stride_v_tok = blocks * int(v_seq_stride); + const int stride_mask_tok = + (bool_mask || float_mask) ? blocks * mask_kv_seq_stride : 0; + // For each key - for (int i = block_idx; i < N; i += blocks) { + int i = block_idx; + for (; i + (TK - 1) * blocks < N; i += blocks * TK) { + // Prefetch: issue all TK K/V loads first so the HW can overlap them. + T k_reg[TK][qk_per_thread]; + T v_reg[TK][v_per_thread]; + for (int t = 0; t < TK; t++) { + const device T* kp = keys + t * stride_k_tok; + for (int j = 0; j < qk_per_thread; j++) { + k_reg[t][j] = kp[j]; + } + } + for (int t = 0; t < TK; t++) { + const device T* vp = values + t * stride_v_tok; + for (int j = 0; j < v_per_thread; j++) { + v_reg[t][j] = vp[j]; + } + } + + for (int t = 0; t < TK; t++) { + int ii = i + t * blocks; + bool use_key = true; + if (do_causal) { + use_key = ii <= (N - q_seq_len + int(q_seq_idx)); + } else if (bool_mask) { + use_key = bmask[t * stride_mask_tok]; + } else if (float_mask) { + use_key = (fmask[t * stride_mask_tok] >= Limits::finite_min); + } + if (use_key) { + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * static_cast(k_reg[t][j]); + } + score = simd_sum(score); + if (float_mask) { + score += static_cast(fmask[t * stride_mask_tok]); + } + + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * static_cast(v_reg[t][j]); + } + } + } + + keys += TK * stride_k_tok; + values += TK * stride_v_tok; + if (bool_mask) { + bmask += TK * stride_mask_tok; + } + if (float_mask) { + fmask += TK * stride_mask_tok; + } + } + + // Tail (fewer than TK remaining keys) + for (; i < N; i += blocks) { bool use_key = true; if (do_causal) { use_key = i <= (N - q_seq_len + int(q_seq_idx)); @@ -427,8 +493,8 @@ template if (use_key) { // Compute the i-th score U score = 0; - for (int i = 0; i < qk_per_thread; i++) { - score += q[i] * keys[i]; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * keys[j]; } score = simd_sum(score); @@ -445,19 +511,19 @@ template sum_exp_score = sum_exp_score * factor + exp_score; // Update the output accumulator - for (int i = 0; i < v_per_thread; i++) { - o[i] = o[i] * factor + exp_score * values[i]; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * values[j]; } } // Move the pointers to the next kv - keys += blocks * int(k_seq_stride); - values += blocks * int(v_seq_stride); + keys += stride_k_tok; + values += stride_v_tok; if (bool_mask) { - bmask += blocks * mask_kv_seq_stride; + bmask += stride_mask_tok; } if (float_mask) { - fmask += blocks * mask_kv_seq_stride; + fmask += stride_mask_tok; } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h index bc833474e..53a4017c3 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h @@ -14,6 +14,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; struct MaxOp { template @@ -77,6 +78,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -103,6 +105,12 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -532,6 +540,20 @@ template < Otile.template row_bin_op(sum_score); threadgroup_barrier(mem_flags::mem_none); + // Store the logsumexp for the backward pass + if (save_lse && sn == 0) { + using stile_t = decltype(Stile); + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + sm + i * stile_t::kFragRows; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = M_LN2_F * (max_score[i] + metal::log2(sum_score[i])); + } + } + } + // Store results O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; @@ -545,4 +567,4 @@ template < } else { Otile.template store(O, params->O_strides[2]); } -} +} \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h index bf2f026ba..43dc2db87 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h @@ -17,6 +17,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; template struct TransformScale { @@ -90,6 +91,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -117,6 +119,13 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + // [B, H, qL], contiguous. + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -475,6 +484,16 @@ template < Otile.template row_bin_op(rcp); + if (save_lse && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + i * otile_t::kFragRowsJump + sm; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = max_score[i] * M_LN2_F + metal::log(sum_score[i]); + } + } + } + // Store results O += tm * int(params->O_strides[2]); @@ -516,6 +535,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* LSE [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -903,6 +923,21 @@ template < rcp[i] = 1.f / sum_score[i]; } + if (save_lse && d_group == 0 && sn == 0) { + const int lse_row_base = int(tid.x) * BQ + tm; + const int lse_head_off = (int(tid.z) * params->H + int(tid.y)) * params->qL; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = lse_row_base + i * stile_t::kFragRowsJump + sm; + if (row < params->qL) { + LSE[lse_head_off + row] = (sum_score[i] == 0) + ? -INFINITY + : float(M_LN2_F * (max_score[i] + metal::log2(sum_score[i]))); + } + } + } + Otile.template row_bin_op(rcp); if (!align_Q && is_last_q) { diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h index f559596b7..f948ec069 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h @@ -46,7 +46,7 @@ implicit_gemm_conv_2d( using loader_a_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DInputBlockLoaderSmallChannels< @@ -84,7 +84,7 @@ implicit_gemm_conv_2d( // Weight loader using loader_b_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DWeightBlockLoaderSmallChannels< diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h index d0797ebea..de623ee4b 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h @@ -186,7 +186,7 @@ implicit_gemm_conv_2d_general( if (offset_n >= gemm_params->N) return; - short diff = gemm_params->N - offset_n; + int diff = gemm_params->N - offset_n; STEEL_PRAGMA_UNROLL for (int i = 0; i < mma_t::TM; i++) { diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h index 1284efe4a..9c134007e 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h @@ -47,8 +47,8 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -197,8 +197,8 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -384,6 +384,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -405,6 +406,7 @@ struct Conv2DWeightBlockLoader { weight_hw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -419,7 +421,7 @@ struct Conv2DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -481,9 +483,9 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; short kdil_d; short kdil_h; @@ -665,9 +667,9 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -887,6 +889,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -908,6 +911,7 @@ struct Conv3DWeightBlockLoader { weight_dhw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -922,7 +926,7 @@ struct Conv3DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h index a90e9c886..a2016fb50 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h @@ -71,6 +71,7 @@ struct Conv2DInputBlockLoaderSmallChannels { // Rows / strided reads within the block STEEL_CONST short n_rows = BROWS / TROWS; + static_assert(n_rows > 0, "BM must be at least tgp_size * vec_size / BK"); // Thread location indices const short thread_idx; @@ -233,6 +234,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -253,6 +255,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { params(params_), weight_hw(thread_idx % TCOLS), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -260,7 +263,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { if (bi >= BROWS || bj >= BCOLS) return; - if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -289,7 +292,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } } else { for (short i = 0; i < BROWS; i += TROWS) { - if (((read_n + i) < params->O)) { + if (((read_n + i) < N)) { STEEL_PRAGMA_UNROLL for (short j = 0; j < n_channels; j++) { dst[i * dst_ld + j] = curr_src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h index b44ff341e..1f4c66029 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h @@ -48,8 +48,8 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -248,8 +248,8 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const int start_row; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h index 5a43e2233..6814ef50d 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h @@ -79,8 +79,8 @@ template < k_start = segments[tid.z]; k_end = segments[tid.z + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid.z * params->batch_stride_d; // Prepare threadgroup mma operation diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h index 9cfece5c5..5731cb8c1 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h @@ -51,8 +51,8 @@ void segmented_mm_nax( k_start = segments[tid_s]; k_end = segments[tid_s + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid_s * params->batch_stride_d; constexpr short SM = BM / WM; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h new file mode 100644 index 000000000..28229a03c --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 0ec148ade..22009b9a8 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -41,15 +41,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -59,8 +74,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -74,7 +89,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -84,8 +99,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -99,7 +114,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -120,15 +135,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -138,8 +168,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -154,7 +184,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -164,8 +194,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -179,7 +209,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -209,13 +239,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) + @@ -311,13 +355,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) + @@ -408,11 +466,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); @@ -498,14 +570,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -598,9 +692,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(); @@ -766,7 +860,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -778,7 +872,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(); @@ -800,19 +895,37 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; y += tid.x * out_vec_size + out_row; - for (int k = 0; k < in_vec_size; k += block_size) { + const int aligned_end = (in_vec_size / block_size) * block_size; + + for (int k = 0; k < aligned_end; k += block_size) { U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -825,9 +938,37 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } + // Tail block: handles K that is a multiple of 512 (dispatch gate) but not + // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for + // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; + // their result[row] stays at 0 and contributes 0 to the simd_sum below. + if (aligned_end < in_vec_size) { + bool in_bounds = + (aligned_end + static_cast(simd_lid) * values_per_thread) < + in_vec_size; + if (in_bounds) { + U sum = load_vector(x, x_thread); + + for (int row = 0; row < results_per_simdgroup; row++) { + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; + + U s = sl[0]; + U b = bl[0]; + result[row] += + qdot(wl, x_thread, s, b, sum); + } + } + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1206,6 +1347,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1214,7 +1356,7 @@ METAL_FUNC void qmm_t_impl( const device T* scales, const device T* biases, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -1239,7 +1381,7 @@ METAL_FUNC void qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1615,6 +1757,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( @@ -1655,7 +1798,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1673,6 +1816,7 @@ template < int group_size, const int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv( @@ -1959,7 +2103,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -1990,7 +2134,7 @@ template < const device T* scales [[buffer(1)]], const device T* biases [[buffer(2)]], const device T* x [[buffer(3)]], - device T* y [[buffer(4)]], + device float* y [[buffer(4)]], const constant int& K [[buffer(5)]], const constant int& N [[buffer(6)]], const constant int& M [[buffer(7)]], @@ -2018,7 +2162,7 @@ template < biases += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2349,7 +2493,7 @@ template < s_strides, b_strides, tid); - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2630,14 +2774,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; @@ -2741,7 +2894,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/Source/Cmlx/mlx-generated/quantized_nax.cpp b/Source/Cmlx/mlx-generated/quantized_nax.cpp index 05277f64d..8f5d2297c 100644 --- a/Source/Cmlx/mlx-generated/quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/quantized_nax.cpp @@ -44,15 +44,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -62,8 +77,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -77,7 +92,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -87,8 +102,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -102,7 +117,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -123,15 +138,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -141,8 +171,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -157,7 +187,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -167,8 +197,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -182,7 +212,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -212,13 +242,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) + @@ -314,13 +358,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) + @@ -411,11 +469,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 == 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); + } + } - 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); @@ -500,14 +580,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -600,9 +702,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(); @@ -728,9 +830,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); diff --git a/Source/Cmlx/mlx-generated/reduce.cpp b/Source/Cmlx/mlx-generated/reduce.cpp index cac2d8de5..3cb5e0ca0 100644 --- a/Source/Cmlx/mlx-generated/reduce.cpp +++ b/Source/Cmlx/mlx-generated/reduce.cpp @@ -202,7 +202,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); @@ -257,6 +257,7 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], + const constant int64_t& column_offset [[buffer(11)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -278,7 +279,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT out_idx = gid.y + gsize.y * IdxT(gid.z); @@ -328,7 +329,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += out_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -396,7 +397,8 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], - const constant size_t& out_size [[buffer(11)]], + const constant int64_t& column_offset [[buffer(11)]], + const constant size_t& out_size [[buffer(12)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -421,7 +423,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT full_idx = gid.y + gsize.y * IdxT(gid.z); @@ -468,7 +470,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += full_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -492,8 +494,9 @@ template < template [[kernel]] void init_reduce( device T* out [[buffer(0)]], - uint tid [[thread_position_in_grid]]) { - out[tid] = Op::init; + uint2 index [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + out[index.x + grid_dim.x * int64_t(index.y)] = Op::init; } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/scatter.cpp b/Source/Cmlx/mlx-generated/scatter.cpp index b8bf7c108..ed991a65d 100644 --- a/Source/Cmlx/mlx-generated/scatter.cpp +++ b/Source/Cmlx/mlx-generated/scatter.cpp @@ -77,8 +77,8 @@ METAL_FUNC void scatter_impl( for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) { LocT out_idx = out_offset; for (int i = 0; i < NIDX; ++i) { - auto idx_loc = indices.row_contiguous[i] - ? ind_idx + LocT idx_loc = indices.row_contiguous[i] + ? LocT(ind_idx) : elem_to_loc( ind_idx, &indices.shapes[indices.ndim * i], @@ -89,7 +89,7 @@ METAL_FUNC void scatter_impl( out_idx += static_cast(idx_val) * static_cast(out_strides[ax]); } - auto upd_idx = ind_idx * static_cast(upd_size) + gid.x; + LocT upd_idx = ind_idx * static_cast(upd_size) + gid.x; if constexpr (!UPD_ROW_CONTIG) { upd_idx = elem_to_loc(upd_idx, upd_shape, upd_strides, upd_ndim); } diff --git a/Source/Cmlx/mlx-generated/steel_attention.cpp b/Source/Cmlx/mlx-generated/steel_attention.cpp index acd115969..31469e831 100644 --- a/Source/Cmlx/mlx-generated/steel_attention.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention.cpp @@ -1795,6 +1795,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; struct MaxOp { template @@ -1858,6 +1859,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -1884,6 +1886,12 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -2313,6 +2321,20 @@ template < Otile.template row_bin_op(sum_score); threadgroup_barrier(mem_flags::mem_none); + // Store the logsumexp for the backward pass + if (save_lse && sn == 0) { + using stile_t = decltype(Stile); + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + sm + i * stile_t::kFragRows; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = M_LN2_F * (max_score[i] + metal::log2(sum_score[i])); + } + } + } + // Store results O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; diff --git a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp index 3d4595d3c..222bb185d 100644 --- a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp @@ -1305,6 +1305,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; template struct TransformScale { @@ -1378,6 +1379,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -1405,6 +1407,13 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + // [B, H, qL], contiguous. + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -1763,6 +1772,16 @@ template < Otile.template row_bin_op(rcp); + if (save_lse && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + i * otile_t::kFragRowsJump + sm; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = max_score[i] * M_LN2_F + metal::log(sum_score[i]); + } + } + } + // Store results O += tm * int(params->O_strides[2]); @@ -1804,6 +1823,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* LSE [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -2191,6 +2211,21 @@ template < rcp[i] = 1.f / sum_score[i]; } + if (save_lse && d_group == 0 && sn == 0) { + const int lse_row_base = int(tid.x) * BQ + tm; + const int lse_head_off = (int(tid.z) * params->H + int(tid.y)) * params->qL; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = lse_row_base + i * stile_t::kFragRowsJump + sm; + if (row < params->qL) { + LSE[lse_head_off + row] = (sum_score[i] == 0) + ? -INFINITY + : float(M_LN2_F * (max_score[i] + metal::log2(sum_score[i]))); + } + } + } + Otile.template row_bin_op(rcp); if (!align_Q && is_last_q) { diff --git a/Source/Cmlx/mlx-generated/steel_conv.cpp b/Source/Cmlx/mlx-generated/steel_conv.cpp index 2af0e5605..85aca4026 100644 --- a/Source/Cmlx/mlx-generated/steel_conv.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv.cpp @@ -59,7 +59,7 @@ implicit_gemm_conv_2d( using loader_a_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DInputBlockLoaderSmallChannels< @@ -97,7 +97,7 @@ implicit_gemm_conv_2d( // Weight loader using loader_b_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DWeightBlockLoaderSmallChannels< diff --git a/Source/Cmlx/mlx-generated/steel_conv_general.cpp b/Source/Cmlx/mlx-generated/steel_conv_general.cpp index 0024a77f5..a3e98ac47 100644 --- a/Source/Cmlx/mlx-generated/steel_conv_general.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv_general.cpp @@ -71,8 +71,8 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -271,8 +271,8 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const int start_row; @@ -595,7 +595,7 @@ implicit_gemm_conv_2d_general( if (offset_n >= gemm_params->N) return; - short diff = gemm_params->N - offset_n; + int diff = gemm_params->N - offset_n; STEEL_PRAGMA_UNROLL for (int i = 0; i < mma_t::TM; i++) { diff --git a/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp b/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp index 848428230..1e1d37269 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp @@ -92,8 +92,8 @@ template < k_start = segments[tid.z]; k_end = segments[tid.z + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid.z * params->batch_stride_d; // Prepare threadgroup mma operation diff --git a/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp index 519087e1b..2e7adfb0c 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp @@ -64,8 +64,8 @@ void segmented_mm_nax( k_start = segments[tid_s]; k_end = segments[tid_s + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid_s * params->batch_stride_d; constexpr short SM = BM / WM; diff --git a/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp new file mode 100644 index 000000000..d9253fc8b --- /dev/null +++ b/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp @@ -0,0 +1,242 @@ +namespace mlx::core::metal { + +const char* steel_gemm_thin_nax() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/utils.cpp b/Source/Cmlx/mlx-generated/utils.cpp index c8cf785e0..047bcd788 100644 --- a/Source/Cmlx/mlx-generated/utils.cpp +++ b/Source/Cmlx/mlx-generated/utils.cpp @@ -681,10 +681,20 @@ constexpr complex_t operator*(complex_t a, complex_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}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -692,11 +702,7 @@ template < 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 = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md deleted file mode 100644 index 98620ba94..000000000 --- a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md +++ /dev/null @@ -1,30 +0,0 @@ ---- -name: Bug report -about: Create a report about a bug you've encountered -title: "[BUG] " -labels: '' -assignees: '' - ---- - -☑️ I understand it is strictly prohibited to use AI to write issues. - -**Describe the bug** -A clear and concise description of what the bug is. - -**To Reproduce** - -Include code snippet -```python - -``` - -**Expected behavior** -A clear and concise description of what you expected to happen. - -**Desktop (please complete the following information):** - - OS Version: [e.g. MacOS 14.1.2] - - Version [e.g. 0.7.0] - -**Additional context** -Add any other context about the problem here. diff --git a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml deleted file mode 100644 index 3ba13e0ce..000000000 --- a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml +++ /dev/null @@ -1 +0,0 @@ -blank_issues_enabled: false diff --git a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md deleted file mode 100644 index eb410efbb..000000000 --- a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md +++ /dev/null @@ -1,10 +0,0 @@ ---- -name: Other -about: Any other issue -title: '' -labels: '' -assignees: '' - ---- - -☑️ I understand it is strictly prohibited to use AI to write issues. diff --git a/Source/Cmlx/mlx/.github/actions/build-docs/action.yml b/Source/Cmlx/mlx/.github/actions/build-docs/action.yml deleted file mode 100644 index 4d4286e3c..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-docs/action.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: 'Build Documentation' -description: 'Build documentation' - -runs: - using: "composite" - steps: - - name: Setup machine - id: setup - uses: ./.github/actions/setup - with: - ccache-key: 'release' - ccache-save: false - - - name: Install dependencies - shell: bash - env: - CMAKE_ARGS: ${{ steps.setup.outputs.cmake-args }} - run: | - sudo apt-get install -y doxygen - uv pip install -r docs/requirements.txt - uv pip install . -v - - - name: Build documentation - shell: bash - run: | - cd docs - doxygen - make html O=-W - - - name: Create artifact tar - shell: bash - run: tar -cf artifact.tar -C docs --dereference build/html index.html - - # Do it manually because upload-pages-artifact requires gtar - - name: Upload artifact - id: upload-artifact - uses: actions/upload-artifact@v7 - with: - name: github-pages - path: artifact.tar - retention-days: 1 - if-no-files-found: error diff --git a/Source/Cmlx/mlx/.github/actions/build-macos/action.yml b/Source/Cmlx/mlx/.github/actions/build-macos/action.yml deleted file mode 100644 index 7bcd4e0d0..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-macos/action.yml +++ /dev/null @@ -1,39 +0,0 @@ -name: 'Build macOS' -description: 'Build C++ and Python binaries for testing on macOS' - -inputs: - cmake-args: - description: 'The args for generating CMake project' - required: true - macos-target: - description: 'The target macOS version to build for' - required: false - default: '26.2' - -runs: - using: 'composite' - steps: - - name: Install dependencies - shell: bash - run: uv pip install build setuptools - - - name: Build wheel - shell: bash - env: - DEBUG: 1 - CMAKE_ARGS: ${{ inputs.cmake-args }} - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: python -m build -w - - - name: Build CPP only - shell: bash - env: - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - if ${{ contains(inputs.cmake-args, 'CMAKE_BUILD_TYPE') }} ; then - cmake . -B build ${{ inputs.cmake-args }} - else - cmake . -B build ${{ inputs.cmake-args }} \ - -DCMAKE_BUILD_TYPE=Debug - fi - cmake --build build -j $(sysctl -n hw.physicalcpu) diff --git a/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml deleted file mode 100644 index c8f86420a..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml +++ /dev/null @@ -1,106 +0,0 @@ -name: 'Build wheel' -description: 'Build the Python wheels for release on all platforms' - -inputs: - cmake-args: - description: 'The args for generating CMake project' - required: true - build-frontend: - description: 'Build the frontend mlx package' - required: false - default: 'true' - build-backend: - description: 'Build the backend mlx-cpu/mlx-cuda/mlx-metal packages' - required: false - default: 'true' - macos-target: - description: 'The target macOS version to build for' - required: false - default: '26.2' - arch-tag: - description: 'Platform architecture tag' - required: false - default: |- - ${{ case(runner.arch == 'x64', 'x86_64', - runner.arch == 'x86', 'i686', - runner.arch == 'arm', 'armv7l', - runner.arch == 'arm64', 'aarch64', - 'unknown') - }} - -runs: - using: 'composite' - steps: - - name: Install dependencies - shell: bash - run: | - uv pip install build setuptools - if ${{ runner.os == 'Linux' }} ; then - uv pip install auditwheel patchelf - fi - mkdir -p wheelhouse - - - name: Build frontend package - if: inputs.build-frontend == 'true' - shell: bash - env: - CMAKE_ARGS: ${{ inputs.cmake-args }} - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - python setup.py clean --all - MLX_BUILD_FRONTEND_PACKAGE=1 python -m build -w - - - name: Post-process frontend package - if: inputs.build-frontend == 'true' - shell: bash - run: | - if ${{ runner.os == 'Linux' }} ; then - auditwheel repair dist/mlx-*.whl \ - --plat manylinux_2_35_${{ inputs.arch-tag }} \ - --exclude libmlx.so* \ - --only-plat - else - mv dist/mlx-*.whl wheelhouse/ - fi - - - name: Build backend package - if: inputs.build-backend == 'true' - shell: bash - env: - CMAKE_ARGS: ${{ inputs.cmake-args }} - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - python setup.py clean --all - MLX_BUILD_BACKEND_PACKAGE=1 python -m build -w - - - name: Post-process backend package - if: inputs.build-backend == 'true' - shell: bash - run: | - if ${{ runner.os == 'Linux' }} ; then - if [ -f dist/mlx_cpu*.whl ]; then - auditwheel repair dist/mlx_cpu*.whl \ - --plat manylinux_2_35_${{ inputs.arch-tag }} - fi - if [ -f dist/mlx_cuda*.whl ]; then - auditwheel repair dist/mlx_cuda*.whl \ - --plat manylinux_2_35_${{ inputs.arch-tag }} \ - --exclude libcublas* \ - --exclude libcuda* \ - --exclude libcudnn* \ - --exclude libcufft* \ - --exclude libcusolver* \ - --exclude libnccl* \ - --exclude libnvrtc* - fi - else - if [ -f dist/mlx_cpu*.whl ]; then - mv dist/mlx_cpu*.whl wheelhouse/ - fi - if [ -f dist/mlx_cuda*.whl ]; then - mv dist/mlx_cuda*.whl wheelhouse/ - fi - if [ -f dist/mlx_metal*.whl ]; then - mv dist/mlx_metal*.whl wheelhouse/ - fi - fi diff --git a/Source/Cmlx/mlx/.github/actions/setup/action.yml b/Source/Cmlx/mlx/.github/actions/setup/action.yml deleted file mode 100644 index 6ba85028b..000000000 --- a/Source/Cmlx/mlx/.github/actions/setup/action.yml +++ /dev/null @@ -1,270 +0,0 @@ -name: 'Setup environment' -description: 'Install dependencies for Linux and Windows builds' - -inputs: - toolkit: - description: 'Which toolkit to install' - required: false - default: 'cpu' - python-version: - description: 'Version of python to set up' - required: false - default: '3.12' - use-ccache: - description: 'Whether to enable ccache' - required: false - default: 'true' - ccache-key: - description: 'Extra key to use in the key of ccache' - required: false - default: 'test' - ccache-save: - description: 'Whether the ccache in this workflow will be saved' - required: false - default: 'auto' - ccache-toolkit: - description: 'Use this instead of inputs.toolkit in the ccache key' - required: false - default: '' - -outputs: - cmake-args: - description: 'The args for generating CMake project' - value: ${{ steps.cmake-args.outputs.cmakeArgs }} - -runs: - using: 'composite' - steps: - - name: Install Linux dependencies - if: runner.os == 'Linux' - shell: bash - run: | - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - gdb g++ ninja-build unzip \ - libblas-dev liblapack-dev liblapacke-dev \ - openmpi-bin openmpi-common libopenmpi-dev - - - name: Install macOS dependencies - if: runner.os == 'macOS' - shell: bash - run: | - brew update - brew trust aws/tap # suppress warning in github actions - brew install openmpi - xcodebuild -showComponent MetalToolchain - sysctl -a | grep machdep.cpu - - - name: Setup Windows environment - if: runner.os == 'Windows' - shell: cmd - run: | - :: Find out path to Visual Studio. - pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\" - for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x - popd - :: Import Visual stubs vars. - call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" ${{ runner.arch }} - :: Avoid using default TMP which uses short path and causes mismatch of CCACHE_BASEDIR. - set TMP=%RUNNER_TEMP% - set TEMP=%RUNNER_TEMP% - set UV_CACHE_DIR=%RUNNER_TEMP% - :: The cuda headers are downloaded. - set CCACHE_COMPILERCHECK=content - set CCACHE_SLOPPINESS=include_file_ctime,include_file_mtime - :: Export to all steps. - >>%GITHUB_ENV% set - - - uses: astral-sh/setup-uv@v8.2.0 - with: - enable-cache: false - quiet: true - - - name: Use ccache - if: inputs.use-ccache == 'true' - uses: hendrikmuhs/ccache-action@v1.2.23 - with: - 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'), - '300MB', - '60MB'), - case(startsWith(inputs.toolkit, 'cuda'), - 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. - update-package-index: false - - - name: Cache JIT-compiled CUDA kernels - if: runner.os == 'Linux' && startsWith(inputs.toolkit, 'cuda') - uses: actions/cache@v5 - with: - path: /tmp/mlx-ptx-cache - key: >- - ptx-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}- - ${{ hashFiles('mlx/backend/cuda/**') }} - - - name: Setup Python venv - if: runner.os != 'Windows' - shell: bash - run: | - uv venv --python ${{ inputs.python-version }} --managed-python - # Make sure all builds use the same cmake binary. - uv pip install cmake - source .venv/bin/activate - echo PATH=$PATH >> $GITHUB_ENV - # Search python packages in .venv - echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV - # Use PTX ccache for CUDA. - if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then - echo MLX_PTX_CACHE_DIR=/tmp/mlx-ptx-cache >> $GITHUB_ENV - fi - - - name: Setup Python venv (Windows) - if: runner.os == 'Windows' - shell: cmd - run: | - uv venv --python ${{ inputs.python-version }}${{ runner.arch == 'arm64' && '-arm64' || ''}} || exit /b - uv pip install cmake - call ".venv/Scripts/activate.bat" - >>%GITHUB_ENV% set - - - name: Install CUDA toolkit (Linux) - if: runner.os == 'Linux' && startsWith(inputs.toolkit, 'cuda') - shell: bash - env: - PACKAGES: | - { - "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6", - "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9", - "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0" - } - run: | - # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is - # Jetson specific. SBSA means Arm Server Base System Architecture. - ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }} - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - libnccl2 libnccl-dev \ - ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} - echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH - - - name: Install CUDA Toolkit (Windows) - if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') - shell: powershell - env: - INSTALLERS: | - { - "cuda-12.6": "https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_561.17_windows.exe", - "cuda-12.9": "https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda_12.9.1_576.57_windows.exe", - "cuda-13.0": "https://developer.download.nvidia.com/compute/cuda/13.0.2/local_installers/cuda_13.0.2_windows.exe" - } - PACKAGES: | - { - "cuda-12.6": ["cudart_12.6", "nvcc_12.6", "cublas_12.6", "cublas_dev_12.6", "cufft_12.6", "cufft_dev_12.6", "cusolver_12.6", "cusolver_dev_12.6", "nvrtc_12.6", "nvrtc_dev_12.6"], - "cuda-12.9": ["cudart_12.9", "nvcc_12.9", "cublas_12.9", "cublas_dev_12.9", "cufft_12.9", "cufft_dev_12.9", "cusolver_12.9", "cusolver_dev_12.9", "nvrtc_12.9", "nvrtc_dev_12.9"], - "cuda-13.0": ["cudart_13.0", "nvcc_13.0", "cublas_13.0", "cublas_dev_13.0", "cufft_13.0", "cufft_dev_13.0", "cusolver_13.0", "cusolver_dev_13.0", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"], - } - run: | - $ErrorActionPreference = "Stop" - $cudaUrl = "${{ fromJson(env.INSTALLERS)[inputs.toolkit] }}" - $cudaInstaller = "./install.exe" - - echo "Downloading '$cudaUrl'..." - $webClient = New-Object System.Net.WebClient - $webClient.DownloadFile($cudaUrl, $cudaInstaller) - echo "Downloaded: $cudaInstaller ($(([math]::Round((Get-Item $cudaInstaller).Length / 1MB, 2))) MB)" - - $args = "-s ${{ join(fromJson(env.PACKAGES)[inputs.toolkit], ' ') }}" - echo "Running '$cudaInstaller $args'..." - Start-Process -FilePath $cudaInstaller -ArgumentList "$args" -NoNewWindow -Wait - $cudaPath = (Resolve-Path "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*").path - echo "$cudaPath\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append - - - name: Install cuDNN (Windows) - if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') - id: cudnn - shell: powershell - env: - ARCHIVES: | - { - "cuda-12.6": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda12-archive.zip", - "cuda-12.9": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda12-archive.zip", - "cuda-13.0": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda13-archive.zip" - } - run: | - $ErrorActionPreference = "Stop" - $cudnnUrl = "${{ fromJson(env.ARCHIVES)[inputs.toolkit] }}" - $cudnnZip = "cudnn.zip" - - echo "Downloading '$cudnnUrl'..." - $webClient = New-Object System.Net.WebClient - $webClient.DownloadFile($cudnnUrl, $cudnnZip) - echo "Downloaded: $cudnnZip ($(([math]::Round((Get-Item $cudnnZip).Length / 1MB, 2))) MB)" - - echo "Extracing..." - Expand-Archive -Path $cudnnZip -DestinationPath cudnn-extracted - $cudnnDir = (Get-ChildItem -Path cudnn-extracted -Directory)[0].FullName - echo "cudnnDir=$($cudnnDir -replace '\\', '/')" | Out-File -FilePath $env:GITHUB_OUTPUT - - - name: Generate CMake args - id: cmake-args - shell: bash - run: | - cmakeArgs=( - "-G Ninja" - ) - if ${{ runner.os == 'macOS' }} ; then - if ${{ inputs.toolkit == 'cpu' }} ; then - cmakeArgs+=("-DMLX_BUILD_METAL=OFF") - else - cmakeArgs+=("-DMLX_BUILD_METAL=ON") - if ${{ inputs.toolkit == 'jit' }} ; then - cmakeArgs+=("-DMLX_METAL_JIT=ON") - fi - fi - fi - if ${{ runner.os == 'Windows' }} ; then - cmakeArgs+=( - "-DCMAKE_C_COMPILER=cl" - "-DCMAKE_CXX_COMPILER=cl" - "-DCMAKE_RC_COMPILER=rc" - ) - else - cmakeArgs+=( - "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" - ) - fi - if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then - cmakeArgs+=("-DMLX_BUILD_CUDA=ON") - # Some machines have no GPU. - if ! __nvcc_device_query ; then - # Use a fallback arch for testing. - cmakeArgs+=("-DMLX_CUDA_ARCHITECTURES=80") - # Can not build tests and stubs when the built executables can not run. - cmakeArgs+=("-DMLX_BUILD_TESTS=OFF") - cmakeArgs+=("-DMLX_BUILD_PYTHON_STUBS=OFF") - fi - echo - # Set cuDNN paths. - if ${{ runner.os == 'Windows' }} ; then - cmakeArgs+=( - "-DCUDNN_INCLUDE_PATH=${{ steps.cudnn.outputs.cudnnDir }}/include" - "-DCUDNN_LIBRARY_PATH=${{ steps.cudnn.outputs.cudnnDir }}/lib/x64" - ) - fi - else - cmakeArgs+=("-DMLX_BUILD_CUDA=OFF") - fi - # Pass to following steps. - IFS=" " - echo ${cmakeArgs[*]} - echo "cmakeArgs=${cmakeArgs[*]}" >> $GITHUB_OUTPUT diff --git a/Source/Cmlx/mlx/.github/actions/test-linux/action.yml b/Source/Cmlx/mlx/.github/actions/test-linux/action.yml deleted file mode 100644 index 7c5f43221..000000000 --- a/Source/Cmlx/mlx/.github/actions/test-linux/action.yml +++ /dev/null @@ -1,69 +0,0 @@ -name: 'Run tests' -description: 'Run Python and C++ tests on Linux' - -runs: - using: 'composite' - steps: - - name: Check GPU support - id: gpu-check - shell: bash - run: | - if __nvcc_device_query ; then - echo "good=true" >> $GITHUB_OUTPUT - else - echo "good=false" >> $GITHUB_OUTPUT - fi - echo - - - name: Run MPI tests - if: steps.gpu-check.outputs.good == 'false' - shell: bash - run: mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - - - name: Run distributed tests - if: steps.gpu-check.outputs.good == 'false' - shell: bash - run: | - mlx.launch --verbose -n 8 python python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) - if grep -Fq '[WARN]' stderr.log ; then - grep -F '[WARN]' stderr.log - echo "Distributed ring test failed"; - exit 1; - fi - - - name: Run Python tests - CPU - if: steps.gpu-check.outputs.good == 'false' - shell: bash - env: - DEVICE: cpu - run: python python/tests/run.py -v - - - name: Run Python tests - GPU - if: steps.gpu-check.outputs.good == 'true' - shell: bash - env: - DEVICE: gpu - run: python python/tests/run.py -v - - - name: Run CPP tests - CPU - shell: bash - env: - DEVICE: cpu - run: ./build/cpp/mlx/tests/tests - - - name: Run CPP tests - GPU - if: steps.gpu-check.outputs.good == 'true' - shell: bash - env: - DEVICE: gpu - run: ./build/cpp/mlx/tests/tests -sfe="*linalg_tests.cpp" - - - name: Show stack trace on crash - if: failure() - shell: bash - run: | - set +e - sleep 10 - if coredumpctl list; then - coredumpctl debug --debugger-arguments="-batch -ex 'thread apply all bt'" - fi diff --git a/Source/Cmlx/mlx/.github/actions/test-macos/action.yml b/Source/Cmlx/mlx/.github/actions/test-macos/action.yml deleted file mode 100644 index 55e5acc5c..000000000 --- a/Source/Cmlx/mlx/.github/actions/test-macos/action.yml +++ /dev/null @@ -1,76 +0,0 @@ -name: 'Run tests' -description: 'Run Python and C++ tests on macOS' - -inputs: - toolkit: - description: 'Which toolkit to test' - required: false - default: 'cpu' - -runs: - using: 'composite' - steps: - - name: Install tests dependencies - shell: bash - run: uv pip install tensorflow - - - name: Run tests - shell: bash - env: - METAL_DEBUG_ERROR_MODE: 0 - METAL_DEVICE_WRAPPER_TYPE: 1 - run: | - run_tests() { - for WHEEL in dist/*.whl; do - break - done - - echo "::group::Install $WHEEL" - uv pip install $WHEEL - echo "::endgroup::" - - echo "::group::Run CPP tests" - ./build/tests/tests - echo "::endgroup::" - - echo "::group::Run distributed tests" - # TODO: Re-enable for GPU after openmpi updating their libprrte - # dependency: https://github.com/openpmix/prrte/issues/2416 - if ${{ inputs.toolkit == 'cpu' }} ; then - mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib uv run python/tests/mpi_test_distributed.py - fi - mlx.launch --verbose -n 8 uv run python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) - if $(grep "\[WARN\]" stderr.log); then - echo "Distributed ring test failed" - exit 1 - fi - echo "::endgroup::" - - echo "::group::Run Python tests" - uv run python/tests/run.py -v - echo "::endgroup::" - - if ${{ inputs.toolkit != 'cpu' }} ; then - echo "::group::Build example extension" - pushd examples/extensions - uv pip install -r requirements.txt - uv run setup.py build_ext --inplace - uv run test.py - popd - echo "::endgroup::" - fi - } - - if [[ -d artifact ]] ; then - for SUBDIR in artifact/*; do - rm -rf build dist - mv $SUBDIR/* . - # Downloaded binaries all lose executable permission. - find build -type f -exec chmod 755 {} \; - # Path to metallib is hard-coded so move it besides executable. - mv build/mlx/backend/metal/kernels/mlx.metallib build/tests - run_tests - done - else - run_tests - fi diff --git a/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml deleted file mode 100644 index e06087650..000000000 --- a/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml +++ /dev/null @@ -1,49 +0,0 @@ -name: 'Test wheel' -description: 'Run tests with the built wheels' - -inputs: - toolkit: - description: 'Which toolkit to test' - required: false - default: 'cpu' - -runs: - using: 'composite' - steps: - - name: Get Python version - id: python - shell: bash - run: python -c "import sys; print(f'version={sys.version_info.major}.{sys.version_info.minor}')" >> $GITHUB_OUTPUT - - - name: Download frontend packages - uses: actions/download-artifact@v8 - with: - pattern: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ steps.python.outputs.version }} - path: wheelhouse - - - name: Download backend packages - uses: actions/download-artifact@v8 - with: - pattern: backend-${{ inputs.toolkit}}-${{ runner.os }}-${{ runner.arch }} - path: wheelhouse - - - name: Test local packages - shell: bash - run: | - # Fall back to PyPI if PyTorch index has issues. - uv pip install torch --torch-backend=cpu - uv pip install numpy - if ${{ inputs.toolkit == 'cpu' }} ; then - uv pip install wheelhouse/mlx_cpu*.whl - uv pip install wheelhouse/mlx-*.whl - elif ${{ startsWith(inputs.toolkit, 'cuda') }} ; then - uv pip install wheelhouse/mlx_cuda*.whl - uv pip install wheelhouse/mlx-*.whl - elif ${{ inputs.toolkit == 'metal' }} ; then - uv pip install wheelhouse/mlx_metal-*-macosx_26_0_arm64.whl - uv pip install wheelhouse/mlx-*-macosx_26_0_arm64.whl - else - echo "No matching backend wheel to install" - exit 1 - fi - uv run python/tests/run.py -v --failfast diff --git a/Source/Cmlx/mlx/.github/actions/test-windows/action.yml b/Source/Cmlx/mlx/.github/actions/test-windows/action.yml deleted file mode 100644 index 6d89c4e5b..000000000 --- a/Source/Cmlx/mlx/.github/actions/test-windows/action.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: 'Run tests' -description: 'Run Python and C++ tests on Windows' - -runs: - using: 'composite' - steps: - - name: Run Python tests - CPU - shell: bash - env: - DEVICE: cpu - run: uv run python/tests/run.py -v - - - name: Run CPP tests - CPU - shell: bash - env: - DEVICE: cpu - run: ./build/cpp/mlx/tests.exe -tce="*gguf*" diff --git a/Source/Cmlx/mlx/.github/dependabot.yml b/Source/Cmlx/mlx/.github/dependabot.yml deleted file mode 100644 index 5ace4600a..000000000 --- a/Source/Cmlx/mlx/.github/dependabot.yml +++ /dev/null @@ -1,6 +0,0 @@ -version: 2 -updates: - - package-ecosystem: "github-actions" - directory: "/" - schedule: - interval: "weekly" diff --git a/Source/Cmlx/mlx/.github/pull_request_template.md b/Source/Cmlx/mlx/.github/pull_request_template.md deleted file mode 100644 index a7f928e8b..000000000 --- a/Source/Cmlx/mlx/.github/pull_request_template.md +++ /dev/null @@ -1,2 +0,0 @@ -- ☑️ I understand it is strictly prohibited to use AI to write PR description -- AI usage disclosure: diff --git a/Source/Cmlx/mlx/.github/scripts/build-sanitizer-tests.sh b/Source/Cmlx/mlx/.github/scripts/build-sanitizer-tests.sh deleted file mode 100755 index 9c911890b..000000000 --- a/Source/Cmlx/mlx/.github/scripts/build-sanitizer-tests.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/bin/bash -set -ex - -export CMAKE_C_COMPILER=/usr/bin/clang -export CMAKE_CXX_COMPILER=/usr/bin/clang++ -BASE_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=DEBUG -DCMAKE_COMPILE_WARNING_AS_ERROR=ON" -if [[ "$(uname -s)" != "Darwin" ]]; then - BASE_CMAKE_ARGS+=" -DMLX_BUILD_METAL=OFF" -fi - -run_test() { - local sanitizer_name=$1 - local cmake_sanitizer_flag="-DUSE_${sanitizer_name}=ON" - echo " Running tests with: ${sanitizer_name}" - - case "$sanitizer_name" in - ASAN) - export ASAN_OPTIONS="detect_leaks=0" - ;; - UBSAN) - export UBSAN_OPTIONS="halt_on_error=0:print_stacktrace=1" - ;; - TSAN) - export TSAN_OPTIONS="" - ;; - esac - - rm -rf build - mkdir -p build - pushd build > /dev/null - - cmake .. ${BASE_CMAKE_ARGS} ${cmake_sanitizer_flag} - make -j $(nproc) - ./tests/tests - - popd > /dev/null - unset ${sanitizer_name}_OPTIONS -} - -sanitizer_arg=$(echo "$1" | tr '[:lower:]' '[:upper:]') - -if [[ "$sanitizer_arg" == "ASAN" || "$sanitizer_arg" == "UBSAN" || "$sanitizer_arg" == "TSAN" ]]; then - run_test "$sanitizer_arg" - echo " ${sanitizer_arg} test run completed successfully." -else - echo "Error: Invalid sanitizer '$1'. Please use one of: ASAN, UBSAN, TSAN." - exit 1 -fi diff --git a/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js b/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js deleted file mode 100644 index 296c8f2be..000000000 --- a/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js +++ /dev/null @@ -1,232 +0,0 @@ -#!/usr/bin/env node - -/** - * find-eligible-contributors.js - * - * Finds users who have opened pull requests on a GitHub repo, - * are NOT collaborators on that repo, and match: - * - fewer than 5 currently open PRs - * - more than 2 merged PRs - * - merge rate (merged / closed) > 50% - * - "low priority" PRs < 50% - * - * Usage: - * GITHUB_TOKEN=xxxx node find-eligible-contributors.js - */ - -const GITHUB_API = 'https://api.github.com'; - -function authHeaders(token) { - return { - Authorization: `Bearer ${token}`, - Accept: 'application/vnd.github+json', - 'X-GitHub-Api-Version': '2022-11-28', - 'User-Agent': 'find-eligible-contributors-script', - }; -} - -async function fetchAllPages(url, token, { stopWhen } = {}) { - const results = []; - let page = 1; - const perPage = 100; - const concurrentRequests = 5; - - while (true) { - const responses = await Promise.all( - Array.from({ length: concurrentRequests }, async (_, index) => { - const pageUrl = new URL(url); - pageUrl.searchParams.set('per_page', String(perPage)); - pageUrl.searchParams.set('page', String(page + index)); - - const res = await fetch(pageUrl, { headers: authHeaders(token) }); - if (!res.ok) { - const body = await res.text(); - throw new Error(`GitHub API error ${res.status} for ${pageUrl}: ${body}`); - } - return res.json(); - }) - ); - - let end = false; - for (const data of responses) { - if (!Array.isArray(data) || data.length === 0) { - end = true; - continue; - } - if (data.length < perPage) { - end = true; - } - - for (const item of data) { - if (stopWhen && stopWhen(item)) break; - results.push(item); - } - } - - if (end) { - break; - } - - page += concurrentRequests; - } - - return results; -} - -async function fetchPullRequests(owner, repo, token) { - const url = `${GITHUB_API}/repos/${owner}/${repo}/pulls?state=all&sort=created&direction=desc`; - - const oneYearAgo = new Date(); - oneYearAgo.setFullYear(oneYearAgo.getFullYear() - 1); - - return fetchAllPages(url, token, { - stopWhen: (pr) => new Date(pr.created_at) < oneYearAgo, - }); -} - -async function fetchCollaborators(owner, repo, token) { - const url = `${GITHUB_API}/repos/${owner}/${repo}/collaborators`; - const collaborators = await fetchAllPages(url, token); - return new Set(collaborators.map((c) => c.login)); -} - -/** - * Aggregate per-user PR stats from the list of PRs. - * Returns Map - */ -function aggregateUserStats(pullRequests) { - const stats = new Map(); - - for (const pr of pullRequests) { - const login = pr.user?.login; - if (!login) continue; - - if (!stats.has(login)) { - stats.set(login, { - open: 0, - closed: 0, - merged: 0, - total: 0, - lowPriority: 0, - mostRecentCreatedAt: null, - }); - } - const s = stats.get(login); - - s.total += 1; - if (pr.labels?.some((label) => label.name === 'low priority')) { - s.lowPriority += 1; - } - - if (pr.state === 'open') { - s.open += 1; - } else if (pr.state === 'closed') { - s.closed += 1; - if (pr.merged_at) { - s.merged += 1; - } - } - - const createdAt = new Date(pr.created_at); - if (!s.mostRecentCreatedAt || createdAt > s.mostRecentCreatedAt) { - s.mostRecentCreatedAt = createdAt; - } - } - - return stats; -} - -/** - * Apply eligibility filters: - * - fewer than 5 currently open PRs - * - more than 2 merged PRs - * - merge rate (merged / closed) > 50% - * - "low priority" PRs < 50% - * - not a collaborator - * - has an open PR or at least one PR within the past 6 weeks - */ -function filterEligibleUsers(stats, collaboratorLogins) { - const SIX_WEEKS_MS = 6 * 7 * 24 * 60 * 60 * 1000; - const cutoff = new Date((new Date).getTime() - SIX_WEEKS_MS); - const eligible = []; - - for (const [login, s] of stats.entries()) { - if (s.open >= 5) continue; - if (s.merged <= 2) continue; - if (collaboratorLogins.has(login)) continue; - if (s.open == 0 && s.mostRecentCreatedAt < cutoff) continue; - if ((s.merged / s.closed) < 0.5) continue; - if ((s.lowPriority / s.total) > 0.5) continue; - - eligible.push(login); - } - return eligible; -} - -/** - * Given a repo (owner/repo) and a GitHub token, returns the login names of - * eligible external contributors matching all filters. - * - * @param {Object} params - * @param {string} params.owner - Repo owner (user or org). - * @param {string} params.repo - Repo name. - * @param {string} params.token - GitHub token with repo read (and ideally push) access. - * @returns {Promise} - */ -async function findEligibleContributors({ owner, repo, token } = {}) { - if (!owner || !repo) { - throw new Error('findEligibleContributors requires both "owner" and "repo".'); - } - if (!token) { - throw new Error('findEligibleContributors requires a "token".'); - } - - const [pullRequests, collaboratorLogins] = await Promise.all([ - fetchPullRequests(owner, repo, token), - fetchCollaborators(owner, repo, token), - ]); - - const stats = aggregateUserStats(pullRequests); - return filterEligibleUsers(stats, collaboratorLogins); -} - -// ---- CLI entry point ---- - -function isRunAsCLI() { - return import.meta.url === `file://${process.argv[1]}`; -} - -async function runCLI() { - const [, , owner, repo] = process.argv; - if (!owner || !repo) { - console.error('Usage: node find-eligible-contributors.js '); - process.exit(1); - } - const token = process.env.GITHUB_TOKEN; - if (!token) { - console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); - process.exit(1); - } - - console.error(`Fetching pull requests and collaborators for ${owner}/${repo} in parallel...`); - const eligibleLogins = await findEligibleContributors({ owner, repo, token }); - - console.log(JSON.stringify(eligibleLogins, null, 2)); - console.error(`\n${eligibleLogins.length} user(s) match the criteria.`); -} - -if (isRunAsCLI()) { - runCLI().catch((err) => { - console.error('Fatal error:', err); - process.exit(1); - }); -} - -export { findEligibleContributors }; diff --git a/Source/Cmlx/mlx/.github/scripts/setup+build-cpp-linux-fedora-container.sh b/Source/Cmlx/mlx/.github/scripts/setup+build-cpp-linux-fedora-container.sh deleted file mode 100755 index e93757d47..000000000 --- a/Source/Cmlx/mlx/.github/scripts/setup+build-cpp-linux-fedora-container.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash -set -ex - -# [Setup] Install dependencies inside the container. -dnf update -y -dnf install -y \ - blas-devel \ - lapack-devel \ - openblas-devel \ - make \ - cmake \ - clang \ - git -dnf clean all - -# [C++] CI Build Sanity Check: Verifies code compilation, not for release. -export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" -export DEBUG=1 -export CMAKE_C_COMPILER=/usr/bin/clang -export CMAKE_CXX_COMPILER=/usr/bin/clang++ - -mkdir -p build -pushd build -cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG -make -j $(nproc) -./tests/tests -popd diff --git a/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js b/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js deleted file mode 100644 index 0f0f0817a..000000000 --- a/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js +++ /dev/null @@ -1,161 +0,0 @@ -#!/usr/bin/env node - -/** - * update-bypass-list.js - * - * Syncs a repo's PR interaction-limits bypass list with the current - * set of eligible external contributors. - * - * Usage: - * GITHUB_TOKEN=xxxx node update-bypass-list.js [--dry-run] - */ - -import { findEligibleContributors } from './find-eligible-contributors.js'; - -const GITHUB_API = 'https://api.github.com'; - -function authHeaders(token) { - return { - Authorization: `Bearer ${token}`, - Accept: 'application/vnd.github+json', - 'X-GitHub-Api-Version': '2026-03-10', - 'User-Agent': 'update-bypass-list-script', - 'Content-Type': 'application/json', - }; -} - -function bypassListUrl(owner, repo) { - return `${GITHUB_API}/repos/${owner}/${repo}/interaction-limits/pulls/bypass-list`; -} - -async function fetchBypassList(owner, repo, token) { - const url = bypassListUrl(owner, repo); - const res = await fetch(url, { headers: authHeaders(token) }); - - if (res.status === 404) { - return new Set(); - } - - if (!res.ok) { - const body = await res.text(); - throw new Error(`Failed to fetch bypass list (${res.status}): ${body}`); - } - - const data = await res.json(); - const logins = data.map((user) => user?.login).filter(Boolean); - - return new Set(logins); -} - -async function addUsersToBypassList(owner, repo, token, usernames) { - if (usernames.length === 0) return; - - const res = await fetch(bypassListUrl(owner, repo), { - method: 'PUT', - headers: authHeaders(token), - body: JSON.stringify({ users: usernames }), - }); - - if (!res.ok) { - const body = await res.text(); - throw new Error(`Failed to add users [${usernames.join(', ')}] (${res.status}): ${body}`); - } -} - -async function removeUsersFromBypassList(owner, repo, token, usernames) { - if (usernames.length === 0) return; - - const res = await fetch(bypassListUrl(owner, repo), { - method: 'DELETE', - headers: authHeaders(token), - body: JSON.stringify({ users: usernames }), - }); - - if (!res.ok) { - const body = await res.text(); - throw new Error(`Failed to remove users [${usernames.join(', ')}] (${res.status}): ${body}`); - } -} - -function diffLists(eligibleLogins, currentBypassLogins) { - const toAdd = [...eligibleLogins].filter((login) => !currentBypassLogins.has(login)); - const toRemove = [...currentBypassLogins].filter((login) => !eligibleLogins.has(login)); - return { toAdd, toRemove }; -} - -/** - * Update the bypass list to eligible contributors. - * - * @param {Object} params - * @param {string} params.owner - * @param {string} params.repo - * @param {string} params.token - * @param {boolean} [params.dryRun] - If true, only computes the diff, makes no API writes. - */ -async function updateBypassList({ owner, repo, token, dryRun = false }) { - console.error(`Getting eligible external contributors for ${owner}/${repo}...`); - const eligibleLogins = new Set(await findEligibleContributors({ owner, repo, token })); - console.error(`Found ${eligibleLogins.size} eligible user(s).`); - - console.error(`Fetching current bypass list...`); - const currentBypassLogins = await fetchBypassList(owner, repo, token); - console.error(`Current bypass list has ${currentBypassLogins.size} user(s).`); - - const { toAdd, toRemove } = diffLists(eligibleLogins, currentBypassLogins); - - console.error(`\nUsers to add (${toAdd.length}): ${toAdd.join(', ') || '(none)'}`); - console.error(`Users to remove (${toRemove.length}): ${toRemove.join(', ') || '(none)'}`); - - if (dryRun) { - console.error('\nDry run mode: no changes will be made.'); - return; - } - - try { - await addUsersToBypassList(owner, repo, token, toAdd); - if (toAdd.length > 0) console.error(`Added ${toAdd.length} user(s) to bypass list.`); - } catch (err) { - console.error(`Error adding users: ${err.message}`); - } - - try { - await removeUsersFromBypassList(owner, repo, token, toRemove); - if (toRemove.length > 0) console.error(`Removed ${toRemove.length} user(s) from bypass list.`); - } catch (err) { - console.error(`Error removing users: ${err.message}`); - } -} - -// ---- CLI entry point ---- - -function isRunAsCLI() { - return import.meta.url === `file://${process.argv[1]}`; -} - -async function runCLI() { - const args = process.argv.slice(2); - const dryRun = args.includes('--dry-run'); - const [owner, repo] = args.filter((a) => !a.startsWith('--')); - - if (!owner || !repo) { - console.error('Usage: node update-bypass-list.js [--dry-run]'); - process.exit(1); - } - - const token = process.env.GITHUB_TOKEN; - if (!token) { - console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); - process.exit(1); - } - - await updateBypassList({ owner, repo, token, dryRun }); -} - -if (isRunAsCLI()) { - runCLI().catch((err) => { - console.error('Fatal error:', err); - process.exit(1); - }); -} - -export { updateBypassList }; diff --git a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml b/Source/Cmlx/mlx/.github/workflows/build_and_test.yml deleted file mode 100644 index a4b5ddf8f..000000000 --- a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml +++ /dev/null @@ -1,205 +0,0 @@ -name: Build and Test - -on: - workflow_dispatch: - pull_request: - push: - branches: - - main - # For testing CI without starting a pull request: - - test/* - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} - -jobs: - check_lint: - name: Check Lint - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v7 - - uses: pre-commit/action@v3.0.1 - - build_and_test: - name: ${{ matrix.os }} (${{ matrix.toolkit }}, ${{ matrix.arch }}) - if: github.repository == 'ml-explore/mlx' - needs: check_lint - strategy: - fail-fast: false - matrix: - os: ['Linux', 'Windows'] - arch: ['x86_64', 'aarch64'] - toolkit: ['cpu', 'cuda-12.6', 'cuda-12.9', 'cuda-13.0'] - exclude: - # CUDA does not support Windows on arm. - - os: 'Windows' - arch: 'aarch64' - toolkit: 'cuda-12.6' - - os: 'Windows' - arch: 'aarch64' - toolkit: 'cuda-12.9' - - os: 'Windows' - arch: 'aarch64' - toolkit: 'cuda-13.0' - # CUDA 12.6 does not compile with CUTLASS on Windows. - - os: 'Windows' - arch: 'x86_64' - toolkit: 'cuda-12.6' - runs-on: |- - ${{ case(matrix.os == 'Windows', - case(matrix.arch == 'aarch64', 'windows-11-arm', - startsWith(matrix.toolkit, 'cuda'), 'windows-2022-large', - 'windows-2022'), - case(matrix.arch == 'x86_64' && startsWith(matrix.toolkit, 'cuda'), 'gpu-t4-4-core', - matrix.arch == 'aarch64', 'ubuntu-22.04-arm', - 'ubuntu-22.04')) - }} - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - id: setup - with: - toolkit: ${{ matrix.toolkit }} - - uses: ./.github/actions/build - with: - cmake-args: ${{ steps.setup.outputs.cmake-args }} - # For MSVC, Ninja/Release is the only config supported by ccache. - debug: ${{ matrix.os != 'Windows' }} - - name: Check generated Python stubs with ty - if: matrix.os == 'Linux' && matrix.arch == 'x86_64' && matrix.toolkit == 'cpu' - run: uvx ty check python/mlx/core - - uses: ./.github/actions/test-linux - if: matrix.os == 'Linux' && (matrix.toolkit == 'cpu' || matrix.arch == 'x86_64') - - uses: ./.github/actions/test-windows - if: matrix.os == 'Windows' && matrix.toolkit == 'cpu' - - mac_build: - name: macOS (${{ matrix.macos-target }}, ${{ matrix.toolkit }}) - if: github.repository == 'ml-explore/mlx' - strategy: - fail-fast: false - matrix: - macos-target: ['14.0', '15.0', '26.2'] - toolkit: ['cpu', 'metal', 'jit'] - runs-on: 'macos-26-xlarge' - needs: check_lint - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - id: setup - with: - toolkit: ${{ matrix.toolkit }} - ccache-key: 'test-${{ matrix.macos-target }}' - # Use the compile cache of metal build for all builds. - ccache-save: ${{ matrix.toolkit == 'metal' }} - ccache-toolkit: 'metal' - - uses: ./.github/actions/build-macos - with: - cmake-args: ${{ steps.setup.outputs.cmake-args }} - macos-target: ${{ matrix.macos-target }} - - uses: ./.github/actions/test-macos - if: matrix.toolkit == 'cpu' - with: - toolkit: 'cpu' - - uses: actions/upload-artifact@v7 - if: matrix.toolkit != 'cpu' - with: - name: mlx-${{ matrix.toolkit }}-macos${{ matrix.macos-target }} - path: | - dist/mlx-*.whl - build/mlx/backend/metal/kernels/mlx.metallib - build/tests/tests - if-no-files-found: error - - mac_test: - name: Test macOS (${{ matrix.toolkit }}) - if: github.repository == 'ml-explore/mlx' - runs-on: [self-hosted, macos] - needs: mac_build - strategy: - fail-fast: false - matrix: - toolkit: ['metal', 'jit'] - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - with: - toolkit: ${{ matrix.toolkit }} - use-ccache: false - - uses: actions/download-artifact@v8 - with: - path: artifact - pattern: mlx-${{ matrix.toolkit }}-* - - run: ls -lhR artifact - - uses: ./.github/actions/test-macos - with: - toolkit: ${{ matrix.toolkit }} - - build_documentation: - name: Build Documentation - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22.04 - needs: check_lint - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/build-docs - - linux_sanitizer_build_and_test: - name: Linux Sanitizer Tests (${{ matrix.sanitizer }}) - needs: check_lint - strategy: - fail-fast: false - matrix: - sanitizer: [ASAN, UBSAN] - # todo 12/16/2025: enable TSAN later + consider enabling ASAN for GPU backend tests. - # sanitizer: [ASAN, UBSAN, TSAN] - runs-on: ubuntu-22.04-arm - steps: - - name: Checkout code - uses: actions/checkout@v7 - - - name: Install Dependencies - run: | - export DEBIAN_FRONTEND=noninteractive - sudo apt-get update -y - sudo apt-get install -y \ - build-essential \ - libblas-dev \ - liblapacke-dev \ - libopenblas-dev \ - cmake \ - clang \ - git - sudo apt-get clean - sudo rm -rf /var/lib/apt/lists/* - - - name: Linux Build and Test with ${{ matrix.sanitizer }} - run: | - bash .github/scripts/build-sanitizer-tests.sh ${{ matrix.sanitizer }} - - linux_fedora_build_cpp: - name: Linux Fedora (${{ matrix.arch }}) - needs: check_lint - strategy: - fail-fast: false - matrix: - include: - - host: ubuntu-22.04 - arch: x86_64 - - host: ubuntu-22.04-arm - arch: aarch64 - - runs-on: ${{ matrix.host }} - container: - image: fedora:42 - steps: - - name: Checkout code - uses: actions/checkout@v7 - - - name: CPP Build Test - No Release - run: | - bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh diff --git a/Source/Cmlx/mlx/.github/workflows/documentation.yml b/Source/Cmlx/mlx/.github/workflows/documentation.yml deleted file mode 100644 index fad164799..000000000 --- a/Source/Cmlx/mlx/.github/workflows/documentation.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: Documentation - -on: - workflow_dispatch: - -permissions: - contents: read - -jobs: - build: - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/build-docs - - deploy: - needs: build - permissions: - pages: write - id-token: write - runs-on: ubuntu-latest - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v5 diff --git a/Source/Cmlx/mlx/.github/workflows/release.yml b/Source/Cmlx/mlx/.github/workflows/release.yml deleted file mode 100644 index 0445da036..000000000 --- a/Source/Cmlx/mlx/.github/workflows/release.yml +++ /dev/null @@ -1,321 +0,0 @@ -name: 'Release build' -description: 'Build Python wheels for nightly or offical releases' - -on: - push: - tags: - - 'v*' - workflow_dispatch: - inputs: - publish: - description: 'Publish to PyPI' - required: false - type: boolean - dev-release: - description: 'Development release (DEV_RELEASE=1)' - required: false - type: boolean - schedule: - - cron: 33 6 * * * - -# In jobs we must use |*publish| instead of |inputs.publish| because we can not -# set default value for workflow_dispatch inputs reliably. -env: - publish: &publish ${{ inputs.publish || github.event_name == 'push' }} - pypi-env: &pypi-env ${{ (inputs.publish || github.event_name == 'push') && 'pypi' || 'dry-run' }} - -permissions: - contents: read - -jobs: - build_documentation: - name: Build documentation - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/build-docs - - deploy_documentation: - name: Deploy documentation - if: *publish - needs: build_documentation - permissions: - pages: write - id-token: write - runs-on: ubuntu-latest - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v5 - - build_frontend: - name: ${{ matrix.os }} (python-${{ matrix.python-version }}, ${{ matrix.arch }}) - strategy: - matrix: - os: ['Linux', 'Windows'] - arch: ['x86_64', 'aarch64'] - python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t'] - # There is no cp310 binary for Windows on arm. - exclude: - - os: 'Windows' - arch: 'aarch64' - python-version: '3.10' - runs-on: |- - ${{ case(matrix.os == 'Windows', - case(matrix.arch == 'aarch64', 'windows-11-arm', - 'windows-2022'), - case(matrix.arch == 'aarch64', 'ubuntu-22.04-arm', - 'ubuntu-22.04')) - }} - env: &build-env - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev-release && 1 || 0 }} - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - id: setup - with: - python-version: ${{ matrix.python-version }} - ccache-key: 'release' - ccache-save: false - - uses: ./.github/actions/build-wheel - with: - cmake-args: ${{ steps.setup.outputs.cmake-args }} - build-backend: false - - run: unzip -l wheelhouse/*.whl - - uses: actions/upload-artifact@v7 - with: - name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} - path: wheelhouse/mlx-*.whl - if-no-files-found: error - - build_backend: - name: ${{ matrix.os }} (${{ matrix.toolkit }}, ${{ matrix.arch }}) - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - os: ['Linux', 'Windows'] - arch: ['x86_64', 'aarch64'] - toolkit: ['cpu', 'cuda-12.9', 'cuda-13.0'] - exclude: - # CUDA does not support Windows on arm. - - os: 'Windows' - arch: 'aarch64' - toolkit: 'cuda-12.9' - - os: 'Windows' - arch: 'aarch64' - toolkit: 'cuda-13.0' - runs-on: |- - ${{ case(matrix.os == 'Windows', - case(matrix.arch == 'aarch64', 'windows-11-arm', - 'windows-2022-large'), - case(matrix.arch == 'aarch64', 'ubuntu-22-large-arm', - 'ubuntu-22-large')) - }} - env: *build-env - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - id: setup - with: - toolkit: ${{ matrix.toolkit }} - ccache-key: 'release' - - uses: ./.github/actions/build-wheel - with: - cmake-args: ${{ steps.setup.outputs.cmake-args }} - build-frontend: false - - run: unzip -l wheelhouse/*.whl - - uses: actions/upload-artifact@v7 - with: - name: backend-${{ matrix.toolkit }}-${{ runner.os }}-${{ runner.arch }} - path: wheelhouse/*.whl - if-no-files-found: ignore - - build_mac_wheels: - name: macOS (python-${{ matrix.python-version }}) - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - python-version: *pyver - runs-on: 'macos-26-xlarge' - env: *build-env - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - id: setup - with: - toolkit: 'metal' - python-version: ${{ matrix.python-version }} - ccache-key: 'release' - ccache-save: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 14 package - uses: ./.github/actions/build-wheel - with: - macos-target: '14.0' - cmake-args: ${{ steps.setup.outputs.cmake-args }} - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 15 package - uses: ./.github/actions/build-wheel - with: - macos-target: '15.0' - cmake-args: ${{ steps.setup.outputs.cmake-args }} - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 26 package - uses: ./.github/actions/build-wheel - with: - macos-target: '26.2' - cmake-args: ${{ steps.setup.outputs.cmake-args }} - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Display content of the wheels - run: for WHEEL in wheelhouse/*.whl; do unzip -l $WHEEL; echo; done - - name: Upload frontend packages - uses: actions/upload-artifact@v7 - with: - name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} - path: wheelhouse/mlx-*.whl - if-no-files-found: error - - name: Upload backend packages - if: matrix.python-version == '3.10' - uses: actions/upload-artifact@v7 - with: - name: backend-metal-${{ runner.os }}-${{ runner.arch }} - path: wheelhouse/mlx_metal-*.whl - if-no-files-found: error - - test_wheel: - name: Test (${{ matrix.os }}, ${{ matrix.toolkit }}, ${{ matrix.arch }}) - if: github.repository == 'ml-explore/mlx' - needs: [build_frontend, build_backend, build_mac_wheels] - strategy: - matrix: - os: ['Linux', 'Windows'] - arch: ['aarch64'] - toolkit: ['cpu'] - include: - - os: 'Linux' - arch: 'x86_64' - toolkit: 'cpu' - - os: 'Linux' - arch: 'x86_64' - toolkit: 'cuda-12.9' - - os: 'Linux' - arch: 'x86_64' - toolkit: 'cuda-13.0' - - os: 'Windows' - arch: 'x86_64' - toolkit: 'cpu' - - os: 'macOS' - arch: 'aarch64' - toolkit: 'metal' - runs-on: |- - ${{ case(matrix.os == 'Windows', case(matrix.arch == 'aarch64', 'windows-11-arm', - 'windows-2022'), - matrix.os == 'macOS', fromJson('["self-hosted","macos"]'), - case(matrix.arch == 'x86_64' && startsWith(matrix.toolkit, 'cuda'), 'gpu-t4-4-core', - matrix.arch == 'aarch64', 'ubuntu-22.04-arm', - 'ubuntu-22.04')) - }} - steps: - - uses: actions/checkout@v7 - - uses: ./.github/actions/setup - with: - toolkit: ${{ matrix.toolkit }} - use-ccache: false - - uses: ./.github/actions/test-wheel - with: - toolkit: ${{ matrix.toolkit }} - - pypi-publish-frontend: - name: Publish mlx to PyPI - runs-on: ubuntu-latest - needs: [test_wheel] - permissions: - id-token: write - environment: - name: *pypi-env - url: https://pypi.org/p/mlx - steps: - - uses: actions/download-artifact@v8 - with: - pattern: frontend-* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: *publish - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-cuda: - name: Publish mlx-cuda to PyPI - runs-on: ubuntu-latest - needs: [test_wheel] - permissions: - id-token: write - environment: - name: *pypi-env - url: https://pypi.org/p/mlx-cuda - steps: - - uses: actions/download-artifact@v8 - with: - pattern: backend-cuda* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: *publish - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-cpu: - name: Publish mlx-cpu to PyPI - runs-on: ubuntu-latest - needs: [test_wheel] - permissions: - id-token: write - environment: - name: *pypi-env - url: https://pypi.org/p/mlx-cpu - steps: - - uses: actions/download-artifact@v8 - with: - pattern: backend-cpu-* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: *publish - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-metal: - name: Publish mlx-metal to PyPI - runs-on: ubuntu-latest - needs: [test_wheel] - permissions: - id-token: write - environment: - name: *pypi-env - url: https://pypi.org/p/mlx-metal - steps: - - uses: actions/download-artifact@v8 - with: - pattern: backend-metal-* - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: *publish - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ diff --git a/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml b/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml deleted file mode 100644 index ec4025893..000000000 --- a/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml +++ /dev/null @@ -1,25 +0,0 @@ -name: 'Update bypass list' -description: 'Update interaction-limits bypass list with eligible contributors' - -on: - workflow_dispatch: - pull_request_target: - types: - - closed - schedule: - - cron: 33 6 * * * - -permissions: - contents: write - -jobs: - update_bypass_list: - name: 'Update bypass list' - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v7 - - env: - # Can only use PAT with permissions: admin:org_hook, public_repo - GITHUB_TOKEN: ${{ secrets.ZCBENZ_TOKEN_UPDATE_BYPASS_LIST }} - run: node .github/scripts/update-bypass-list.js ml-explore mlx diff --git a/Source/Cmlx/mlx/.gitignore b/Source/Cmlx/mlx/.gitignore index f5d859ff3..015940e75 100644 --- a/Source/Cmlx/mlx/.gitignore +++ b/Source/Cmlx/mlx/.gitignore @@ -12,6 +12,7 @@ __pycache__/ # Distribution / packaging /build/ +/examples/**/build/ python/mlx/core python/mlx/share python/mlx/include @@ -79,4 +80,5 @@ uv.lock .cache/ # vim *.swp -wip \ No newline at end of file +.idea +wip diff --git a/Source/Cmlx/mlx/.pre-commit-config.yaml b/Source/Cmlx/mlx/.pre-commit-config.yaml index 859d0f660..bff933ce1 100644 --- a/Source/Cmlx/mlx/.pre-commit-config.yaml +++ b/Source/Cmlx/mlx/.pre-commit-config.yaml @@ -10,19 +10,15 @@ repos: hooks: - id: clang-format files: \.(h|cpp)$ -# Using this mirror lets us use mypyc-compiled black, which is about 2x faster -- repo: https://github.com/psf/black-pre-commit-mirror - rev: 26.1.0 - hooks: - - id: black - -- repo: https://github.com/pycqa/isort - rev: 7.0.0 - hooks: - - id: isort - args: - - --profile=black - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 hooks: - id: cmake-format +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.16.10 + hooks: + - id: ruff-check + types_or: [python, pyi] + args: [--fix] + - id: ruff-format + types_or: [python, pyi] diff --git a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md index a97754922..e8cb16af5 100644 --- a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md +++ b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md @@ -12,6 +12,7 @@ MLX was developed with contributions from the following individuals: - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example. - Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`. - Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented pooling layers and ``Upsample``. +- Aditya Dhulipala: Added `mlx.core.linalg.pinv`. - Hinrik Snær Guðmundsson: Added `atleast_1d`, `atleast_2d`, `atleast_3d` ops. - Luca Arnaboldi: Added `Ceil` and `Floor` ops; implemented pickling, copy and deepcopy for mlx arrays. - Brian Keene & Atila Orhon, with Argmax Inc.: Added `fast.scaled_dot_product_attention` diff --git a/Source/Cmlx/mlx/CMakeLists.txt b/Source/Cmlx/mlx/CMakeLists.txt index 01b329b2f..3f2330483 100644 --- a/Source/Cmlx/mlx/CMakeLists.txt +++ b/Source/Cmlx/mlx/CMakeLists.txt @@ -396,7 +396,7 @@ if(MLX_BUILD_PYTHON_BINDINGS) FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v3.0.1 + GIT_TAG v3.1.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) diff --git a/Source/Cmlx/mlx/benchmarks/numpy/time_utils.py b/Source/Cmlx/mlx/benchmarks/numpy/time_utils.py index 163f007dc..738bd04d5 100644 --- a/Source/Cmlx/mlx/benchmarks/numpy/time_utils.py +++ b/Source/Cmlx/mlx/benchmarks/numpy/time_utils.py @@ -13,7 +13,7 @@ def time_fn(fn, *args): num_iters = 100 tic = time.perf_counter() for _ in range(num_iters): - x = fn(*args) + _x = fn(*args) toc = time.perf_counter() msec = 1e3 * (toc - tic) / num_iters diff --git a/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemm.py b/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemm.py index ee358a95d..bdf6c46ff 100644 --- a/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemm.py +++ b/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemm.py @@ -2,7 +2,6 @@ import argparse import math -import os import subprocess import time @@ -148,7 +147,8 @@ def bench_shape(B, M, N, K, np_dtype, transpose="nn"): if not np.allclose(c_mlx, c_npy.astype(np_dtype), atol=atol): print( - f"Failed at {(B, M, N, K)} [transpose = {transpose}] with max(|a - b|) = {np.max(np.abs(c_npy - c_mlx))}" + f"Failed at {(B, M, N, K)} [transpose = {transpose}] with max(|a - b|) = " + f"{np.max(np.abs(c_npy - c_mlx))}" ) return time_mlx, time_torch @@ -185,7 +185,8 @@ def get_gflop_count(B, M, N, K): diff = gflops_mx / gflops_pt - 1.0 print( - f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%" + f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, " + f"{gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%" ) if gflops_pt >= 2.0 * gflops_mx: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemv.py b/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemv.py index 3cfc5eba4..94b5d169c 100644 --- a/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemv.py +++ b/Source/Cmlx/mlx/benchmarks/python/blas/bench_gemv.py @@ -114,7 +114,8 @@ def bench_lens(in_vec_len, out_vec_len, np_dtype, transpose=False): if not np.allclose(c_mlx, c_npy, atol=2e-5): print( - f"Failed at {shape_mat} [transpose = {transpose}] with max(|a - b|) = {np.max(np.abs(c_npy - c_mlx))}" + f"Failed at {shape_mat} [transpose = {transpose}] with max(|a - b|) = " + f"{np.max(np.abs(c_npy - c_mlx))}" ) return time_mlx, time_torch diff --git a/Source/Cmlx/mlx/benchmarks/python/block_masked_mm_bench.py b/Source/Cmlx/mlx/benchmarks/python/block_masked_mm_bench.py index bcfeb1c4f..58abd3fc3 100644 --- a/Source/Cmlx/mlx/benchmarks/python/block_masked_mm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/block_masked_mm_bench.py @@ -101,7 +101,10 @@ def main(): "1024x1024x1024x32x0.0," "1024x1024x1024x32x0.9" ), - help="Comma-separated MxNxKxBSxSparsity list. Sparsity=fraction of blocks zeroed.", + help=( + "Comma-separated MxNxKxBSxSparsity list. " + "Sparsity=fraction of blocks zeroed." + ), ) parser.add_argument( "--dtype", diff --git a/Source/Cmlx/mlx/benchmarks/python/comparative/bench_mlx.py b/Source/Cmlx/mlx/benchmarks/python/comparative/bench_mlx.py index 4e6ba04f8..dff927186 100644 --- a/Source/Cmlx/mlx/benchmarks/python/comparative/bench_mlx.py +++ b/Source/Cmlx/mlx/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 ), @@ -175,13 +189,6 @@ def relu(x): mx.eval(y) -def leaky_relu(x: mx.array): - y = x - for i in range(100): - y = nn.leaky_relu(y) - mx.eval(y) - - def prelu(x: mx.array): y = x for i in range(100): @@ -189,13 +196,6 @@ def prelu(x: mx.array): mx.eval(y) -def softplus(x: mx.array): - y = x - for i in range(100): - y = nn.softplus(y) - mx.eval(y) - - def mish(x: mx.array): y = x for i in range(100): @@ -287,7 +287,6 @@ def rope(x): *_, N, D = x.shape ys = [] for i in range(10): - shape = x.shape x = mx.reshape(x, (-1, N, D)) positions = mx.arange(N) freqs = mx.exp(mx.arange(0.0, D // 2) / math.log(10000 / (D // 2 - 1))) diff --git a/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py b/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py index aa81369fe..8fde17745 100644 --- a/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py +++ b/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py @@ -48,15 +48,23 @@ def search(x): def make_predicate(positive_filter, negative_filter): if positive_filter is not None: positive_filter_search = make_regex_search(positive_filter) - positive_filter = lambda x: all(positive_filter_search(x)) + + def positive_filter(x): + return all(positive_filter_search(x)) else: - positive_filter = lambda x: True + + def positive_filter(x): + return True if negative_filter is not None: negative_filter_search = make_regex_search(negative_filter) - negative_filter = lambda x: not any(negative_filter_search(x)) + + def negative_filter(x): + return not any(negative_filter_search(x)) else: - negative_filter = lambda x: True + + def negative_filter(x): + return True def predicate(x): return positive_filter(x) and negative_filter(x) @@ -83,13 +91,19 @@ def predicate(x): _filter = make_predicate(args.filter, args.negative_filter) if args.mlx_dtypes: - compare_filtered = lambda x: ( - compare_mlx_dtypes(x.split() + rest, args.mlx_dtypes[0], args.mlx_dtypes[1]) - if _filter(x) - else None - ) + + def compare_filtered(x): + return ( + compare_mlx_dtypes( + x.split() + rest, args.mlx_dtypes[0], args.mlx_dtypes[1] + ) + if _filter(x) + else None + ) else: - compare_filtered = lambda x: compare(x.split() + rest) if _filter(x) else None + + def compare_filtered(x): + return compare(x.split() + rest) if _filter(x) else None # Binary ops compare_filtered("add --size 10x1024x128 --size 1x1024x128 --cpu") @@ -150,13 +164,15 @@ def predicate(x): "matmul --size 16x768x768 --size 16x768x768 --transpose= --transpose 0,2,1" ) compare_filtered( - "matmul --size 16x768x768 --size 16x768x768 --transpose= --transpose 0,2,1 --cpu" + "matmul --size 16x768x768 --size 16x768x768 --transpose= --transpose 0,2,1 " + "--cpu" ) compare_filtered( "matmul --size 16x768x128 --size 16x768x128 --transpose= --transpose 0,2,1" ) compare_filtered( - "matmul --size 16x768x128 --size 16x768x128 --transpose= --transpose 0,2,1 --cpu" + "matmul --size 16x768x128 --size 16x768x128 --transpose= --transpose 0,2,1 " + "--cpu" ) compare_filtered("matmul --size 512x8192 --size 8192x512") compare_filtered("matmul --size 512x8192 --size 8192x512 --cpu") @@ -185,7 +201,8 @@ def predicate(x): compare_filtered("matmul --size 32x1x1000 --size 32x1000x128 --cpu") compare_filtered("matmul --size 32x1x1000 --size 32x1000x128") compare_filtered( - "matmul --size 32x1x1000 --size 32x128x1000 --transpose= --transpose 0,2,1 --cpu" + "matmul --size 32x1x1000 --size 32x128x1000 --transpose= --transpose 0,2,1 " + "--cpu" ) compare_filtered( "matmul --size 32x1x1000 --size 32x128x1000 --transpose= --transpose 0,2,1" @@ -218,7 +235,8 @@ def predicate(x): compare_filtered("step --size 32x16x1024 --cpu") compare_filtered("selu --size 32x16x1024") compare_filtered("selu --size 32x16x1024 --cpu") - # compare_filtered("mish --size 32x16x1024") NOTE: Torch does not implement Mish in MPS atm + # NOTE: Torch does not implement Mish in MPS atm + # compare_filtered("mish --size 32x16x1024") compare_filtered("mish --size 32x16x1024 --cpu") compare_filtered("prelu --size 32x16x1024") compare_filtered("prelu --size 32x16x1024 --cpu") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv1d_bench.py b/Source/Cmlx/mlx/benchmarks/python/conv1d_bench.py index 0306a3e08..e1e15302f 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv1d_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv1d_bench.py @@ -1,6 +1,5 @@ import argparse import math -import os import subprocess import time @@ -53,10 +52,10 @@ def pt_conv_1D(a, b): return pt_conv_1D -def bench_shape(N, iH, C, wH, O, strides, padding, np_dtype, groups): +def bench_shape(N, iH, C, wH, OC, strides, padding, np_dtype, groups): scale = 1.0 / math.sqrt(wH * C) a_np = np.random.uniform(0, 0.5, (N, iH, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, wH, int(C / groups))).astype(np_dtype) + b_np = np.random.uniform(-scale, scale, (OC, wH, int(C / groups))).astype(np_dtype) a_mx = mx.array(a_np) b_mx = mx.array(b_np) @@ -83,7 +82,9 @@ def bench_shape(N, iH, C, wH, O, strides, padding, np_dtype, groups): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, iH, C)}, {(O, wH, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, iH, C)}, {(OC, wH, C)} [strides = {strides}, padding = " + f"{padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -108,15 +109,16 @@ def bench_shape(N, iH, C, wH, O, strides, padding, np_dtype, groups): for dtype in dtypes: print("(N, iH, C), (O, wH, C), dtype, stride, pads, groups, diff%") - for N, iH, C, wH, O, strides, padding, groups in shapes: + for N, iH, C, wH, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, iH, C, wH, O, strides, padding, np_dtype, groups + N, iH, C, wH, OC, strides, padding, np_dtype, groups ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {iH:3d}, {C:3d}), ({O:3d}, {wH:2d}, {C:3d}), {dtype}, {strides:5d}, {padding:4d}, {groups:6d}, {100. * diff:+5.2f}%" + f"({N}, {iH:3d}, {C:3d}), ({OC:3d}, {wH:2d}, {C:3d}), {dtype}, " + f"{strides:5d}, {padding:4d}, {groups:6d}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: diff --git a/Source/Cmlx/mlx/benchmarks/python/conv2d_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv2d_bench_cpu.py index d560ae1b3..46575f440 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv2d_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv2d_bench_cpu.py @@ -47,10 +47,10 @@ def pt_conv_2D(a, b): return pt_conv_2D -def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kH * kH * C) a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kH, kW, int(C / groups))).astype( np_dtype ) @@ -77,7 +77,9 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, H, W, C)}, {(OC, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -111,17 +113,19 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): for dtype in dtypes: print( - "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, " + "diff%" ) - for N, H, W, C, kH, kW, O, strides, padding, groups in shapes: + for N, H, W, C, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype + N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kH:2d}, {kW:2d}, {C:3d}), " + f"{dtype}, {strides}, {padding}, {groups:7d}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv2d_train_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv2d_train_bench_cpu.py index d85587909..c572226ad 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv2d_train_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv2d_train_bench_cpu.py @@ -64,7 +64,7 @@ def step(params, image): mx.eval(state) end_time = time.perf_counter() - print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") + print(f"{i:3d}, time={(end_time - start_time) * 1000:7.2f} ms") total_time += (end_time - start_time) * 1000 return total_time @@ -119,7 +119,7 @@ def loss_fn(pred_image, image): end_time = time.perf_counter() - print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") + print(f"{i:3d}, time={(end_time - start_time) * 1000:7.2f} ms") total_time += (end_time - start_time) * 1000 return total_time @@ -130,13 +130,13 @@ def main(): time_mlx = bench_mlx(steps) time_torch = bench_torch(steps) - print(f"average time of MLX: {time_mlx/steps:9.2f} ms") + print(f"average time of MLX: {time_mlx / steps:9.2f} ms") print(f"total time of MLX: {time_mlx:9.2f} ms") - print(f"average time of PyTorch: {time_torch/steps:9.2f} ms") + print(f"average time of PyTorch: {time_torch / steps:9.2f} ms") print(f"total time of PyTorch: {time_torch:9.2f} ms") diff = time_torch / time_mlx - 1.0 - print(f"torch/mlx diff: {100. * diff:+5.2f}%") + print(f"torch/mlx diff: {100.0 * diff:+5.2f}%") if __name__ == "__main__": diff --git a/Source/Cmlx/mlx/benchmarks/python/conv2d_transpose_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv2d_transpose_bench_cpu.py index 28f31f66a..fcade3f9a 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv2d_transpose_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv2d_transpose_bench_cpu.py @@ -50,10 +50,10 @@ def pt_conv_transpose_2D(a, b): return pt_conv_transpose_2D -def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kH * kH * C) a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (int(O / groups), kH, kW, C)).astype( + b_np = np.random.uniform(-scale, scale, (int(OC / groups), kH, kW, C)).astype( np_dtype ) @@ -82,7 +82,9 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, H, W, C)}, {(OC, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -113,17 +115,19 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): for dtype in dtypes: print( - "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, " + "diff%" ) - for N, H, W, C, kH, kW, O, strides, padding, groups in shapes: + for N, H, W, C, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype + N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kH:2d}, {kW:2d}, {C:3d}), " + f"{dtype}, {strides}, {padding}, {groups:7d}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv3d_bench.py b/Source/Cmlx/mlx/benchmarks/python/conv3d_bench.py index 91c62dfc3..5d44d0bd1 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv3d_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv3d_bench.py @@ -47,11 +47,11 @@ def pt_conv_3D(a, b, b_prime): return pt_conv_3D -def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kD * kH * kW * C) a_np = np.random.uniform(0, 0.5, (N, D, H, W, C)) - b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups))) - b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(O / groups))) + b_np = np.random.uniform(-scale, scale, (OC, kD, kH, kW, int(C / groups))) + b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(OC / groups))) a_np, b_np, b_prime_np = map(lambda x: x.astype(np_dtype), (a_np, b_np, b_prime_np)) a_mx, b_mx, b_prime_mx = map(lambda x: mx.array(x), (a_np, b_np, b_prime_np)) @@ -97,7 +97,7 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} " + f"Failed at {(N, D, H, W, C)}, {(OC, kD, kH, kW, C)} " f"[strides = {strides}, padding = {padding}, groups = {groups}] " f"with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" ) @@ -130,23 +130,24 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype ) for dtype in dtypes: - print(f"\n{'=' * 120}" f"\n dtype: {dtype}" f"\n{'=' * 120}") + print(f"\n{'=' * 120}\n dtype: {dtype}\n{'=' * 120}") print( f"{'(N, D, H, W, C)':<26s} {'( O, kD, kH, kW, C)':<24s} " f"{'stride':<12s} {'pads':<12s} {'groups':>6s} " f"{'diff%':>7s} " f"{'MLX peak':>9s} {'MLX act':>8s} {'PT cur':>8s} {'PT drv':>8s}" ) - for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes: + for N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch, mlx_peak, mlx_act, pt_cur, pt_drv = bench_shape( - N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype + N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), " + f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kD:2d}, {kH:2d}, " + f"{kW:2d}, {C:3d}), " f"{strides}, {padding}, {groups:6d}, " - f"{100. * diff:+6.1f}% " + f"{100.0 * diff:+6.1f}% " f"{mlx_peak:8.1f} {mlx_act:7.1f} {pt_cur:7.1f} {pt_drv:7.1f}" ) diff --git a/Source/Cmlx/mlx/benchmarks/python/conv3d_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv3d_bench_cpu.py index d3b263da8..39d004294 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv3d_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv3d_bench_cpu.py @@ -47,10 +47,10 @@ def pt_conv_3D(a, b): return pt_conv_3D -def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kD * kH * kW * C) a_np = np.random.uniform(0, 0.5, (N, D, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kD, kH, kW, int(C / groups))).astype( np_dtype ) @@ -77,7 +77,9 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, D, H, W, C)}, {(OC, kD, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -94,17 +96,20 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype for dtype in dtypes: print( - "(N, D, H, W, C), ( O, kD, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, D, H, W, C), ( O, kD, kH, kW, C), dtype, stride, " + "pads, groups, diff%" ) - for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes: + for N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype + N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kD:2d}, {kH:2d}, " + f"{kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, " + f"{100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv3d_train_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv3d_train_bench_cpu.py index fee99ee76..c25926746 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv3d_train_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv3d_train_bench_cpu.py @@ -64,7 +64,7 @@ def step(params, image): mx.eval(state) end_time = time.perf_counter() - print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") + print(f"{i:3d}, time={(end_time - start_time) * 1000:7.2f} ms") total_time += (end_time - start_time) * 1000 return total_time @@ -119,7 +119,7 @@ def loss_fn(pred_image, image): end_time = time.perf_counter() - print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") + print(f"{i:3d}, time={(end_time - start_time) * 1000:7.2f} ms") total_time += (end_time - start_time) * 1000 return total_time @@ -130,13 +130,13 @@ def main(): time_mlx = bench_mlx(steps) time_torch = bench_torch(steps) - print(f"average time of MLX: {time_mlx/steps:9.2f} ms") + print(f"average time of MLX: {time_mlx / steps:9.2f} ms") print(f"total time of MLX: {time_mlx:9.2f} ms") - print(f"average time of PyTorch: {time_torch/steps:9.2f} ms") + print(f"average time of PyTorch: {time_torch / steps:9.2f} ms") print(f"total time of PyTorch: {time_torch:9.2f} ms") diff = time_torch / time_mlx - 1.0 - print(f"torch/mlx diff: {100. * diff:+5.2f}%") + print(f"torch/mlx diff: {100.0 * diff:+5.2f}%") if __name__ == "__main__": diff --git a/Source/Cmlx/mlx/benchmarks/python/conv3d_transpose_bench_cpu.py b/Source/Cmlx/mlx/benchmarks/python/conv3d_transpose_bench_cpu.py index 742d9d46b..b23da4d32 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv3d_transpose_bench_cpu.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv3d_transpose_bench_cpu.py @@ -51,10 +51,10 @@ def pt_conv_3D(a, b): return pt_conv_3D -def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kD * kH * kW * C) a_np = np.random.uniform(0, 0.5, (N, D, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kD, kH, kW, int(C / groups))).astype( np_dtype ) @@ -83,7 +83,9 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, D, H, W, C)}, {(OC, kD, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -100,17 +102,20 @@ def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype for dtype in dtypes: print( - "(N, D, H, W, C), ( O, kD, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, D, H, W, C), ( O, kD, kH, kW, C), dtype, stride, " + "pads, groups, diff%" ) - for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes: + for N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype + N, D, H, W, C, kD, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kD:2d}, {kH:2d}, " + f"{kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, " + f"{100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv_bench.py b/Source/Cmlx/mlx/benchmarks/python/conv_bench.py index f421963ad..ca9426364 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv_bench.py @@ -1,6 +1,5 @@ import argparse import math -import os import subprocess import time @@ -53,10 +52,10 @@ def pt_conv_2D(a, b): return pt_conv_2D -def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kH * kH * C) a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kH, kW, int(C / groups))).astype( np_dtype ) @@ -85,7 +84,9 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, H, W, C)}, {(OC, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -119,17 +120,19 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): for dtype in dtypes: print( - "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, " + "diff%" ) - for N, H, W, C, kH, kW, O, strides, padding, groups in shapes: + for N, H, W, C, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype + N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kH:2d}, {kW:2d}, {C:3d}), " + f"{dtype}, {strides}, {padding}, {groups:7d}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv_transpose_bench.py b/Source/Cmlx/mlx/benchmarks/python/conv_transpose_bench.py index 22522d8ed..938dd81df 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv_transpose_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv_transpose_bench.py @@ -1,7 +1,5 @@ import argparse import math -import os -import subprocess import time import mlx.core as mx @@ -54,10 +52,10 @@ def pt_conv_transpose_2D(a, b): return pt_conv_transpose_2D -def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kH * kH * C) a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kH, kW, int(C / groups))).astype( np_dtype ) @@ -88,7 +86,9 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, H, W, C)}, {(OC, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -119,17 +119,19 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): for dtype in dtypes: print( - "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, diff%" + "(N, H, W, C), ( O, kH, kW, C), dtype, stride, pads, groups, " + "diff%" ) - for N, H, W, C, kH, kW, O, strides, padding, groups in shapes: + for N, H, W, C, kH, kW, OC, strides, padding, groups in shapes: np_dtype = getattr(np, dtype) time_mlx, time_torch = bench_shape( - N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype + N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%" + f"({N}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kH:2d}, {kW:2d}, {C:3d}), " + f"{dtype}, {strides}, {padding}, {groups:7d}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/conv_unaligned_bench.py b/Source/Cmlx/mlx/benchmarks/python/conv_unaligned_bench.py index 981d7b48b..f95e43da6 100644 --- a/Source/Cmlx/mlx/benchmarks/python/conv_unaligned_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/conv_unaligned_bench.py @@ -47,10 +47,10 @@ def pt_conv_2D(a, b): return pt_conv_2D -def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): +def bench_shape(N, H, W, C, kH, kW, OC, strides, padding, groups, np_dtype): scale = 1.0 / math.sqrt(kH * kH * C) a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) - b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( + b_np = np.random.uniform(-scale, scale, (OC, kH, kW, int(C / groups))).astype( np_dtype ) @@ -79,7 +79,9 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): if not np.allclose(out_pt, out_mx, atol=atol): print( - f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" + f"Failed at {(N, H, W, C)}, {(OC, kH, kW, C)} [strides = {strides}, " + f"padding = {padding}, groups = {groups}] with max(|a - b|) = " + f"{np.max(np.abs(out_pt - out_mx))}" ) return time_mlx, time_torch @@ -94,14 +96,15 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): (4, 32, 32, 370, 7, 7, 128), (2, 320, 640, 21, 7, 7, 21), ) - for N, H, W, C, kh, kw, O in shapes: + for N, H, W, C, kh, kw, OC in shapes: time_mlx, time_torch = bench_shape( - N, H, W, C, kh, kw, O, (1, 1), (0, 0), 1, dtype + N, H, W, C, kh, kw, OC, (1, 1), (0, 0), 1, dtype ) diff = time_torch / time_mlx - 1.0 print( - f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kh:2d}, {kw:2d}, {C:3d}), {dtype}, {100. * diff:+5.2f}%" + f"({N}, {H:3d}, {W:3d}, {C:3d}), ({OC:3d}, {kh:2d}, {kw:2d}, {C:3d}), " + f"{dtype}, {100.0 * diff:+5.2f}%" ) if time_mlx >= 2.0 * time_torch: print("ATTENTION ^^^^^^^") diff --git a/Source/Cmlx/mlx/benchmarks/python/distributed_bench.py b/Source/Cmlx/mlx/benchmarks/python/distributed_bench.py index b6fcef19e..0b5936041 100644 --- a/Source/Cmlx/mlx/benchmarks/python/distributed_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/distributed_bench.py @@ -26,7 +26,7 @@ def time_fn(fn, *args, **kwargs): num_iters = 100 tic = time.perf_counter() for _ in range(num_iters): - x = mx.eval(fn(*args, **kwargs)) + _x = mx.eval(fn(*args, **kwargs)) toc = time.perf_counter() msec = 1e3 * (toc - tic) / num_iters diff --git a/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py b/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py index 277897699..1e0dfcef9 100644 --- a/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py @@ -3,23 +3,16 @@ import itertools import os import time -from datetime import datetime -from typing import Optional, Tuple import mlx.core as mx -import numpy as np -RED_BOLD = "\033[1;31m" -GREEN = "\033[0;32m" RESET = "\033[0m" - N_warmup = 8 N_iter_bench = 80 N_iter_func = 5 -# similar to ./blas/bench_gemm.py def bench(f, *args): for _ in range(N_warmup): f(*args) @@ -30,7 +23,7 @@ def bench(f, *args): f(*args) mx.synchronize() e = time.perf_counter_ns() - return (e - s) * 1e-9 # total seconds for N_iter_bench * N_iter_func calls + return (e - s) * 1e-9 def do_kernel_bench(f, *args): @@ -43,7 +36,23 @@ def do_kernel_bench(f, *args): return ys -def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): +def do_grad_bench(f, *args): + ys = [] + for _ in range(N_iter_func): + ys.extend(f(*args)) + mx.eval(ys) + return ys + + +def make_grad_fn(): + def f(q, k, v, g, b, h0): + out, state = mx.fast.gated_delta_update(q, k, v, g, b, h0) + return out.sum() + state.sum() + + return mx.grad(f, argnums=(0, 1, 2, 3, 4, 5)) + + +def make_inputs(B, T, Hk, Hv, Dk, Dv): mx.random.seed(42) q = mx.random.normal(shape=(B, T, Hk, Dk)) k = mx.random.normal(shape=(B, T, Hk, Dk)) @@ -51,12 +60,18 @@ def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): v = mx.random.normal(shape=(B, T, Hv, Dv)) g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) + mx.eval(q, k, v, g, b, h0) + return q, k, v, g, b, h0 + + +def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): + q, k, v, g, b, h0 = make_inputs(B, T, Hk, Hv, Dk, Dv) shape_str = f"B={B} T={T} Hk={Hk} Hv={Hv} Dk={Dk} Dv={Dv}" denom = N_iter_bench * N_iter_func os.environ["GATED_DELTA_CHUNK"] = "0" - h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0)) ms_seq = ( bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0) @@ -68,7 +83,6 @@ def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): for C in (c for c in chunk_sizes if c != 0): try: os.environ["GATED_DELTA_CHUNK"] = str(C) - h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0)) ms_c = ( bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0) @@ -86,18 +100,11 @@ def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"): if run_full: Bs = [1, 4, 8, 16] - Ts = [8, 64, 256, 512, 1024, 2048, 4096] - Hks = [16] - Hvs = [32] - Dks = [128] - Dvs = [128] + Ts = [8, 64, 256, 512, 1024, 2048] else: Bs = [1, 8, 16] Ts = [8, 512, 1024, 2048] - Hks = [16] - Hvs = [32] - Dks = [128] - Dvs = [128] + Hks, Hvs, Dks, Dvs = [16], [32], [128], [128] chunk_sizes = [0, 8, 16] non_zero_Cs = [C for C in chunk_sizes if C != 0] @@ -106,13 +113,13 @@ def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"): f"C={C} (speedup)" for C in non_zero_Cs ] - col_widths = [6, 6, 6, 6, 6, 6, 15] + [25] * (len(non_zero_Cs)) + col_widths = [6, 6, 6, 6, 6, 6, 15] + [25] * len(non_zero_Cs) fmt = "".join(f"{{:<{w}}}" for w in col_widths) rows = [] print(fmt.format(*headers)) - print("-" * (sum(col_widths))) + print("-" * sum(col_widths)) for B, T, Hk, Hv, Dk, Dv in itertools.product(Bs, Ts, Hks, Hvs, Dks, Dvs): shapes_s, base_time_s, speedups, base_time = benchmark_shape( @@ -135,11 +142,96 @@ def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"): print(f"\nResults also written to {csv_path}") +def benchmark_variants_shape(B, T, Hk, Hv, Dk, Dv, do_backward, variants): + q, k, v, g, b, h0 = make_inputs(B, T, Hk, Hv, Dk, Dv) + denom = N_iter_bench * N_iter_func + + if do_backward: + fn = make_grad_fn() + runner = do_grad_bench + else: + fn = mx.fast.gated_delta_update + runner = do_kernel_bench + + def time_one(variant): + os.environ["GATED_DELTA_VJP_FALLBACK"] = "1" if variant == "fallback" else "0" + C = "16" if variant == "nax" else "0" + if do_backward: + os.environ["GATED_DELTA_CHUNK"] = "16" + os.environ["GATED_DELTA_CHUNK_VJP"] = C + else: + os.environ["GATED_DELTA_CHUNK"] = C + mx.eval(*fn(q, k, v, g, b, h0)) + return bench(runner, fn, q, k, v, g, b, h0) / denom * 1e3 + + times = [] + for variant in variants: + try: + times.append(time_one(variant)) + except Exception as ex: + print(f" {variant} failed: {ex}") + times.append(float("nan")) + mx.clear_cache() + return times + + +def run_variants_benchmark(run_full, do_backward=False, do_fallback=False): + if run_full: + Bs = [1, 4, 8, 16] + Ts = [8, 32, 64, 128, 256, 512, 1024, 2048, 4096] + else: + Bs = [1, 8] + Ts = [8, 512, 1024] + if do_fallback: + Ts = [8, 32, 64, 128, 256, 512] + Hks, Hvs, Dks, Dvs = [16], [32], [128], [128] + + variants = ["seq", "nax"] + if do_fallback: + variants = ["fallback"] + variants + + headers = ["B", "T", "Hk", "Hv", "Dk", "Dv"] + headers += [f"{v} (ms)" for v in variants] + headers += [f"{v} (speedup)" for v in variants[1:]] + + col_widths = [6, 6, 6, 6, 6, 6] + [16] * len(variants) + [16] * (len(variants) - 1) + fmt = "".join(f"{{:<{w}}}" for w in col_widths) + + mode = "BACKWARD" if do_backward else "FORWARD" + print(f"\n=== {mode}: {' vs '.join(variants)} ===") + print(fmt.format(*headers)) + print("-" * sum(col_widths)) + + for B, T, Hk, Hv, Dk, Dv in itertools.product(Bs, Ts, Hks, Hvs, Dks, Dvs): + try: + times = benchmark_variants_shape( + B, T, Hk, Hv, Dk, Dv, do_backward, variants + ) + except Exception as ex: + print(f" B={B} T={T} failed: {ex}") + mx.clear_cache() + continue + + base = times[0] + row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}"] + row += [f"{t:.3f}" for t in times] + row += [f"{base / t:.2f}x" if t > 0 else "nan" for t in times[1:]] + print(fmt.format(*row)) + print(RESET, end="") + + if __name__ == "__main__": parser = argparse.ArgumentParser(description="Gated delta benchmark") parser.add_argument("--full", "-f", action="store_true") parser.add_argument("--csv", "-c", action="store_true") parser.add_argument("--csv_out", "-co", default="benchmark_results.csv") + parser.add_argument("--fallback", "-fb", action="store_true") + parser.add_argument("--backward", "-bw", action="store_true") args = parser.parse_args() - run_benchmark(args.full, to_csv=args.csv, csv_path=args.csv_out) + if args.backward or args.fallback: + run_variants_benchmark( + args.full, do_backward=args.backward, do_fallback=args.fallback + ) + else: + run_benchmark(args.full, to_csv=args.csv, csv_path=args.csv_out) diff --git a/Source/Cmlx/mlx/benchmarks/python/gather_mm_bench.py b/Source/Cmlx/mlx/benchmarks/python/gather_mm_bench.py index ffeb73487..6783b841c 100644 --- a/Source/Cmlx/mlx/benchmarks/python/gather_mm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/gather_mm_bench.py @@ -7,7 +7,7 @@ D = 1024 M = 1024 E = 32 -I = 4 +num_indices = 4 def gather_sort(x, indices): @@ -38,8 +38,8 @@ def time_gather_mm(): 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 - indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32) - sorted_indices = mx.sort(indices.flatten()).reshape(N, I) + indices = (mx.random.uniform(shape=(N, num_indices)) * E).astype(mx.uint32) + sorted_indices = mx.sort(indices.flatten()).reshape(N, num_indices) mx.eval(x, w1, w2, indices, sorted_indices) def gather_mm(x, w1, w2, indices, sort): @@ -57,7 +57,7 @@ def gather_mm(x, w1, w2, indices, sort): time_fn(gather_mm, x, w1, w2, sorted_indices, False) time_fn(gather_mm, x, w1, w2, indices, True) - x = mx.random.normal((N * I, D)) / 1024**0.5 + x = mx.random.normal((N * num_indices, D)) / 1024**0.5 w1 = mx.random.normal((M, D)) / 1024**0.5 w2 = mx.random.normal((D, M)) / 1024**0.5 mx.eval(x, w1, w2) diff --git a/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py index c05db9199..713878a83 100644 --- a/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py @@ -7,7 +7,7 @@ D = 1024 M = 1024 E = 32 -I = 4 +num_indices = 4 def gather_sort(x, indices): @@ -46,8 +46,8 @@ def time_gather_qmm(): 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) - sorted_indices = mx.sort(indices.flatten()).reshape(N, I) + indices = (mx.random.uniform(shape=(N, num_indices)) * E).astype(mx.uint32) + sorted_indices = mx.sort(indices.flatten()).reshape(N, num_indices) mx.eval(x, w1, w2, indices, sorted_indices) def gather_mm(x, w1, w2, indices, sort): @@ -65,7 +65,7 @@ def gather_mm(x, w1, w2, indices, sort): time_fn(gather_mm, x, w1, w2, sorted_indices, False) time_fn(gather_mm, x, w1, w2, indices, True) - x = mx.random.normal((N * I, D)) / 1024**0.5 + x = mx.random.normal((N * num_indices, D)) / 1024**0.5 w1 = mx.random.normal((M, D)) / 1024**0.5 w2 = mx.random.normal((D, M)) / 1024**0.5 w1 = mx.quantize(w1) @@ -82,13 +82,13 @@ def 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 + N, E, num_indices = 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) + indices = (mx.random.uniform(shape=(N, num_indices)) * E).astype(mx.uint32) mx.eval(x, w1, w2, indices) def gather_mm(x, w1, w2, indices): diff --git a/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py index 83257a4cd..fddbf05bf 100644 --- a/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py @@ -91,9 +91,7 @@ def benchmark(shape, workload, tokens): global_scale_w, rhs, sorted_, - msg=( - f"{workload} routes={rhs.size} " f"N={shape['output']} K={shape['input']}" - ), + msg=(f"{workload} routes={rhs.size} N={shape['output']} K={shape['input']}"), ) diff --git a/Source/Cmlx/mlx/benchmarks/python/large_gemm_bench.py b/Source/Cmlx/mlx/benchmarks/python/large_gemm_bench.py index 6f3b24865..9f084aae2 100644 --- a/Source/Cmlx/mlx/benchmarks/python/large_gemm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/large_gemm_bench.py @@ -113,7 +113,7 @@ def bench_gemm(M, N, K, dtype, rtol): print( f"{M:5d} {N:5d} {K:6d} " - f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f} " - f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f} " + f"{mlx_mean * 1000:7.2f}±{mlx_std * 1000:5.2f} " + f"{torch_mean * 1000:7.2f}±{torch_std * 1000:5.2f} " f"{speedup:8.2f}x" ) diff --git a/Source/Cmlx/mlx/benchmarks/python/layer_norm_bench.py b/Source/Cmlx/mlx/benchmarks/python/layer_norm_bench.py index 29925de0b..61e9bb8c5 100644 --- a/Source/Cmlx/mlx/benchmarks/python/layer_norm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/layer_norm_bench.py @@ -3,12 +3,10 @@ from functools import partial import mlx.core as mx -import mlx.nn as nn from time_utils import time_fn def layer_norm(x, w, b, eps): - ot = x.dtype x = x.astype(mx.float32) mu = mx.mean(x, -1, keepdims=True) v = mx.var(x, -1, keepdims=True) @@ -22,8 +20,13 @@ def layer_norm(x, w, b, eps): def time_layer_norm(N, dt): L = 1024 - f1 = lambda x, w, b, y: (layer_norm(x, w, b, 1e-5) * y).sum() - f2 = lambda x, w, b, y: (mx.fast.layer_norm(x, w, b, 1e-5) * y).sum() + + def f1(x, w, b, y): + return (layer_norm(x, w, b, 1e-5) * y).sum() + + def f2(x, w, b, y): + return (mx.fast.layer_norm(x, w, b, 1e-5) * y).sum() + g1 = mx.grad(f1, argnums=(0, 1, 2)) g2 = mx.grad(f2, argnums=(0, 1, 2)) @@ -52,8 +55,12 @@ def layer_norm_grad_loop(g, x, w, b): time_fn(layer_norm_grad_loop, mx.compile(g1), x, w, b) time_fn(layer_norm_grad_loop, mx.compile(g2), x, w, b) - f1 = lambda x, y: (layer_norm(x, None, None, 1e-5) * y).sum() - f2 = lambda x, y: (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum() + def f1(x, y): + return (layer_norm(x, None, None, 1e-5) * y).sum() + + def f2(x, y): + return (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum() + g1 = mx.grad(f1, argnums=(0,)) g2 = mx.grad(f2, argnums=(0,)) diff --git a/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py index 3939d2a8a..a0fd75d94 100644 --- a/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py @@ -48,8 +48,7 @@ def benchmark(shape, tokens): global_scale_x, global_scale_w, msg=( - f"mode=nvfp4 dtype=bfloat16 tokens={tokens} " - f"N={output_dims} K={input_dims}" + f"mode=nvfp4 dtype=bfloat16 tokens={tokens} N={output_dims} K={input_dims}" ), ) diff --git a/Source/Cmlx/mlx/benchmarks/python/rms_norm_bench.py b/Source/Cmlx/mlx/benchmarks/python/rms_norm_bench.py index 50f3a40b0..85ab8fba9 100644 --- a/Source/Cmlx/mlx/benchmarks/python/rms_norm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/rms_norm_bench.py @@ -1,7 +1,6 @@ # Copyright © 2023-2024 Apple Inc. import mlx.core as mx -import mlx.nn as nn from time_utils import time_fn @@ -16,8 +15,12 @@ def rms_norm(x, w, eps): def time_rms_norm(): - f1 = lambda x, w, y: (rms_norm(x, w, 1e-5) * y).sum() - f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, 1e-5) * y).sum() + def f1(x, w, y): + return (rms_norm(x, w, 1e-5) * y).sum() + + def f2(x, w, y): + return (mx.fast.rms_norm(x, w, 1e-5) * y).sum() + g1 = mx.grad(f1, argnums=(0, 1)) g2 = mx.grad(f2, argnums=(0, 1)) @@ -37,8 +40,12 @@ def rms_norm_loop(g, x, w): time_fn(rms_norm_loop, mx.compile(g1), x, w) time_fn(rms_norm_loop, mx.compile(g2), x, w) - f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum() - f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum() + def f1(x, y): + return (rms_norm(x, None, 1e-5) * y).sum() + + def f2(x, y): + return (mx.fast.rms_norm(x, None, 1e-5) * y).sum() + g1 = mx.grad(f1, argnums=(0,)) g2 = mx.grad(f2, argnums=(0,)) diff --git a/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py b/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py index 7dfc7e0d1..1d1c7b979 100644 --- a/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py @@ -74,7 +74,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None): scores = q @ mx.swapaxes(k, -1, -2) if mask is not None: - if mask == "causal": q_offset = max(0, kL - L) q_indices = mx.arange(q_offset, q_offset + L) @@ -149,19 +148,110 @@ def bench_shape( if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol, rtol=atol): print( - f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}" + f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: " + f"{n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] " + f"with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}" ) return time_mlx_fused, time_mlx_unfused +def set_vjp_fallback(enabled): + os.environ["MLX_SDPA_VJP_FALLBACK"] = "1" if enabled else "0" + + +def mlx_fused_attn_grads(q, k, v, scale, cotan, mask=None, transpose=False): + def f(q_, k_, v_): + return do_attention(mlx_fused_attn, q_, k_, v_, scale, mask, transpose) + + _, grads = mx.vjp(f, [q, k, v], [cotan]) + return grads + + +def do_attention_vjp_bench(q, k, v, scale, cotan, mask=None, transpose=False): + def f(q_, k_, v_): + return do_attention(mlx_fused_attn, q_, k_, v_, scale, mask, transpose) + + dq = q + + for i in range(N_iter_func): + _, (dq, dk, dv) = mx.vjp(f, [dq, k, v], [cotan]) + + mx.eval([dq, dk, dv]) + return dq + + +def peak_mem_vjp(q, k, v, scale, cotan, mask=None, transpose=False): + mx.clear_cache() + mx.reset_peak_memory() + grads = mlx_fused_attn_grads(q, k, v, scale, cotan, mask, transpose) + mx.eval(grads) + peak = mx.get_peak_memory() / float(1024.0**3) + del grads + mx.clear_cache() + return peak + + +def max_rel_diff(a_list, b_list): + rel = 0.0 + for a, b in zip(a_list, b_list): + denom = mx.maximum(mx.max(mx.abs(a)), mx.array(1e-6, mx.float32)) + rel = max(rel, (mx.max(mx.abs(a - b)) / denom).item()) + return rel + + +def bench_shape_vjp( + B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None +): + q_mx, k_mx, v_mx, scale, mask = prepare_inputs( + B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, mask_in, transpose, dtype + ) + cotan = mx.array(np.random.normal(0.0, 1.0, q_mx.shape).astype(getattr(np, dtype))) + mx.eval(cotan) + + set_vjp_fallback(True) + time_fallback = bench( + do_attention_vjp_bench, q_mx, k_mx, v_mx, scale, cotan, mask, transpose + ) + mem_fallback = peak_mem_vjp(q_mx, k_mx, v_mx, scale, cotan, mask, transpose) + g_fallback = mlx_fused_attn_grads(q_mx, k_mx, v_mx, scale, cotan, mask, transpose) + mx.eval(g_fallback) + + set_vjp_fallback(False) + time_vjp = bench( + do_attention_vjp_bench, q_mx, k_mx, v_mx, scale, cotan, mask, transpose + ) + mem_vjp = peak_mem_vjp(q_mx, k_mx, v_mx, scale, cotan, mask, transpose) + g_vjp = mlx_fused_attn_grads(q_mx, k_mx, v_mx, scale, cotan, mask, transpose) + mx.eval(g_vjp) + + rel = max_rel_diff(g_fallback, g_vjp) + # nax truncates so the accuracy is not full fp32 + atol = 5e-3 + + if rel > atol: + print( + f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: " + f"{n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] " + f"with max rel = {rel:3.2e}" + ) + + return time_vjp, time_fallback, mem_vjp, mem_fallback, rel + + def get_gflop_count(B, M, N, K): return float(2.0 * N_iter_bench * N_iter_func * B * M * N * K) / float(1024.0**3) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Run gemm benchmarks") - + parser.add_argument( + "-bw", + "--backward", + action="store_true", + help="benchmark the vjp against the fallback", + ) + args = parser.parse_args() dtypes = ("float16", "float32")[:1] transposes = (False,) @@ -228,10 +318,46 @@ def get_gflop_count(B, M, N, K): shapes = shapes_64 + shapes_72 + shapes_80 + shapes_96 + shapes_128 + shapes_256 + if args.backward: + masks = [None, "causal"] + + print( + " B, qsl, ksl, hdim, n_qh, n_kvh, t, dtype, mask, t_fall, " + "t_vjp, speedup, m_fall, m_vjp, mem_x, rel" + ) + + for dtype in ["float32"]: + for transpose in transposes: + for B, qsl, ksl, head_dim, n_q_heads, n_kv_heads in shapes: + for mask_in in masks: + t_vjp, t_fall, m_vjp, m_fall, rel = bench_shape_vjp( + B, + qsl, + ksl, + head_dim, + n_q_heads, + n_kv_heads, + dtype, + transpose, + mask_in, + ) + speedup = t_fall / t_vjp + mem_x = m_fall / max(m_vjp, 1e-9) + t_str = 1 if transpose else 0 + print( + f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, " + f"{n_q_heads:4d}, {n_kv_heads:5d}, {t_str:1d}, {dtype}, " + f"{str(mask_in):>8}, {t_fall: 2.3f}, {t_vjp: 2.3f}, " + f"{speedup:6.2f}x, {m_fall:6.2f}, {m_vjp:6.2f}, " + f"{mem_x:5.2f}x, {rel:3.1e}" + ) + exit(0) + masks = [None, "bool", "causal"] print( - " B, qsl, ksl, hdim, n_qh, n_kvh, t, dtype, mask, t_unfs, t_fuse, diff%" + " B, qsl, ksl, hdim, n_qh, n_kvh, t, dtype, mask, t_unfs, t_fuse, " + "diff%" ) for dtype in dtypes: @@ -252,5 +378,8 @@ def get_gflop_count(B, M, N, K): diff = time_mlx_unfused / time_mlx_fused - 1.0 t_str = 1 if transpose else 0 print( - f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:1d}, {dtype}, {str(mask_in):>8}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%" + f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, " + f"{n_kv_heads:5d}, {t_str:1d}, {dtype}, {str(mask_in):>8}, " + f"{time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, " + f"{100.0 * diff:+5.2f}%" ) diff --git a/Source/Cmlx/mlx/benchmarks/python/sdpa_vector_bench.py b/Source/Cmlx/mlx/benchmarks/python/sdpa_vector_bench.py index 546bff84c..983093029 100644 --- a/Source/Cmlx/mlx/benchmarks/python/sdpa_vector_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/sdpa_vector_bench.py @@ -1,95 +1,209 @@ +"""Benchmark single-token decode through the Metal SDPA vector kernels. + +Run the same command after building the PR parent and head revisions: + + python benchmarks/python/sdpa_vector_bench.py + +The default shapes select the generic ``sdpa_vector_2pass_1`` kernel changed +by this PR. Use ``--shape QH,KVH,CTX,D`` to benchmark additional shapes. +""" + import argparse import math +import statistics +import time import mlx.core as mx -from time_utils import time_fn - -L = 16384 -H = 32 -H_k = H // 4 -D = 128 -V = 128 -dtype = mx.float16 -loops = 10 - - -def upproject(x, w): - if w is None: - return x - else: - return x @ w.T - - -def attention(q, k, v, mask=None, w=None): - def _sdpa(q, k, v): - B, Hq, L, D = q.shape - _, Hk, S, _ = k.shape - _, _, _, V = v.shape - q = q.reshape(B, Hk, Hq // Hk, L, D) - k = k[:, :, None, :, :] - v = v[:, :, None, :, :] - s = q @ k.transpose(0, 1, 2, 4, 3) - if mask is not None: - m = mx.broadcast_to(mask, (B, Hq, L, S)).reshape(B, Hk, Hq // Hk, L, S) - s = mx.where(m, s, mx.finfo(s.dtype).min) - p = mx.softmax(s.astype(mx.float32), axis=-1).astype(s.dtype) - o = p @ v - return o.reshape(B, Hq, L, V) - - for i in range(loops): - q = _sdpa(q, k, v) - q = upproject(q, w) - return q - - -def sdpa(q, k, v, mask=None, w=None): - for i in range(loops): - q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask) - q = upproject(q, w) - return q - - -def time_self_attention_primitives(): - mx.random.seed(3) - q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype) - k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype) - v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype) - w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None - mx.eval(q, k, v, w) - time_fn(attention, q, k, v, w=w) - - -def time_self_attention_sdpa(): - mx.random.seed(3) - q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype) - k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype) - v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype) - w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None - mx.eval(q, k, v, w) - time_fn(sdpa, q, k, v, w=w) - - -def time_self_attention_sdpa_with_mask(): - mx.random.seed(3) - q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype) - k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype) - v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype) - w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None - mask = mx.full((L,), True) - mask[L // 2 :] = False - mx.eval(q, k, v, mask, w) - - def sdpa_mask(*args): - return sdpa(*args, mask=mask, w=w) - - def attention_mask(*args): - return attention(*args, mask=mask, w=w) - - time_fn(attention_mask, q, k, v) - time_fn(sdpa_mask, q, k, v) + +DEFAULT_SHAPES = ( + # q_heads, kv_heads, context, head_dim + (32, 32, 4096, 128), + (32, 32, 32768, 128), + (4, 4, 65536, 128), + (8, 4, 4096, 128), + (8, 4, 32768, 128), + (16, 4, 4096, 128), + (16, 4, 32768, 128), + # GQA=8 selects the generic 2-pass kernel below context 8192. + (32, 4, 4096, 128), +) + + +def parse_shape(value): + try: + shape = tuple(int(x) for x in value.split(",")) + except ValueError as error: + raise argparse.ArgumentTypeError( + "shape must be q_heads,kv_heads,context,head_dim" + ) from error + if len(shape) != 4: + raise argparse.ArgumentTypeError( + "shape must be q_heads,kv_heads,context,head_dim" + ) + q_heads, kv_heads, context, head_dim = shape + if min(shape) <= 0 or q_heads % kv_heads != 0: + raise argparse.ArgumentTypeError( + "shape values must be positive and q_heads divisible by kv_heads" + ) + return shape + + +def kernel_path(q_heads, kv_heads, context, head_dim, architecture): + gqa = q_heads // kv_heads + device_class = architecture[-1] if architecture else "" + uses_two_pass = (device_class in ("d", "s") and context >= 1024) or ( + kv_heads < q_heads and context >= 4096 + ) + if not uses_two_pass: + return "vector-1pass" + + specialized_gqa = (gqa == 8 and head_dim in (64, 128)) or ( + gqa in (12, 16) and head_dim == 128 + ) + if specialized_gqa and context >= 8192: + return "gqa-specialized" + return "generic-2pass" + + +def reference_sdpa(q, k, v, scale): + n_repeats = q.shape[1] // k.shape[1] + if n_repeats > 1: + k = mx.repeat(k, n_repeats, axis=1) + v = mx.repeat(v, n_repeats, axis=1) + scores = (q @ mx.swapaxes(k, -1, -2)) * scale + probabilities = mx.softmax(scores.astype(mx.float32), axis=-1).astype(q.dtype) + return probabilities @ v + + +def run_chain(q, k, v, scale, chain): + out = q + for _ in range(chain): + out = mx.fast.scaled_dot_product_attention(out, k, v, scale=scale) + mx.eval(out) + + +def measure(q, k, v, scale, chain, warmup, iterations, rounds): + for _ in range(warmup): + run_chain(q, k, v, scale, chain) + + round_medians = [] + best = math.inf + for _ in range(rounds): + samples = [] + for _ in range(iterations): + start = time.perf_counter_ns() + run_chain(q, k, v, scale, chain) + elapsed = (time.perf_counter_ns() - start) * 1e-9 / chain + samples.append(elapsed) + best = min(best, min(samples)) + round_medians.append(statistics.median(samples)) + + return ( + best, + statistics.median(round_medians), + min(round_medians), + max(round_medians), + ) + + +def benchmark_shape(shape, args, architecture): + q_heads, kv_heads, context, head_dim = shape + scale = 1.0 / math.sqrt(head_dim) + + mx.random.seed(args.seed) + q = mx.random.normal((1, q_heads, 1, head_dim)).astype(mx.float16) + k = mx.random.normal((1, kv_heads, context, head_dim)).astype(mx.float16) + v = mx.random.normal((1, kv_heads, context, head_dim)).astype(mx.float16) + mx.eval(q, k, v) + + actual = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + expected = reference_sdpa(q, k, v, scale) + mx.eval(actual, expected) + error = mx.abs(actual.astype(mx.float32) - expected.astype(mx.float32)) + max_abs_error = error.max().item() + relative_linf_error = max_abs_error / max(mx.abs(expected).max().item(), 1e-9) + + best, median, median_low, median_high = measure( + q, + k, + v, + scale, + args.chain, + args.warmup, + args.iterations, + args.rounds, + ) + + # Logical traffic of the generic kernel: every query head streams one K + # and one V head. This is intentionally called logical bandwidth because + # shared KV heads can be served from cache on GQA shapes. + itemsize = 2 + logical_bytes = q_heads * context * head_dim * itemsize * 2 + logical_gbps = logical_bytes / median / 1e9 + + return { + "gqa": q_heads // kv_heads, + "path": kernel_path(*shape, architecture), + "best_us": best * 1e6, + "median_us": median * 1e6, + "median_low_us": median_low * 1e6, + "median_high_us": median_high * 1e6, + "logical_gbps": logical_gbps, + "max_abs_error": max_abs_error, + "relative_linf_error": relative_linf_error, + } + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Benchmark single-token decode through the Metal sdpa_vector_2pass kernels." + ) + ) + parser.add_argument( + "--shape", + action="append", + type=parse_shape, + metavar="QH,KVH,CTX,D", + help="shape to benchmark; may be repeated", + ) + parser.add_argument("--chain", type=int, default=8) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=30) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + shapes = args.shape or DEFAULT_SHAPES + device_info = mx.device_info(mx.gpu) + architecture = device_info.get("architecture", "") + print( + f"device={device_info.get('device_name', 'unknown')} " + f"architecture={architecture}" + ) + print( + f"chain={args.chain} warmup={args.warmup} " + f"iterations={args.iterations} rounds={args.rounds}" + ) + print( + f"{'QH':>3} {'KVH':>3} {'GQA':>3} {'CTX':>6} {'D':>3} " + f"{'path':>15} {'best us':>9} {'median us [range]':>27} " + f"{'logical GB/s':>12} {'rel L-inf':>9}" + ) + for shape in shapes: + result = benchmark_shape(shape, args, architecture) + q_heads, kv_heads, context, head_dim = shape + print( + f"{q_heads:>3} {kv_heads:>3} {result['gqa']:>3} " + f"{context:>6} {head_dim:>3} {result['path']:>15} " + f"{result['best_us']:>9.1f} " + f"{result['median_us']:>9.1f} " + f"[{result['median_low_us']:>6.1f}," + f"{result['median_high_us']:>6.1f}] " + f"{result['logical_gbps']:>12.1f} " + f"{result['relative_linf_error']:>9.2e}" + ) if __name__ == "__main__": - time_self_attention_sdpa() - time_self_attention_primitives() - time_self_attention_sdpa_with_mask() + main() diff --git a/Source/Cmlx/mlx/benchmarks/python/segmented_mm_bench.py b/Source/Cmlx/mlx/benchmarks/python/segmented_mm_bench.py index 0823c2e04..9977a0b39 100644 --- a/Source/Cmlx/mlx/benchmarks/python/segmented_mm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/segmented_mm_bench.py @@ -133,7 +133,8 @@ def main(): mlx_dtype = MLX_DTYPES[args.dtype] print( - f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}" + f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} " + f"segments={args.segments}" ) headers = [ diff --git a/Source/Cmlx/mlx/benchmarks/python/slice_update_bench.py b/Source/Cmlx/mlx/benchmarks/python/slice_update_bench.py index 60f12b981..a1a78ac29 100644 --- a/Source/Cmlx/mlx/benchmarks/python/slice_update_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/slice_update_bench.py @@ -105,5 +105,6 @@ def slice_update(dst, updates, slice_range): ) print( f"{dtype:<12} {str(dst_shape):<25} {str(update_shape):<20} " - f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} {torch_bw:<12.2f}" + f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} " + f"{torch_bw:<12.2f}" ) diff --git a/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py b/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py index 805bc7c1d..01904698b 100644 --- a/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py @@ -36,7 +36,7 @@ def scatter_unsort(x, inv_order, shape=None): return x -def time_gather_mm(name, D, M, E, I): +def time_gather_mm(name, D, M, E, num_indices): w1 = mx.random.normal((E, M, D), dtype=mx.bfloat16, scale=D**-0.5) w2 = mx.random.normal((E, M, D), dtype=mx.bfloat16, scale=D**-0.5) w3 = mx.random.normal((E, D, M), dtype=mx.bfloat16, scale=M**-0.5) @@ -64,8 +64,10 @@ def gather_mm(x, w1, w2, w3, indices, sort): for N in SEQ_LENS: x = mx.random.normal((N, 1, 1, D), dtype=mx.bfloat16) scores = mx.random.uniform(shape=(N, E)) - indices = mx.argpartition(scores, E - I, axis=-1)[:, -I:].astype(mx.uint32) - sorted_indices = mx.sort(indices.flatten()).reshape(N, I) + indices = mx.argpartition(scores, E - num_indices, axis=-1)[ + :, -num_indices: + ].astype(mx.uint32) + sorted_indices = mx.sort(indices.flatten()).reshape(N, num_indices) mx.eval(x, indices, sorted_indices) label = f"{name} N={N}" diff --git a/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py index 2e80ec2fa..7ff8ab215 100644 --- a/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py @@ -36,7 +36,7 @@ def scatter_unsort(x, inv_order, shape=None): return x -def time_gather_qmm(name, D, M, E, I, mode): +def time_gather_qmm(name, D, M, E, num_indices, mode): w1 = mx.random.normal((E, M, D), dtype=mx.bfloat16, scale=D**-0.5) w2 = mx.random.normal((E, M, D), dtype=mx.bfloat16, scale=D**-0.5) w3 = mx.random.normal((E, D, M), dtype=mx.bfloat16, scale=M**-0.5) @@ -59,7 +59,9 @@ def gather_qmm(x, w1, w2, w3, indices, sort): for N in SEQ_LENS: x = mx.random.normal((N, 1, 1, D), dtype=mx.bfloat16) scores = mx.random.uniform(shape=(N, E)) - indices = mx.argpartition(scores, E - I, axis=-1)[:, -I:].astype(mx.uint32) + indices = mx.argpartition(scores, E - num_indices, axis=-1)[ + :, -num_indices: + ].astype(mx.uint32) mx.eval(x, indices) label = f"{name} {mode} N={N}" diff --git a/Source/Cmlx/mlx/benchmarks/python/time_utils.py b/Source/Cmlx/mlx/benchmarks/python/time_utils.py index 181f7d3fe..e03b7fca5 100644 --- a/Source/Cmlx/mlx/benchmarks/python/time_utils.py +++ b/Source/Cmlx/mlx/benchmarks/python/time_utils.py @@ -19,7 +19,7 @@ def time_fn(fn, *args, **kwargs): num_iters = 100 tic = time.perf_counter() for _ in range(num_iters): - x = mx.eval(fn(*args, **kwargs)) + _x = mx.eval(fn(*args, **kwargs)) toc = time.perf_counter() msec = 1e3 * (toc - tic) / num_iters diff --git a/Source/Cmlx/mlx/cmake/extension.cmake b/Source/Cmlx/mlx/cmake/extension.cmake index 2f244a14f..1318eb656 100644 --- a/Source/Cmlx/mlx/cmake/extension.cmake +++ b/Source/Cmlx/mlx/cmake/extension.cmake @@ -1,47 +1,119 @@ include(CMakeParseArguments) +function(_mlx_metal_supports_logging OUTPUT_VARIABLE) + execute_process( + COMMAND xcrun -sdk macosx metal --help + RESULT_VARIABLE _MLX_METAL_HELP_RESULT + OUTPUT_VARIABLE _MLX_METAL_HELP_OUTPUT + ERROR_VARIABLE _MLX_METAL_HELP_ERROR + TIMEOUT 10) + set(_MLX_METAL_HELP "${_MLX_METAL_HELP_OUTPUT}\n${_MLX_METAL_HELP_ERROR}") + string(FIND "${_MLX_METAL_HELP}" "-fmetal-enable-logging" + _MLX_METAL_LOGGING_OPTION_INDEX) + if(_MLX_METAL_HELP_RESULT STREQUAL "0" + AND NOT _MLX_METAL_LOGGING_OPTION_INDEX EQUAL -1) + set(${OUTPUT_VARIABLE} + TRUE + PARENT_SCOPE) + else() + set(${OUTPUT_VARIABLE} + FALSE + PARENT_SCOPE) + endif() +endfunction() + # clang format off # # ############################################################################## # Build metal library # -# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib -# from list ${SOURCES}, including list ${INCLUDE_DIRS}, depends on list ${DEPS} +# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib. +# Each source in ${SOURCES} is compiled to an AIR file before the AIR files are +# linked into the metallib, including list ${INCLUDE_DIRS}, depends on list +# ${DEPS}, and passes ${COMPILE_OPTIONS} to the Metal compiler. # # Args: TARGET: Custom target to be added for the metal library TITLE: Name of # the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List # of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency -# files (like headers) DEBUG: Boolean, if true, enables debug compile options -# for this specific library. If not provided, uses global MLX_METAL_DEBUG. +# files (like headers) COMPILE_OPTIONS: Additional Metal compiler options +# DEPLOYMENT_TARGET: Optional macOS deployment target for this library DEBUG: +# Boolean, if true, enables debug compile options for this specific library. If +# not provided, uses global MLX_METAL_DEBUG. # # clang format on macro(mlx_build_metallib) # Parse args - set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG) - set(multiValueArgs SOURCES INCLUDE_DIRS DEPS) + set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG DEPLOYMENT_TARGET) + set(multiValueArgs SOURCES INCLUDE_DIRS DEPS COMPILE_OPTIONS) cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) # Set output 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 - -Wmetal-addr-spaces) - if(MLX_METAL_DEBUG OR MTLLIB_DEBUG) - set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only - -frecord-sources) + set(_MTLLIB_COMPILE_OPTIONS + -Wall -Wextra -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions + -Wmetal-addr-spaces) + if(MLX_METAL_DEBUG + OR MTLLIB_DEBUG + OR CMAKE_BUILD_TYPE STREQUAL "Debug") + set(_MTLLIB_COMPILE_OPTIONS ${_MTLLIB_COMPILE_OPTIONS} -gline-tables-only + -frecord-sources) + endif() + if(CMAKE_BUILD_TYPE STREQUAL "Debug") + _mlx_metal_supports_logging(_MTLLIB_SUPPORTS_LOGGING) + if(_MTLLIB_SUPPORTS_LOGGING) + list(APPEND _MTLLIB_COMPILE_OPTIONS -fmetal-enable-logging) + endif() endif() - # Prepare metallib build command + set(_MTLLIB_LINK_OPTIONS) + if(MTLLIB_DEPLOYMENT_TARGET) + set(_MTLLIB_DEPLOYMENT_OPTION + "-mmacosx-version-min=${MTLLIB_DEPLOYMENT_TARGET}") + list(APPEND _MTLLIB_COMPILE_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION}) + list(APPEND _MTLLIB_LINK_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION}) + elseif(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") + set(_MTLLIB_DEPLOYMENT_OPTION + "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") + list(APPEND _MTLLIB_COMPILE_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION}) + list(APPEND _MTLLIB_LINK_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION}) + endif() + + # Compile each Metal source separately for incremental builds + set(_MTLLIB_AIR_TARGETS) + set(_MTLLIB_INCLUDE_OPTIONS ${MTLLIB_INCLUDE_DIRS}) + list(TRANSFORM _MTLLIB_INCLUDE_OPTIONS PREPEND "-I") + set(_MTLLIB_SOURCE_INDEX 0) + foreach(_MTLLIB_SOURCE IN LISTS MTLLIB_SOURCES) + get_filename_component(_MTLLIB_SOURCE_STEM "${_MTLLIB_SOURCE}" NAME_WE) + set(_MTLLIB_AIR_NAME + "${MTLLIB_TARGET}_${_MTLLIB_SOURCE_INDEX}_${_MTLLIB_SOURCE_STEM}.air") + set(_MTLLIB_AIR_TARGET "${CMAKE_CURRENT_BINARY_DIR}/${_MTLLIB_AIR_NAME}") + set(_MTLLIB_DEPFILE "${_MTLLIB_AIR_TARGET}.d") + add_custom_command( + OUTPUT ${_MTLLIB_AIR_TARGET} + COMMAND + xcrun -sdk macosx metal ${_MTLLIB_INCLUDE_OPTIONS} + ${_MTLLIB_COMPILE_OPTIONS} ${MTLLIB_COMPILE_OPTIONS} -MMD -MF + ${_MTLLIB_DEPFILE} -MT ${_MTLLIB_AIR_TARGET} -c ${_MTLLIB_SOURCE} -o + ${_MTLLIB_AIR_TARGET} + DEPENDS ${MTLLIB_DEPS} ${_MTLLIB_SOURCE} + DEPFILE ${_MTLLIB_DEPFILE} + COMMAND_EXPAND_LISTS + COMMENT "Building ${_MTLLIB_AIR_NAME}" + VERBATIM) + list(APPEND _MTLLIB_AIR_TARGETS ${_MTLLIB_AIR_TARGET}) + math(EXPR _MTLLIB_SOURCE_INDEX "${_MTLLIB_SOURCE_INDEX} + 1") + endforeach() + + # Link the AIR files into a metallib add_custom_command( OUTPUT ${MTLLIB_BUILD_TARGET} - COMMAND - xcrun -sdk macosx metal - "$" - ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET} - DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES} - COMMAND_EXPAND_LISTS + COMMAND xcrun -sdk macosx metal ${_MTLLIB_LINK_OPTIONS} + ${_MTLLIB_AIR_TARGETS} -o ${MTLLIB_BUILD_TARGET} + DEPENDS ${_MTLLIB_AIR_TARGETS} COMMENT "Building ${MTLLIB_TITLE}.metallib" VERBATIM) diff --git a/Source/Cmlx/mlx/docs/src/conf.py b/Source/Cmlx/mlx/docs/src/conf.py index b95a3fe74..62c7d7e0b 100644 --- a/Source/Cmlx/mlx/docs/src/conf.py +++ b/Source/Cmlx/mlx/docs/src/conf.py @@ -2,8 +2,6 @@ # -*- coding: utf-8 -*- -import os -import subprocess import mlx.core as mx diff --git a/Source/Cmlx/mlx/docs/src/dev/extensions.rst b/Source/Cmlx/mlx/docs/src/dev/extensions.rst index 5f5fab642..50df005f1 100644 --- a/Source/Cmlx/mlx/docs/src/dev/extensions.rst +++ b/Source/Cmlx/mlx/docs/src/dev/extensions.rst @@ -549,24 +549,20 @@ Building and Binding Let's look at the overall directory structure first. -| extensions +| cmake_extension | ├── axpby | │ ├── axpby.cpp | │ ├── axpby.h | │ └── axpby.metal -| ├── mlx_sample_extensions -| │ └── __init__.py | ├── bindings.cpp | ├── CMakeLists.txt | └── setup.py -* ``extensions/axpby/`` defines the C++ extension library -* ``extensions/mlx_sample_extensions`` sets out the structure for the - associated Python package -* ``extensions/bindings.cpp`` provides Python bindings for our operation -* ``extensions/CMakeLists.txt`` holds CMake rules to build the library and +* ``cmake_extension/axpby/`` defines the C++ extension library +* ``cmake_extension/bindings.cpp`` provides Python bindings for our operation +* ``cmake_extension/CMakeLists.txt`` holds CMake rules to build the library and Python bindings -* ``extensions/setup.py`` holds the ``setuptools`` rules to build and install +* ``cmake_extension/setup.py`` holds the ``setuptools`` rules to build and install the Python package Binding to Python @@ -578,7 +574,7 @@ already provided, adding our :meth:`axpby` is simple. .. code-block:: C++ - NB_MODULE(_ext, m) { + NB_MODULE(mlx_sample_extensions, m) { m.doc() = "Sample extension for MLX"; m.def( @@ -611,13 +607,6 @@ already provided, adding our :meth:`axpby` is simple. Most of the complexity in the above example comes from additional bells and whistles such as the literal names and doc-strings. -.. warning:: - - :mod:`mlx.core` must be imported before importing - :mod:`mlx_sample_extensions` as defined by the nanobind module above to - ensure that the casters for :mod:`mlx.core` components like - :class:`mlx.core.array` are available. - .. _Building with CMake: Building with CMake @@ -678,20 +667,19 @@ Finally, we build the nanobind_ bindings .. code-block:: cmake nanobind_add_module( - _ext - NB_STATIC STABLE_ABI LTO NOMINSIZE + mlx_sample_extensions + NB_STATIC FREE_THREADED LTO NOMINSIZE NB_DOMAIN mlx ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp ) - target_link_libraries(_ext PRIVATE mlx_ext) - - if(BUILD_SHARED_LIBS) - target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path) - endif() + target_link_libraries(mlx_sample_extensions PRIVATE mlx_ext) Building with ``setuptools`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Building with a ``CMakeLists.txt`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + Once we have set out the CMake build rules as described above, we can use the build utilities defined in :mod:`mlx.extension`: @@ -705,42 +693,130 @@ build utilities defined in :mod:`mlx.extension`: name="mlx_sample_extensions", version="0.0.0", description="Sample C++ and Metal extensions for MLX primitives.", - ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")], + ext_modules=[extension.CMakeExtension("mlx_sample_extensions")], cmdclass={"build_ext": extension.CMakeBuild}, - packages=["mlx_sample_extensions"], - package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, + packages=[], extras_require={"dev":[]}, zip_safe=False, python_requires=">=3.10", ) -.. note:: - We treat ``extensions/mlx_sample_extensions`` as the package directory - even though it only contains a ``__init__.py`` to ensure the following: +On macOS and Linux, ``CMakeBuild`` configures a relocatable runtime search path +for all targets. It locates colocated libraries through ``@loader_path`` on +macOS or ``$ORIGIN`` on Linux, and locates the MLX library under ``mlx/lib`` in +the same ``site-packages`` directory. The extension module can therefore load +without importing ``mlx.core`` first. Values provided through ``CMAKE_ARGS`` +can override this default. - * :mod:`mlx.core` must be imported before importing :mod:`_ext` - * The C++ extension library and the metal library are co-located with the python - bindings and copied together if the package is installed +The extension name is the Python import name. This example builds the top-level +module ``mlx_sample_extensions`` without a Python wrapper package. A suffix such +as ``._ext`` is optional: use a dotted name such as ``my_package._ext`` only when +placing the extension inside a Python package. For ``CMakeExtension``, the +nanobind module name must match the final component of the extension name. To build the package, first install the build dependencies with ``pip install -r requirements.txt``. You can then build inplace for development using -``python setup.py build_ext -j8 --inplace`` (in ``extensions/``) +``python setup.py build_ext -j8 --inplace`` (in ``cmake_extension/``) This results in the directory structure: -| extensions -| ├── mlx_sample_extensions -| │ ├── __init__.py -| │ ├── libmlx_ext.dylib # C++ extension library -| │ ├── mlx_ext.metallib # Metal library -| │ └── _ext.cpython-3x-darwin.so # Python Binding +| cmake_extension +| ├── libmlx_ext.dylib # C++ extension library +| ├── mlx_ext.metallib # Metal library +| ├── mlx_sample_extensions.cpython-3x-darwin.so # Python binding | ... When you try to install using the command ``python -m pip install .`` (in -``extensions/``), the package will be installed with the same structure as -``extensions/mlx_sample_extensions`` and the C++ and Metal library will be -copied along with the Python binding since they are specified as -``package_data``. +``cmake_extension/``), the Python module and its C++ and Metal libraries are +installed together directly under ``site-packages``. + +Building without a ``CMakeLists.txt`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For a source-list based build, use ``mlx.extension.MetalExtension`` with +``mlx.extension.BuildExtension``. The helper builds the C++/nanobind module and +a colocated Metal library without requiring a ``CMakeLists.txt`` in the +extension project. + +.. code-block:: python + + from importlib.metadata import version as package_version + + from mlx import extension + from setuptools import setup + + setup( + name="mlx_sample_metal_extension", + ext_modules=[ + extension.MetalExtension( + "mlx_sample_metal_extension", + sources=[ + "activation.cpp", + "silu_and_mul.metal", + ], + extra_compile_args={ + "cxx": ["-O3"], + "metal": ["-O3"], + }, + ) + ], + cmdclass={"build_ext": extension.BuildExtension}, + install_requires=[f"mlx=={package_version('mlx')}"], + packages=[], + ) + +As with ``CMakeExtension``, a ``._ext`` suffix is not required. The final component +of the extension name is also the default Metal library name. This example +builds the top-level module ``mlx_sample_metal_extension``, with +``mlx_sample_metal_extension.metallib`` and +``mlx_sample_metal_extension.pyi`` beside it. These files are included in the +wheel and installed together under ``site-packages``. + +A list passed as ``extra_compile_args`` applies to C++; a dictionary can provide +separate ``cxx`` and ``metal`` options. To disable nanobind stub generation, +configure the command class with +``BuildExtension.with_options(generate_stubs=False)``. + +Pin the ``mlx`` install requirement to the version used for the build because +MLX does not provide a stable C++ ABI across releases. Reading the version from +the installed package metadata avoids importing ``mlx.core`` during setup. + +The generated extension uses a relocatable runtime search path to locate +``libmlx.dylib`` under ``mlx/lib`` in the same ``site-packages`` directory. This +avoids embedding the build machine's absolute MLX library path and lets the +extension module load without importing ``mlx.core`` first. + +The build defines ``MLX_EXTENSION_NAME`` as the final component of the +extension name and ``MLX_METAL_LIBRARY_NAME`` as its string form. Use them so +the binding does not hard-code either the module or Metal library name: + +.. code-block:: C++ + + auto library = + device.get_library(MLX_METAL_LIBRARY_NAME, current_binary_dir()); + + NB_MODULE(MLX_EXTENSION_NAME, m) { + // bindings + } + +Debug builds include Metal line tables and recorded sources. Metal logging is +enabled when the current ``xcrun metal`` compiler advertises support. The +effective Python +``MACOSX_DEPLOYMENT_TARGET`` is forwarded to CMake and applied to both AIR +compilation and metallib linking. By default, the target recorded in the +installed ``libmlx.dylib`` is used. Set ``MACOSX_DEPLOYMENT_TARGET`` or pass +``CMAKE_OSX_DEPLOYMENT_TARGET`` through ``CMAKE_ARGS`` to override it. This can +select a lower target for cross-building, provided the extension is used with +an MLX build for that target at runtime. Explicit targets below macOS 14 are +rejected. + +``BuildExtension`` uses CMake's platform default generator. To use Ninja, +include ``ninja`` in the build requirements and configure the command class +with ``BuildExtension.with_options(use_ninja=True)``. It falls back to the platform +default generator if Ninja is not available. + +The complete standalone example is available in +``examples/metal_extension/``. Usage ----- @@ -827,7 +903,7 @@ Scripts .. admonition:: Download the code - The full example code is available in `mlx `_. + The full example code is available in `mlx `_. .. _Accelerate: https://developer.apple.com/documentation/accelerate/blas?language=objc .. _Metal: https://developer.apple.com/documentation/metal?language=objc diff --git a/Source/Cmlx/mlx/docs/src/python/fft.rst b/Source/Cmlx/mlx/docs/src/python/fft.rst index 78bfe7f8f..0df0b665d 100644 --- a/Source/Cmlx/mlx/docs/src/python/fft.rst +++ b/Source/Cmlx/mlx/docs/src/python/fft.rst @@ -20,6 +20,8 @@ FFT irfft2 rfftn irfftn + stft + istft fftfreq rfftfreq fftshift diff --git a/Source/Cmlx/mlx/docs/src/python/ops.rst b/Source/Cmlx/mlx/docs/src/python/ops.rst index f55ba1b4d..52790680f 100644 --- a/Source/Cmlx/mlx/docs/src/python/ops.rst +++ b/Source/Cmlx/mlx/docs/src/python/ops.rst @@ -134,6 +134,7 @@ Operations meshgrid min minimum + minmax moveaxis multiply nan_to_num @@ -203,6 +204,7 @@ Operations triu trunc unflatten + unique unstack vecdot var diff --git a/Source/Cmlx/mlx/docs/src/usage/distributed.rst b/Source/Cmlx/mlx/docs/src/usage/distributed.rst index 431aef710..861525dcb 100644 --- a/Source/Cmlx/mlx/docs/src/usage/distributed.rst +++ b/Source/Cmlx/mlx/docs/src/usage/distributed.rst @@ -292,6 +292,29 @@ The following JSON defines the valid 4-node mesh from the image above. } ] +Two Macs can also be connected with more than one thunderbolt cable. In that +case, write a list of rdma devices for the peer instead of one device. The +device at position ``k`` in the list must be connected to the device at +position ``k`` in the list of the peer. JACCL sends large messages over all +the cables at the same time. If some pairs have fewer cables than others, JACCL +uses the smallest number of cables for all pairs. The following JSON defines +two Macs connected with two cables. + +.. code-block:: json + + [ + { + "ssh": "m3-ultra-1", + "ips": ["123.123.123.1"], + "rdma": [null, ["rdma_en2", "rdma_en3"]] + }, + { + "ssh": "m3-ultra-2", + "ips": [], + "rdma": [["rdma_en2", "rdma_en3"], null] + } + ] + Even though TCP/IP is not used when communicating with Thunderbolt RDMA, disabling the thunderbolt bridge is still required as well as setting up isolated local networks for each thunderbolt connection. diff --git a/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst b/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst index 3584f4c22..95ee3bdaa 100644 --- a/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst +++ b/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst @@ -87,11 +87,27 @@ handled by the CUDA runtime. Metal ----- +.. envvar:: MLX_METAL_RESIDENCY_REFRESH_INTERVAL_MS + + Re-request residency for existing wired allocations at this interval in + milliseconds. This works around macOS releasing residency while the GPU is + idle. The default is ``0`` (disabled); nonpositive values disable renewal. + For example, ``1000`` renews once per second. This does not change the wired + limit set by :func:`mlx.core.set_wired_limit` or submit GPU work. Renewal + continues while the Metal device exists and can keep memory wired during + long idle periods. Requires a Metal 3 GPU on macOS 15 or iOS 18 or later. + .. 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). +.. envvar:: MLX_CONV_WINOGRAD + + Allow the Winograd algorithm for 3x3 ``conv2d``. The default is ``1``. Set + it to ``0`` to use the GEMM convolution for all inputs. See + :doc:`precision`. + Advanced tuning --------------- @@ -177,8 +193,9 @@ The MLX-prefixed variables in this section are advanced CUDA backend controls. .. envvar:: MLX_ENABLE_CACHE_THRASHING_CHECK - Detect repeated CUDA cache misses and raise an error suggesting a larger - cache capacity. The default is ``1``. + Detect consecutive CUDA cache misses (more than twice the capacity with no + hit in between) 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/Source/Cmlx/mlx/docs/src/usage/export.rst b/Source/Cmlx/mlx/docs/src/usage/export.rst index 9900ed6d2..b06ca9672 100644 --- a/Source/Cmlx/mlx/docs/src/usage/export.rst +++ b/Source/Cmlx/mlx/docs/src/usage/export.rst @@ -207,7 +207,7 @@ The argument to the callback (``args``) is a dictionary which includes a * ``"keyword_inputs"``: The keyword specified inputs to the exported function * ``"outputs"``: The ordered outputs of the exported function * ``"constants"``: Any graph constants -* ``"primitives"``: Inner graph nodes representating the operations +* ``"primitive"``: An inner graph node representing an operation Each type has additional fields in the ``args`` dictionary. diff --git a/Source/Cmlx/mlx/docs/src/usage/precision.rst b/Source/Cmlx/mlx/docs/src/usage/precision.rst index 440a1f705..31c3f2efe 100644 --- a/Source/Cmlx/mlx/docs/src/usage/precision.rst +++ b/Source/Cmlx/mlx/docs/src/usage/precision.rst @@ -19,3 +19,15 @@ To keep these operations in full ``float32``, set Which operations take the reduced-precision path, and how large the difference is, depends on the backend and the hardware. + +On Metal, ``conv2d`` can use the Winograd algorithm for 3x3 convolutions +with stride 1. Winograd is faster for large inputs, but its rounding error +in ``float32`` can be approximately 10 times larger. MLX selects Winograd +from the total input size, which includes the batch size. Thus, the same +image can give different results in a batch of 3 and in a batch of 1. + +To disable Winograd, set :envvar:`MLX_CONV_WINOGRAD` to ``0``: + +.. code-block:: shell + + MLX_CONV_WINOGRAD=0 python my_script.py diff --git a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt b/Source/Cmlx/mlx/examples/cmake_extension/CMakeLists.txt similarity index 90% rename from Source/Cmlx/mlx/examples/extensions/CMakeLists.txt rename to Source/Cmlx/mlx/examples/cmake_extension/CMakeLists.txt index 014f71240..22ce3c300 100644 --- a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt +++ b/Source/Cmlx/mlx/examples/cmake_extension/CMakeLists.txt @@ -1,6 +1,6 @@ cmake_minimum_required(VERSION 3.27) -project(_ext LANGUAGES CXX) +project(mlx_sample_extensions LANGUAGES CXX) # ----------------------------- Setup ----------------------------- set(CMAKE_CXX_STANDARD 20) @@ -23,7 +23,7 @@ find_package( FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v3.0.1 + GIT_TAG v3.1.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) @@ -71,16 +71,12 @@ endif() # ----------------------------- Python Bindings ----------------------------- nanobind_add_module( - _ext + mlx_sample_extensions NB_STATIC - STABLE_ABI + FREE_THREADED LTO NOMINSIZE NB_DOMAIN mlx ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp) -target_link_libraries(_ext PRIVATE mlx_ext) - -if(BUILD_SHARED_LIBS) - target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path) -endif() +target_link_libraries(mlx_sample_extensions PRIVATE mlx_ext) diff --git a/Source/Cmlx/mlx/examples/extensions/README.md b/Source/Cmlx/mlx/examples/cmake_extension/README.md similarity index 64% rename from Source/Cmlx/mlx/examples/extensions/README.md rename to Source/Cmlx/mlx/examples/cmake_extension/README.md index 610de99f4..0599fc139 100644 --- a/Source/Cmlx/mlx/examples/extensions/README.md +++ b/Source/Cmlx/mlx/examples/cmake_extension/README.md @@ -1,4 +1,7 @@ +The example builds `mlx_sample_extensions` as a top-level Python extension +module. No `._ext` suffix or Python wrapper package is required. + ## Build ``` diff --git a/Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp b/Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.cpp similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp rename to Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.cpp diff --git a/Source/Cmlx/mlx/examples/extensions/axpby/axpby.h b/Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.h similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/axpby/axpby.h rename to Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.h diff --git a/Source/Cmlx/mlx/examples/extensions/axpby/axpby.metal b/Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.metal similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/axpby/axpby.metal rename to Source/Cmlx/mlx/examples/cmake_extension/axpby/axpby.metal diff --git a/Source/Cmlx/mlx/examples/extensions/bindings.cpp b/Source/Cmlx/mlx/examples/cmake_extension/bindings.cpp similarity index 95% rename from Source/Cmlx/mlx/examples/extensions/bindings.cpp rename to Source/Cmlx/mlx/examples/cmake_extension/bindings.cpp index 91892fa90..a094d1797 100644 --- a/Source/Cmlx/mlx/examples/extensions/bindings.cpp +++ b/Source/Cmlx/mlx/examples/cmake_extension/bindings.cpp @@ -8,7 +8,7 @@ namespace nb = nanobind; using namespace nb::literals; -NB_MODULE(_ext, m) { +NB_MODULE(mlx_sample_extensions, m) { m.doc() = "Sample extension for MLX"; m.def( diff --git a/Source/Cmlx/mlx/examples/extensions/pyproject.toml b/Source/Cmlx/mlx/examples/cmake_extension/pyproject.toml similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/pyproject.toml rename to Source/Cmlx/mlx/examples/cmake_extension/pyproject.toml diff --git a/Source/Cmlx/mlx/examples/extensions/requirements.txt b/Source/Cmlx/mlx/examples/cmake_extension/requirements.txt similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/requirements.txt rename to Source/Cmlx/mlx/examples/cmake_extension/requirements.txt diff --git a/Source/Cmlx/mlx/examples/extensions/setup.py b/Source/Cmlx/mlx/examples/cmake_extension/setup.py similarity index 76% rename from Source/Cmlx/mlx/examples/extensions/setup.py rename to Source/Cmlx/mlx/examples/cmake_extension/setup.py index be990a02b..f8fe3ada7 100644 --- a/Source/Cmlx/mlx/examples/extensions/setup.py +++ b/Source/Cmlx/mlx/examples/cmake_extension/setup.py @@ -9,10 +9,9 @@ name="mlx_sample_extensions", version="0.0.0", description="Sample C++ and Metal extensions for MLX primitives.", - ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")], + ext_modules=[extension.CMakeExtension("mlx_sample_extensions")], cmdclass={"build_ext": extension.CMakeBuild}, - packages=["mlx_sample_extensions"], - package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, + packages=[], zip_safe=False, python_requires=">=3.10", ) diff --git a/Source/Cmlx/mlx/examples/extensions/test.py b/Source/Cmlx/mlx/examples/cmake_extension/test.py similarity index 100% rename from Source/Cmlx/mlx/examples/extensions/test.py rename to Source/Cmlx/mlx/examples/cmake_extension/test.py diff --git a/Source/Cmlx/mlx/examples/export/eval_mlp.py b/Source/Cmlx/mlx/examples/export/eval_mlp.py index 133ab6917..07981523b 100644 --- a/Source/Cmlx/mlx/examples/export/eval_mlp.py +++ b/Source/Cmlx/mlx/examples/export/eval_mlp.py @@ -2,7 +2,6 @@ import mlx.core as mx import mlx.nn as nn -import mlx.utils class MLP(nn.Module): @@ -19,13 +18,12 @@ def __init__( ] def __call__(self, x): - for l in self.layers[:-1]: - x = nn.relu(l(x)) + for layer in self.layers[:-1]: + x = nn.relu(layer(x)) return self.layers[-1](x) if __name__ == "__main__": - batch_size = 8 input_dim = 32 output_dim = 10 diff --git a/Source/Cmlx/mlx/examples/export/train_mlp.py b/Source/Cmlx/mlx/examples/export/train_mlp.py index 0e19650a1..c188a94b4 100644 --- a/Source/Cmlx/mlx/examples/export/train_mlp.py +++ b/Source/Cmlx/mlx/examples/export/train_mlp.py @@ -20,13 +20,12 @@ def __init__( ] def __call__(self, x): - for l in self.layers[:-1]: - x = nn.relu(l(x)) + for layer in self.layers[:-1]: + x = nn.relu(layer(x)) return self.layers[-1](x) if __name__ == "__main__": - batch_size = 8 input_dim = 32 output_dim = 10 diff --git a/Source/Cmlx/mlx/examples/extensions/mlx_sample_extensions/__init__.py b/Source/Cmlx/mlx/examples/extensions/mlx_sample_extensions/__init__.py deleted file mode 100644 index b56a54fd2..000000000 --- a/Source/Cmlx/mlx/examples/extensions/mlx_sample_extensions/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright © 2023 Apple Inc. - -import mlx.core as mx - -from ._ext import axpby diff --git a/Source/Cmlx/mlx/examples/metal_extension/README.md b/Source/Cmlx/mlx/examples/metal_extension/README.md new file mode 100644 index 000000000..755cf4c3e --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/README.md @@ -0,0 +1,28 @@ +# Activation Metal extension + +A minimal MLX Metal extension exposing `silu_and_mul`. It uses +`mlx.extension.MetalExtension` to build the top-level module +`mlx_sample_metal_extension`, without a `._ext` suffix or Python wrapper package. + +The kernel is adapted from Hugging Face +[`activation_metal`](https://github.com/huggingface/kernels-community/tree/47a3168d0808921eef2f7daca794a4fccae13078/activation/activation_metal) +under Apache-2.0. + +## Build + +```bash +pip install -r requirements.txt +python setup.py build_ext -j8 --inplace +``` + +The build uses CMake's default generator and generates +`mlx_sample_metal_extension.pyi` by default. +To enable Ninja, install `ninja` and use +`BuildExtension.with_options(use_ninja=True)`. To disable stub generation, use +`BuildExtension.with_options(generate_stubs=False)`. + +## Test + +```bash +python test.py +``` diff --git a/Source/Cmlx/mlx/examples/metal_extension/activation.cpp b/Source/Cmlx/mlx/examples/metal_extension/activation.cpp new file mode 100644 index 000000000..7f99b9df5 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/activation.cpp @@ -0,0 +1,139 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "mlx/backend/metal/device.h" +#include "mlx/ops.h" +#include "mlx/primitives.h" + +namespace mx = mlx::core; +namespace nb = nanobind; +using namespace nb::literals; + +namespace { + +std::string current_binary_dir() { + static std::string binary_dir = []() { + Dl_info info; + if (!dladdr(reinterpret_cast(¤t_binary_dir), &info)) { + throw std::runtime_error("Unable to get current binary directory."); + } + return std::filesystem::path(info.dli_fname).parent_path().string(); + }(); + return binary_dir; +} + +class SiluAndMul : public mx::UnaryPrimitive { + public: + explicit SiluAndMul(mx::Stream stream) : mx::UnaryPrimitive(stream) {} + + void eval_cpu(const std::vector&, mx::array&) override { + throw std::runtime_error("SiluAndMul is only implemented for the GPU."); + } + + void eval_gpu(const std::vector& inputs, mx::array& out) override { + const auto& input = inputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + const uint32_t last_dim = input.shape(-1); + const uint32_t d = last_dim / 2; + const uint32_t num_tokens = input.size() / last_dim; + const uint32_t num_chunks = (d + 7) / 8; + + auto& device = mx::metal::device(stream().device); + auto library = + device.get_library(MLX_METAL_LIBRARY_NAME, current_binary_dir()); + auto kernel = device.get_kernel("silu_and_mul_f16", library); + + auto& encoder = mx::metal::get_command_encoder(stream()); + encoder.set_compute_pipeline_state(kernel); + encoder.set_output_array(out, 0); + encoder.set_input_array(input, 1); + encoder.set_bytes(d, 2); + + const size_t group_width = std::min( + static_cast(num_chunks), + static_cast(kernel->maxTotalThreadsPerThreadgroup())); + encoder.dispatch_threads( + MTL::Size(num_chunks, num_tokens, 1), MTL::Size(group_width, 1, 1)); + } + + const char* name() const override { + return "SiluAndMul"; + } + + bool is_equivalent(const mx::Primitive&) const override { + return true; + } +}; + +mx::array silu_and_mul(const mx::array& input, mx::StreamOrDevice s = {}) { + if (input.ndim() == 0) { + throw std::invalid_argument( + "silu_and_mul input must have at least one dimension."); + } + if (input.shape(-1) % 2 != 0) { + throw std::invalid_argument( + "silu_and_mul requires an even-sized last dimension."); + } + if (input.dtype() != mx::float16) { + throw std::invalid_argument("silu_and_mul only supports float16."); + } + + auto stream = mx::to_stream(s); + auto contiguous_input = mx::contiguous(input, false, stream); + auto output_shape = input.shape(); + output_shape.back() /= 2; + return mx::array( + output_shape, + input.dtype(), + std::make_shared(stream), + {contiguous_input}); +} + +} // namespace + +NB_MODULE(MLX_EXTENSION_NAME, m) { + m.doc() = + "Fused float16 SiLU and multiplication implemented with MLX and Metal."; + m.def( + "silu_and_mul", + &silu_and_mul, + "input"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def silu_and_mul(input: mlx.core.array, *, stream: " + "mlx.core.Stream | mlx.core.ThreadLocalStream | mlx.core.Device | " + "None = None) -> mlx.core.array"), + R"pbdoc( + Apply a fused SiLU activation and multiplication. + + Splits the last dimension into equal halves ``x`` and ``y``, then + computes ``silu(x) * y``. The input must be a float16 array with an + even-sized last dimension. + + Args: + input (mlx.core.array): Input array. + stream (Stream or Device, optional): Stream or device on which to + schedule the operation. Defaults to ``None``. + + Returns: + mlx.core.array: An array whose last dimension is half that of the + input. + )pbdoc"); +} diff --git a/Source/Cmlx/mlx/examples/metal_extension/pyproject.toml b/Source/Cmlx/mlx/examples/metal_extension/pyproject.toml new file mode 100644 index 000000000..4a2250f87 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/pyproject.toml @@ -0,0 +1,19 @@ +[build-system] +requires = [ + "setuptools>=80", + "cmake>=3.25", + "mlx>=0.32.0", + "nanobind==3.1.0", +] +build-backend = "setuptools.build_meta" + +[project] +name = "mlx_sample_metal_extension" +version = "0.0.0" +description = "Activation MetalExtension example for MLX." +requires-python = ">=3.10" +dynamic = ["dependencies"] + +[tool.setuptools] +packages = [] +zip-safe = false diff --git a/Source/Cmlx/mlx/examples/metal_extension/requirements.txt b/Source/Cmlx/mlx/examples/metal_extension/requirements.txt new file mode 100644 index 000000000..9b1529014 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/requirements.txt @@ -0,0 +1,4 @@ +setuptools>=42 +cmake>=3.25 +mlx>=0.32.0 +nanobind==3.1.0 diff --git a/Source/Cmlx/mlx/examples/metal_extension/setup.py b/Source/Cmlx/mlx/examples/metal_extension/setup.py new file mode 100644 index 000000000..a87b30256 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/setup.py @@ -0,0 +1,23 @@ +# Copyright © 2026 Apple Inc. + +# from importlib.metadata import version as package_version + +from setuptools import setup + +from mlx import extension + +if __name__ == "__main__": + setup( + ext_modules=[ + extension.MetalExtension( + "mlx_sample_metal_extension", + sources=[ + "activation.cpp", + "silu_and_mul.metal", + ], + extra_compile_args={"cxx": ["-O3"], "metal": ["-O3"]}, + ) + ], + cmdclass={"build_ext": extension.BuildExtension}, + # install_requires=[f"mlx=={package_version('mlx')}"], + ) diff --git a/Source/Cmlx/mlx/examples/metal_extension/silu_and_mul.metal b/Source/Cmlx/mlx/examples/metal_extension/silu_and_mul.metal new file mode 100644 index 000000000..5ee755345 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/silu_and_mul.metal @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// Adapted from huggingface/kernels-community activation_metal at commit +// 47a3168d0808921eef2f7daca794a4fccae13078. +#include +using namespace metal; + +inline float silu(float x) { + return x / (1.0f + metal::exp(-x)); +} + +inline float4 silu(float4 x) { + return x / (1.0f + metal::exp(-x)); +} + +kernel void silu_and_mul_f16( + device half* out [[buffer(0)]], + device const half* input [[buffer(1)]], + constant uint& d [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) { + uint base_idx = gid.x * 8; + if (base_idx >= d) { + return; + } + + uint in_base = gid.y * 2u * d; + uint out_base = gid.y * d; + + if (base_idx + 8 <= d) { + float4 x0 = float4( + *reinterpret_cast(&input[in_base + base_idx])); + float4 x1 = float4( + *reinterpret_cast(&input[in_base + base_idx + 4])); + float4 y0 = float4( + *reinterpret_cast(&input[in_base + d + base_idx])); + float4 y1 = float4(*reinterpret_cast( + &input[in_base + d + base_idx + 4])); + + *reinterpret_cast(&out[out_base + base_idx]) = + half4(silu(x0) * y0); + *reinterpret_cast(&out[out_base + base_idx + 4]) = + half4(silu(x1) * y1); + } else if (base_idx + 4 <= d) { + float4 x = float4( + *reinterpret_cast(&input[in_base + base_idx])); + float4 y = float4( + *reinterpret_cast(&input[in_base + d + base_idx])); + *reinterpret_cast(&out[out_base + base_idx]) = + half4(silu(x) * y); + + for (uint i = base_idx + 4; i < d; i++) { + float x = float(input[in_base + i]); + float y = float(input[in_base + d + i]); + out[out_base + i] = half(silu(x) * y); + } + } else { + for (uint i = base_idx; i < d; i++) { + float x = float(input[in_base + i]); + float y = float(input[in_base + d + i]); + out[out_base + i] = half(silu(x) * y); + } + } +} diff --git a/Source/Cmlx/mlx/examples/metal_extension/test.py b/Source/Cmlx/mlx/examples/metal_extension/test.py new file mode 100644 index 000000000..21630cb16 --- /dev/null +++ b/Source/Cmlx/mlx/examples/metal_extension/test.py @@ -0,0 +1,17 @@ +import mlx.core as mx +from mlx_sample_metal_extension import silu_and_mul + +base = mx.arange(2 * 74 * 3, dtype=mx.float32).reshape(2, 74, 3) +x = mx.sin(base * 0.17).transpose(0, 2, 1).astype(mx.float16) + +d = x.shape[-1] // 2 +left = x[..., :d] +right = x[..., d:] + +output = silu_and_mul(x, stream=mx.gpu) +expected = (left * mx.sigmoid(left)) * right +mx.eval(output, expected) + +correct = mx.allclose(output, expected, rtol=2e-3, atol=2e-3).item() +print(f"silu_and_mul: shape={output.shape}, dtype={output.dtype}, correct={correct}") +assert correct diff --git a/Source/Cmlx/mlx/mlx/allocator.h b/Source/Cmlx/mlx/mlx/allocator.h index c831b7247..6ef859945 100644 --- a/Source/Cmlx/mlx/mlx/allocator.h +++ b/Source/Cmlx/mlx/mlx/allocator.h @@ -3,9 +3,20 @@ #pragma once #include +#include +#include #include "mlx/api.h" +namespace mlx::core::allocator { +class Data; +} // namespace mlx::core::allocator + +namespace mlx::core::cu { +class CommandEncoder; +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); +} // namespace mlx::core::cu + namespace mlx::core::allocator { // Simple wrapper around buffer pointers @@ -51,12 +62,47 @@ class MLX_API Allocator { MLX_API Allocator& allocator(); -inline Buffer malloc(size_t size) { - return allocator().malloc(size); -} +using Deleter = std::function; + +// The allocator owns data from malloc +class Data { + public: + Data(Buffer buffer, Deleter d) : buffer_(buffer), deleter_(std::move(d)) {} + Data(const Data& other) = delete; + Data& operator=(const Data& other) = delete; + Data(Data&& other) noexcept + : buffer_(std::exchange(other.buffer_, Buffer(nullptr))), + deleter_(std::exchange(other.deleter_, nullptr)), + owned_(std::exchange(other.owned_, false)) {} + ~Data() { + if (owned_) { + allocator().free(buffer_); + } else if (deleter_) { + deleter_(buffer_); + } + } + + Buffer buffer() const& { + return buffer_; + } + Buffer buffer() const&& = delete; + + bool is_owned() const { + return owned_; + } + + private: + Buffer buffer_; + Deleter deleter_; + bool owned_{false}; + + explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {} + friend Data malloc(size_t size); + friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder); +}; -inline void free(Buffer buffer) { - allocator().free(buffer); +inline Data malloc(size_t size) { + return Data(allocator().malloc(size)); } // Make a Buffer from a raw pointer of the given size without a copy. If a diff --git a/Source/Cmlx/mlx/mlx/array.cpp b/Source/Cmlx/mlx/mlx/array.cpp index e6dbbe785..cf49a0fe5 100644 --- a/Source/Cmlx/mlx/mlx/array.cpp +++ b/Source/Cmlx/mlx/mlx/array.cpp @@ -60,13 +60,11 @@ std::vector array::make_arrays( array array::unsafe_weak_copy(const array& other) { auto cpy = array(other.shape(), other.dtype(), nullptr, {}); cpy.set_data( - other.buffer(), + Data(other.buffer(), nullptr), other.data_size(), other.strides(), other.flags(), - 0, - [](auto) {}); - cpy.array_desc_->offset = other.array_desc_->offset; + other.offset()); return cpy; } @@ -104,14 +102,13 @@ array::array( allocator::release(buffer); return deleter(ptr); }; - set_data(buffer, std::move(wrapped_deleter)); + set_data(Data(buffer, std::move(wrapped_deleter))); } } -/* Build an array from a shared buffer */ -array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter) +array::array(allocator::Data data, Shape shape, Dtype dtype) : array_desc_(std::make_shared(std::move(shape), dtype)) { - set_data(data, deleter); + set_data(std::move(data)); } void array::detach() { @@ -166,8 +163,8 @@ bool array::is_tracer() const { detail::retain_graph(); } -void array::set_data(allocator::Buffer buffer, Deleter d) { - array_desc_->data = std::make_shared(buffer, d); +void array::set_data(allocator::Data data) { + array_desc_->data = std::make_shared(std::move(data)); array_desc_->offset = 0; array_desc_->data_size = size(); array_desc_->flags.contiguous = true; @@ -177,13 +174,12 @@ void array::set_data(allocator::Buffer buffer, Deleter d) { } void array::set_data( - allocator::Buffer buffer, + allocator::Data data, size_t data_size, Strides strides, Flags flags, - int64_t offset, - Deleter d) { - array_desc_->data = std::make_shared(buffer, d); + int64_t offset) { + array_desc_->data = std::make_shared(std::move(data)); array_desc_->offset = offset; array_desc_->data_size = data_size; array_desc_->strides = std::move(strides); diff --git a/Source/Cmlx/mlx/mlx/array.h b/Source/Cmlx/mlx/mlx/array.h index cc29290d2..fd71104bc 100644 --- a/Source/Cmlx/mlx/mlx/array.h +++ b/Source/Cmlx/mlx/mlx/array.h @@ -19,7 +19,6 @@ namespace mlx::core { // Forward declaration class Primitive; -using Deleter = std::function; using ShapeElem = int32_t; using Shape = SmallVector; using Strides = SmallVector; @@ -69,12 +68,8 @@ class MLX_API array { Dtype dtype, const std::function& deleter); - /* Build an array from a buffer */ - explicit array( - allocator::Buffer data, - Shape shape, - Dtype dtype, - Deleter deleter = allocator::free); + /* Build an array from data */ + explicit array(allocator::Data data, Shape shape, Dtype dtype); /** Assignment to rvalue does not compile. */ array& operator=(const array& other) && = delete; @@ -238,22 +233,7 @@ class MLX_API array { return reinterpret_cast(array_desc_->primitive.get()); } - struct Data { - allocator::Buffer buffer; - Deleter d; - Data(allocator::Buffer buffer, Deleter d = allocator::free) - : buffer(buffer), d(d) {} - // Not copyable - Data(const Data& d) = delete; - Data& operator=(const Data& d) = delete; - Data(Data&& o) : buffer(o.buffer), d(o.d) { - o.buffer = allocator::Buffer(nullptr); - o.d = [](allocator::Buffer) {}; - } - ~Data() { - d(buffer); - } - }; + using Data = allocator::Data; struct Flags { // True iff there are no gaps in the underlying data. Each item @@ -302,7 +282,8 @@ class MLX_API array { /** True indicates the arrays buffer is safe to reuse */ bool is_donatable() const { - return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1); + return array_desc_.use_count() == 1 && + (array_desc_->data.use_count() == 1) && array_desc_->data->is_owned(); } /** The array's siblings. */ @@ -359,11 +340,8 @@ class MLX_API array { return array_desc_->data_size; } - allocator::Buffer& buffer() { - return array_desc_->data->buffer; - } - const allocator::Buffer& buffer() const { - return array_desc_->data->buffer; + allocator::Buffer buffer() const { + return array_desc_->data->buffer(); } size_t buffer_size() const { @@ -447,24 +425,14 @@ class MLX_API array { // Check if the array is a tracer array bool is_tracer() const; - void set_data(allocator::Buffer buffer, Deleter d = allocator::free); + void set_data(allocator::Data data); void set_data( - allocator::Buffer buffer, + allocator::Data data, size_t data_size, Strides strides, Flags flags, - Deleter d) { - set_data(buffer, data_size, std::move(strides), flags, 0, std::move(d)); - } - - void set_data( - allocator::Buffer buffer, - size_t data_size, - Strides strides, - Flags flags, - int64_t offset = 0, - Deleter d = allocator::free); + int64_t offset = 0); void copy_shared_buffer( const array& other, diff --git a/Source/Cmlx/mlx/mlx/backend/common/binary.h b/Source/Cmlx/mlx/mlx/backend/common/binary.h index 78607ef07..a05bc434c 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/binary.h +++ b/Source/Cmlx/mlx/mlx/backend/common/binary.h @@ -39,7 +39,7 @@ inline void set_binary_op_output_data( const array& b, array& out, BinaryOpType bopt, - std::function mallocfn = allocator::malloc) { + std::function mallocfn = allocator::malloc) { bool b_donatable = is_donatable(b, out); bool a_donatable = is_donatable(a, out); switch (bopt) { diff --git a/Source/Cmlx/mlx/mlx/backend/common/common.cpp b/Source/Cmlx/mlx/mlx/backend/common/common.cpp index 5426301a3..dbea6d16d 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/common.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/common.cpp @@ -2,6 +2,7 @@ #include #include "mlx/backend/common/broadcasting.h" +#include "mlx/backend/common/slicing.h" #include "mlx/backend/common/utils.h" #include "mlx/dtype_utils.h" #include "mlx/primitives.h" @@ -20,25 +21,33 @@ void AsStrided::eval(const std::vector& inputs, array& out) { "AsStrided must be used with row contiguous arrays only."); } - auto [no_bsx_size, row_contiguous, col_contiguous] = - check_contiguity(shape_, strides_); + auto [_, row_contiguous, col_contiguous] = check_contiguity(shape_, strides_); int64_t l = 0, h = 0; - bool has_negative_stride = false; + std::vector> dims; for (int i = 0; i < strides_.size(); i++) { auto delta = strides_[i] * (shape_[i] - 1); if (strides_[i] >= 0) { h += delta; } else { l += delta; - has_negative_stride |= shape_[i] > 1; + } + if (shape_[i] > 1 && strides_[i] != 0) { + dims.emplace_back(strides_[i], shape_[i]); } } size_t data_size = out.size() == 0 ? 0 : (h - l) + 1; + std::sort(dims.begin(), dims.end()); + int64_t expected_stride = 1; + bool dense = true; + for (auto [stride, size] : dims) { + dense &= stride == expected_stride; + expected_stride *= size; + } + auto flags = in.flags(); - flags.contiguous = - out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size); + flags.contiguous = out.size() == 0 || dense; flags.row_contiguous = row_contiguous; flags.col_contiguous = col_contiguous; @@ -155,6 +164,12 @@ void shared_buffer_reshape( // becomes col contiguous again. auto max_dim = std::max_element(out.shape().begin(), out.shape().end()); flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim; + } else { + // New shape can change contiguity + auto [_, row_contiguous, col_contiguous] = + check_contiguity(out.shape(), out_strides); + flags.row_contiguous = row_contiguous; + flags.col_contiguous = col_contiguous; } out.copy_shared_buffer(in, out_strides, flags, in.data_size()); } @@ -166,48 +181,29 @@ void Split::eval( auto& in = inputs[0]; - auto compute_new_flags = [](const auto& shape, - const auto& strides, - size_t in_data_size, - auto flags) { - size_t data_size = 1; - size_t f_stride = 1; - size_t b_stride = 1; - flags.row_contiguous = true; - flags.col_contiguous = true; - for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) { - flags.col_contiguous &= strides[i] == f_stride || shape[i] == 1; - flags.row_contiguous &= strides[ri] == b_stride || shape[ri] == 1; - f_stride *= shape[i]; - b_stride *= shape[ri]; - if (strides[i] > 0) { - data_size *= shape[i]; - } + for (int i = 0; i < outputs.size(); i++) { + if (outputs[i].size() == 0) { + outputs[i].set_data(allocator::malloc(0)); + continue; } - if (data_size == 1) { - // Broadcasted scalar array is contiguous. - flags.contiguous = true; - } else if (data_size == in_data_size) { - // Means we sliced a broadcasted dimension so leave the "no holes" flag - // alone. - } else { - // We sliced something. So either we are row or col contiguous or we - // punched a hole. - flags.contiguous &= flags.row_contiguous || flags.col_contiguous; + int64_t start = i == 0 ? 0 : indices_[i - 1]; + int64_t offset = start * in.strides()[axis_]; + + // compute the span + int64_t low_idx = 0; + int64_t high_idx = 0; + for (int j = 0; j < in.ndim(); j++) { + auto delta = in.strides()[j] * (outputs[i].shape()[j] - 1); + if (in.strides()[j] > 0) { + high_idx += delta; + } else { + low_idx += delta; + } } - return std::pair{flags, data_size}; - }; - - std::vector indices(1, 0); - indices.insert(indices.end(), indices_.begin(), indices_.end()); - for (int i = 0; i < indices.size(); i++) { - size_t offset = indices[i] * in.strides()[axis_]; - auto [new_flags, data_size] = compute_new_flags( - outputs[i].shape(), in.strides(), in.data_size(), in.flags()); - outputs[i].copy_shared_buffer( - in, in.strides(), new_flags, data_size, offset); + shared_buffer_slice( + in, in.strides(), offset, (high_idx - low_idx) + 1, outputs[i]); } } diff --git a/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp index 9dc5a4c54..0bcd66fe9 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp @@ -110,7 +110,7 @@ void compiled_allocate_outputs( std::vector& outputs, const std::function& is_constant, bool contiguous, - const std::function& + const std::function& mallocfn /* = allocator::malloc */) { if (contiguous) { int o = 0; diff --git a/Source/Cmlx/mlx/mlx/backend/common/compiled.h b/Source/Cmlx/mlx/mlx/backend/common/compiled.h index 94ba38b8f..c5f91d578 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/compiled.h +++ b/Source/Cmlx/mlx/mlx/backend/common/compiled.h @@ -78,8 +78,7 @@ void compiled_allocate_outputs( std::vector& outputs, const std::function& is_constant, bool contiguous, - const std::function& mallocfn = - allocator::malloc); + const std::function& mallocfn = allocator::malloc); // Collapse contiguous dims ignoring scalars and constants. std::tuple> diff --git a/Source/Cmlx/mlx/mlx/backend/common/copy.h b/Source/Cmlx/mlx/mlx/backend/common/copy.h index 859ce0410..943061b5c 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/copy.h +++ b/Source/Cmlx/mlx/mlx/backend/common/copy.h @@ -26,7 +26,7 @@ inline bool set_copy_output_data( const array& in, array& out, CopyType ctype, - std::function mallocfn = allocator::malloc) { + std::function mallocfn = allocator::malloc) { if (ctype == CopyType::Vector) { // If the input is donateable, we are doing a vector copy and the types // have the same size, then the input buffer can hold the output. diff --git a/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp index 7a2eef51f..d38a43395 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp @@ -203,7 +203,7 @@ std::string make_template_hash(const std::string& template_def) { template_hash.reserve(template_def.size()); for (size_t i = 0; i < template_def.size(); ++i) { auto c = template_def[i]; - if (c == '<' || c == '>') { + if (c == '<' || c == '>' || c == '-') { template_hash += '_'; } else if ( c == ',' && i + 1 < template_def.size() && template_def[i + 1] == ' ') { diff --git a/Source/Cmlx/mlx/mlx/backend/common/slicing.cpp b/Source/Cmlx/mlx/mlx/backend/common/slicing.cpp index cad2fd619..e2723fb8a 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/slicing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/slicing.cpp @@ -30,7 +30,8 @@ void shared_buffer_slice( auto flags = in.flags(); flags.row_contiguous = is_row_contiguous; flags.col_contiguous = is_col_contiguous; - flags.contiguous = (no_bsx_size == data_size); + flags.contiguous = is_row_contiguous || is_col_contiguous || + (flags.contiguous && no_bsx_size == data_size); out.copy_shared_buffer(in, out_strides, flags, data_size, data_offset); } diff --git a/Source/Cmlx/mlx/mlx/backend/common/slicing.h b/Source/Cmlx/mlx/mlx/backend/common/slicing.h index b667d2619..436a855c2 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/slicing.h +++ b/Source/Cmlx/mlx/mlx/backend/common/slicing.h @@ -11,6 +11,13 @@ std::tuple prepare_slice( const Shape& start_indices, const Shape& strides); +void shared_buffer_slice( + const array& in, + const Strides& out_strides, + int64_t data_offset, + size_t data_size, + array& out); + void slice( const array& in, array& out, diff --git a/Source/Cmlx/mlx/mlx/backend/common/ternary.h b/Source/Cmlx/mlx/mlx/backend/common/ternary.h index c63a57261..bd8da2580 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/ternary.h +++ b/Source/Cmlx/mlx/mlx/backend/common/ternary.h @@ -47,7 +47,7 @@ inline void set_ternary_op_output_data( const array& c, array& out, TernaryOpType topt, - std::function mallocfn = allocator::malloc) { + std::function mallocfn = allocator::malloc) { auto maybe_donate = [&out](const array& x) { if (is_donatable(x, out)) { out.copy_shared_buffer(x); diff --git a/Source/Cmlx/mlx/mlx/backend/common/unary.h b/Source/Cmlx/mlx/mlx/backend/common/unary.h index b19fc98ed..b37bb90ab 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/unary.h +++ b/Source/Cmlx/mlx/mlx/backend/common/unary.h @@ -10,7 +10,7 @@ namespace mlx::core { inline void set_unary_output_data( const array& in, array& out, - std::function mallocfn = allocator::malloc) { + std::function mallocfn = allocator::malloc) { if (in.flags().contiguous) { if (is_donatable(in, out)) { out.copy_shared_buffer(in); diff --git a/Source/Cmlx/mlx/mlx/backend/common/utils.cpp b/Source/Cmlx/mlx/mlx/backend/common/utils.cpp index 66764173e..a13c02207 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/utils.cpp @@ -89,13 +89,16 @@ std::pair collapse_contiguous_dims( Strides collapsed_strides; if (shape.size() > 0) { - collapsed_shape.push_back(shape[0]); - collapsed_strides.push_back(strides[0]); - for (int i = 1; i < shape.size(); i++) { + for (int i = 0; i < shape.size(); i++) { if (shape[i] == 1) { continue; - } else if ( - strides[i] * shape[i] != collapsed_strides.back() || + } + if (collapsed_shape.empty()) { + collapsed_shape.push_back(shape[i]); + collapsed_strides.push_back(strides[i]); + continue; + } + if (strides[i] * shape[i] != collapsed_strides.back() || collapsed_shape.back() * static_cast(shape[i]) > size_cap) { collapsed_shape.push_back(shape[i]); collapsed_strides.push_back(strides[i]); @@ -104,6 +107,10 @@ std::pair collapse_contiguous_dims( collapsed_strides.back() = strides[i]; } } + if (collapsed_shape.empty()) { + collapsed_shape.push_back(1); + collapsed_strides.push_back(1); + } } return std::make_pair(collapsed_shape, collapsed_strides); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp index 90b0378f6..150cf5833 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp @@ -119,6 +119,13 @@ void Equal::eval_cpu(const std::vector& inputs, array& out) { } } +void FloorDivide::eval_cpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + auto& a = inputs[0]; + auto& b = inputs[1]; + binary_op_cpu(a, b, out, detail::FloorDivide(), stream()); +} + void Greater::eval_cpu(const std::vector& inputs, array& out) { assert(inputs.size() == 2); comparison_op_cpu(inputs[0], inputs[1], out, detail::Greater(), stream()); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h index fb7db76f9..2eb4afc95 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h @@ -112,7 +112,7 @@ void binary_op_dispatch_dims( const T* b, U* out, int dim, - int size, + size_t size, const Shape& shape, const Strides& a_strides, const Strides& b_strides, @@ -135,7 +135,7 @@ void binary_op_dispatch_dims( ContiguousIterator a_it(shape, a_strides, dim - 3); ContiguousIterator b_it(shape, b_strides, dim - 3); auto stride = out_strides[dim - 4]; - for (int64_t elem = 0; elem < size; elem += stride) { + for (size_t elem = 0; elem < size; elem += stride) { binary_op_dims( a + a_it.loc, b + b_it.loc, @@ -387,6 +387,7 @@ void binary_int_op_cpu( switch (out.dtype()) { case bool_: binary_op(a, b, out, bopt); + break; case uint8: binary_op(a, b, out, bopt); break; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h b/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h index 0bfdc0501..c613bc505 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h @@ -26,6 +26,7 @@ using namespace mlx::core::simd; DEFAULT_BINARY_OP(Add, operator+) DEFAULT_BINARY_OP(ArcTan2, atan2) DEFAULT_BINARY_OP(Divide, divide) +DEFAULT_BINARY_OP(FloorDivide, floor_divide) DEFAULT_BINARY_OP(Multiply, operator*) DEFAULT_BINARY_OP(Subtract, operator-) DEFAULT_BINARY_OP(LogicalAnd, operator&&) diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp index f5b516e44..d6a53a528 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -116,9 +117,14 @@ void* compile( } if (!lib_exists) { - // Open source file and write source code to it - std::string source_file_name = kernel_file_name + ".cpp"; - auto source_file_path = (output_dir / source_file_name).string(); + // Build under a unique name, then move the files into place. Other + // processes then never load a partly written library. + auto tmp_name = fmt::format("tmp_{:08x}", std::random_device{}()); + std::string source_file_name = tmp_name + ".cpp"; + std::string tmp_lib_name = + tmp_name + std::filesystem::path(shared_lib_name).extension().string(); + auto source_file_path = output_dir / source_file_name; + auto tmp_lib_path = output_dir / tmp_lib_name; std::ofstream source_file(source_file_path); source_file << source_code; @@ -127,7 +133,7 @@ void* compile( try { JitCompiler::exec( JitCompiler::build_command( - output_dir, source_file_name, shared_lib_name)); + output_dir, source_file_name, tmp_lib_name)); } catch (const std::exception& error) { throw std::runtime_error( fmt::format( @@ -135,6 +141,19 @@ void* compile( kernel_name, error.what())); } + + std::error_code error; + std::filesystem::rename( + source_file_path, output_dir / (kernel_file_name + ".cpp"), error); + std::filesystem::rename(tmp_lib_path, shared_lib_path, error); + if (error) { + // On Windows the rename can fail while another process uses the file. + if (std::filesystem::exists(shared_lib_path, error)) { + std::filesystem::remove(tmp_lib_path, error); + } else { + shared_lib_path = tmp_lib_path.string(); + } + } } // load library diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp index 17bc6cb9f..442629fbe 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp @@ -299,18 +299,20 @@ void slow_conv_2D( int oH_border_1 = is_idil_one ? ((padding_lo[0] + wt_strides[0] - 1) / wt_strides[0]) : oH; - int oH_border_2 = std::max( - oH_border_1, - (iH + padding_lo[0] - wH * wt_dilation[0]) / wt_strides[0]); + oH_border_1 = std::min(std::max(oH_border_0, oH_border_1), oH); + int oH_border_2 = + (iH + padding_lo[0] - wH * wt_dilation[0]) / wt_strides[0]; + oH_border_2 = std::min(std::max(oH_border_1, oH_border_2), oH); int oH_border_3 = oH; int oW_border_0 = 0; int oW_border_1 = is_idil_one ? ((padding_lo[1] + wt_strides[1] - 1) / wt_strides[1]) : oW; - int oW_border_2 = std::max( - oW_border_1, - (iW + padding_lo[1] - wW * wt_dilation[1]) / wt_strides[1]); + oW_border_1 = std::min(std::max(oW_border_0, oW_border_1), oW); + int oW_border_2 = + (iW + padding_lo[1] - wW * wt_dilation[1]) / wt_strides[1]; + oW_border_2 = std::min(std::max(oW_border_1, oW_border_2), oW); int oW_border_3 = oW; for (int n = 0; n < N; ++n) { @@ -586,27 +588,30 @@ void slow_conv_3D( int oD_border_1 = is_idil_one ? ((padding_lo[0] + wt_strides[0] - 1) / wt_strides[0]) : oD; - int oD_border_2 = std::max( - oD_border_1, - (iD + padding_lo[0] - wD * wt_dilation[0]) / wt_strides[0]); + oD_border_1 = std::min(std::max(oD_border_0, oD_border_1), oD); + int oD_border_2 = + (iD + padding_lo[0] - wD * wt_dilation[0]) / wt_strides[0]; + oD_border_2 = std::min(std::max(oD_border_1, oD_border_2), oD); int oD_border_3 = oD; int oH_border_0 = 0; int oH_border_1 = is_idil_one ? ((padding_lo[1] + wt_strides[1] - 1) / wt_strides[1]) : oH; - int oH_border_2 = std::max( - oH_border_1, - (iH + padding_lo[1] - wH * wt_dilation[1]) / wt_strides[1]); + oH_border_1 = std::min(std::max(oH_border_0, oH_border_1), oH); + int oH_border_2 = + (iH + padding_lo[1] - wH * wt_dilation[1]) / wt_strides[1]; + oH_border_2 = std::min(std::max(oH_border_1, oH_border_2), oH); int oH_border_3 = oH; int oW_border_0 = 0; int oW_border_1 = is_idil_one ? ((padding_lo[2] + wt_strides[2] - 1) / wt_strides[2]) : oW; - int oW_border_2 = std::max( - oW_border_1, - (iW + padding_lo[2] - wW * wt_dilation[2]) / wt_strides[2]); + oW_border_1 = std::min(std::max(oW_border_0, oW_border_1), oW); + int oW_border_2 = + (iW + padding_lo[2] - wW * wt_dilation[2]) / wt_strides[2]; + oW_border_2 = std::min(std::max(oW_border_1, oW_border_2), oW); int oW_border_3 = oW; for (int n = 0; n < N; ++n) { diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp index 381719eb0..ffe290ffa 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp @@ -27,7 +27,13 @@ void copy_vector(const array& src, array& dst) { auto src_ptr = src.data(); auto dst_ptr = dst.data(); auto size = src.data_size(); - std::copy(src_ptr, src_ptr + size, dst_ptr); + if constexpr (std::is_same_v) { + std::copy(src_ptr, src_ptr + size, dst_ptr); + } else { + std::transform(src_ptr, src_ptr + size, dst_ptr, [](SrcT x) { + return static_cast(x); + }); + } } template diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp index 7bf674cfd..109d7dfd0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp @@ -63,12 +63,12 @@ struct EigWork< } void run(T* a, O* values, O* vectors) { - auto eig_tmp = static_cast(buffers[0].buffer.raw_ptr()); + auto eig_tmp = static_cast(buffers[0].buffer().raw_ptr()); T* vec_tmp = nullptr; if (vectors) { - vec_tmp = static_cast(buffers[1].buffer.raw_ptr()); + vec_tmp = static_cast(buffers[1].buffer().raw_ptr()); } - auto work = static_cast(buffers.back().buffer.raw_ptr()); + auto work = static_cast(buffers.back().buffer().raw_ptr()); int n_vecs_l = vectors ? N : 1; int n_vecs_r = 1; @@ -166,9 +166,9 @@ struct EigWork> { &n_vecs_l, nullptr, &n_vecs_r, - static_cast(buffers[0].buffer.raw_ptr()), + static_cast(buffers[0].buffer().raw_ptr()), &lwork, - static_cast(buffers[1].buffer.raw_ptr()), + static_cast(buffers[1].buffer().raw_ptr()), &info); } }; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp index ff5a4805a..41d94205d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp @@ -59,9 +59,9 @@ struct EighWork< vectors, &N, values, - static_cast(buffers[0].buffer.raw_ptr()), + static_cast(buffers[0].buffer().raw_ptr()), &lwork, - static_cast(buffers[1].buffer.raw_ptr()), + static_cast(buffers[1].buffer().raw_ptr()), &liwork, &info); } @@ -116,11 +116,11 @@ struct EighWork> { vectors, &N, values, - static_cast(buffers[0].buffer.raw_ptr()), + static_cast(buffers[0].buffer().raw_ptr()), &lwork, - static_cast(buffers[1].buffer.raw_ptr()), + static_cast(buffers[1].buffer().raw_ptr()), &lrwork, - static_cast(buffers[2].buffer.raw_ptr()), + static_cast(buffers[2].buffer().raw_ptr()), &liwork, &info); if (jobz == 'V') { diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/inverse.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/inverse.cpp index ddc979daa..a02c48b14 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/inverse.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/inverse.cpp @@ -11,14 +11,14 @@ namespace mlx::core { template void general_inv(T* inv, int N) { int info; - auto ipiv = array::Data{allocator::malloc(sizeof(int) * N)}; + auto ipiv = allocator::malloc(sizeof(int) * N); // Compute LU factorization. getrf( /* m = */ &N, /* n = */ &N, /* a = */ inv, /* lda = */ &N, - /* ipiv = */ static_cast(ipiv.buffer.raw_ptr()), + /* ipiv = */ static_cast(ipiv.buffer().raw_ptr()), /* info = */ &info); if (info != 0) { @@ -49,15 +49,15 @@ void general_inv(T* inv, int N) { } const int lwork = workspace_size; - auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)}; + auto scratch = allocator::malloc(sizeof(T) * lwork); // Compute inverse. getri( /* m = */ &N, /* a = */ inv, /* lda = */ &N, - /* ipiv = */ static_cast(ipiv.buffer.raw_ptr()), - /* work = */ static_cast(scratch.buffer.raw_ptr()), + /* ipiv = */ static_cast(ipiv.buffer().raw_ptr()), + /* work = */ static_cast(scratch.buffer().raw_ptr()), /* lwork = */ &lwork, /* info = */ &info); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp index 3299b7471..d2f5481b9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp @@ -35,6 +35,14 @@ static std::pair compute_dynamic_offset( const Strides& strides, const std::vector& axes, Stream stream) { + auto& encoder = cpu::get_command_encoder(stream); + // The kernel reads the indices contiguously + if (!indices.flags().row_contiguous) { + auto copy = contiguous_copy_cpu(indices, stream); + encoder.add_temporary(copy); + return compute_dynamic_offset(copy, strides, axes, stream); + } + array offset({1}, int64, nullptr, {}); bool donate = indices.is_donatable() && (indices.data_size() * indices.itemsize()) >= offset.itemsize(); @@ -44,7 +52,6 @@ static std::pair compute_dynamic_offset( offset.set_data(allocator::malloc(offset.itemsize())); } - auto& encoder = cpu::get_command_encoder(stream); encoder.set_input_array(indices); encoder.set_output_array(offset); auto compute_offset = @@ -380,8 +387,20 @@ void View::eval_cpu(const std::vector& inputs, array& out) { strides[i] *= ibytes; strides[i] /= obytes; } + // Last axis can have any stride + if (obytes < ibytes) { + strides.back() = 1; + } + auto flags = in.flags(); + if (ibytes != obytes) { + // Last axis size can change contiguity + auto [_, row_contiguous, col_contiguous] = + check_contiguity(out.shape(), strides); + flags.row_contiguous = row_contiguous; + flags.col_contiguous = col_contiguous; + } out.copy_shared_buffer( - in, strides, in.flags(), in.data_size() * ibytes / obytes); + in, strides, flags, in.data_size() * ibytes / obytes); } else { auto tmp = array( in.shape(), in.dtype() == bool_ ? uint8 : in.dtype(), nullptr, {}); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp index f438d492a..358544aa8 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp @@ -69,12 +69,11 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) { &N, in_ptr + M * N * i, &lda, - static_cast(tau.raw_ptr()) + num_reflectors * i, - static_cast(work.raw_ptr()), + static_cast(tau.buffer().raw_ptr()) + num_reflectors * i, + static_cast(work.buffer().raw_ptr()), &lwork, &info); } - allocator::free(work); for (int i = 0; i < num_matrices; ++i) { /// num_reflectors x N @@ -102,7 +101,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) { &lwork, &info); lwork = optimal_work; - work = allocator::malloc(sizeof(T) * lwork); + auto orgqr_work = allocator::malloc(sizeof(T) * lwork); // Loop over matrices for (int i = 0; i < num_matrices; ++i) { @@ -113,8 +112,8 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) { &num_reflectors, in_ptr + M * N * i, &lda, - static_cast(tau.raw_ptr()) + num_reflectors * i, - static_cast(work.raw_ptr()), + static_cast(tau.buffer().raw_ptr()) + num_reflectors * i, + static_cast(orgqr_work.buffer().raw_ptr()), &lwork, &info); } @@ -128,10 +127,6 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) { } } } - - // Cleanup - allocator::free(work); - allocator::free(tau); }); encoder.add_temporary(in); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp index 15e00cd91..f029561c7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp +++ b/Source/Cmlx/mlx/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."); } } @@ -1181,15 +1186,24 @@ void quantize( w_min = std::min(w_min, (float)w[w_idx + j]); } bool mask = std::abs(w_min) > std::abs(w_max); - float scale = std::max((w_max - w_min) / n_bins, eps); - scale = mask ? scale : -scale; - - float edge = mask ? w_min : w_max; - float q0 = std::rint(edge / scale); - float bias = 0; - if (q0 != 0) { - scale = edge / q0; - bias = edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = std::max(w_max - w_min, eps); + bias = w_min; + } else { + scale = std::max((w_max - w_min) / n_bins, eps); + scale = mask ? scale : -scale; + + float edge = mask ? w_min : w_max; + float q0 = std::rint(edge / scale); + bias = 0; + if (q0 != 0) { + scale = edge / q0; + bias = edge; + } } size_t out_idx = i * int_per_group; for (int j = 0; j < int_per_group / bytes_per_pack; ++j) { diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp index 04965690b..cd4911e89 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp @@ -108,7 +108,12 @@ void strided_reduce( }; template -void contiguous_reduce(const T* x, U* accumulator, int size, Op op, U init) { +void contiguous_reduce( + const T* x, + U* accumulator, + int64_t size, + Op op, + U init) { constexpr int N = std::min(simd::max_size, simd::max_size); simd::Simd accumulator_v(init); while (size >= N) { @@ -125,11 +130,11 @@ void contiguous_reduce(const T* x, U* accumulator, int size, Op op, U init) { // Helper for the ndimensional strided loop void nd_loop( - std::function callback, + std::function callback, const Shape& shape, const Strides& strides) { - std::function loop_inner; - loop_inner = [&](int dim, int offset) { + std::function loop_inner; + loop_inner = [&](int dim, int64_t offset) { if (dim < shape.size() - 1) { auto size = shape[dim]; auto stride = strides[dim]; @@ -181,16 +186,16 @@ void reduction_op( auto [shape, strides] = shapes_without_reduction_axes(x, axes); if (plan.shape.size() == 0) { for (int i = 0; i < out.size(); i++, out_ptr++) { - int offset = elem_to_loc(i, shape, strides); + int64_t offset = elem_to_loc(i, shape, strides); *out_ptr = init; contiguous_reduce(in_ptr + offset, out_ptr, reduction_size, Op{}, init); } } else { for (int i = 0; i < out.size(); i++, out_ptr++) { - int offset = elem_to_loc(i, shape, strides); + int64_t offset = elem_to_loc(i, shape, strides); *out_ptr = init; nd_loop( - [&](int extra_offset) { + [&](int64_t extra_offset) { contiguous_reduce( in_ptr + offset + extra_offset, out_ptr, @@ -229,7 +234,7 @@ void reduction_op( if (plan.shape.size() == 0) { for (int i = 0; i < out.size(); i += reduction_stride) { - int offset = elem_to_loc(i, shape, strides); + int64_t offset = elem_to_loc(i, shape, strides); std::fill_n(out_ptr, reduction_stride, init); strided_reduce( in_ptr + offset, out_ptr, reduction_size, reduction_stride, Op{}); @@ -237,10 +242,10 @@ void reduction_op( } } else { for (int i = 0; i < out.size(); i += reduction_stride) { - int offset = elem_to_loc(i, shape, strides); + int64_t offset = elem_to_loc(i, shape, strides); std::fill_n(out_ptr, reduction_stride, init); nd_loop( - [&](int extra_offset) { + [&](int64_t extra_offset) { strided_reduce( in_ptr + offset + extra_offset, out_ptr, @@ -260,10 +265,10 @@ void reduction_op( auto [shape, strides] = shapes_without_reduction_axes(x, axes); for (int i = 0; i < out.size(); i++, out_ptr++) { - int offset = elem_to_loc(i, shape, strides); + int64_t offset = elem_to_loc(i, shape, strides); U val = init; nd_loop( - [&](int extra_offset) { + [&](int64_t extra_offset) { val = Op{}(val, *(in_ptr + offset + extra_offset)); }, plan.shape, @@ -293,7 +298,7 @@ void float_reduction( struct AndReduce { template bool operator()(bool x, T y) { - return x & (y != 0); + return x & static_cast(y); } bool operator()(bool x, bool y) { @@ -319,7 +324,7 @@ struct AndReduce { struct OrReduce { template bool operator()(bool x, T y) { - return x | (y != 0); + return x | static_cast(y); } bool operator()(bool x, bool y) { @@ -517,12 +522,11 @@ void Reduce::eval_cpu(const std::vector& inputs, array& out) { break; case uint64: case int64: - // 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 complex64: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case float64: reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h index 5d6128328..649365373 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h @@ -243,6 +243,27 @@ Simd divide(Simd a, Simd b) { } } +// Signed types round the quotient towards minus infinity. +template +Simd floor_divide(Simd a, Simd b) { + if constexpr (std::is_integral_v) { + // Swap zero divisors for 1 before dividing so the quotient is defined. + auto zero = b == Simd(0); + auto q = a / select(zero, Simd(1), b); + q = select(zero, Simd(0), q); + if constexpr (is_signed_v) { + // a - b * q is the remainder and needs no division by zero. + auto mask = b != 0 && (a - b * q) != 0 && (a < 0 != b < 0); + q = select(mask, q - Simd(1), q); + } + return q; + } else if constexpr (is_complex) { + return a / b; + } else { + return floor(divide(a, b)); + } +} + template Simd remainder(Simd a, Simd b) { Simd r; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h index 626dd93df..b24bc9d01 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h @@ -218,6 +218,29 @@ Simd divide(Simd a_, Simd b_) { return a / b; } +template +Simd floor_divide(Simd a_, Simd b_) { + T a = a_.value; + T b = b_.value; + if constexpr (std::is_integral_v) { + // Integer division by zero gives 0. See divide above. + if (b == 0) { + return T(0); + } + T q = a / b; + if constexpr (is_signed_v) { + if (a % b != 0 && (a < 0 != b < 0)) { + q -= 1; + } + } + return q; + } else if constexpr (is_complex) { + return a / b; + } else { + return floor(divide(a_, b_)); + } +} + template Simd remainder(Simd a_, Simd b_) { T a = a_.value; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h index 5d32042cc..ad258296e 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h @@ -80,7 +80,7 @@ DEFINE_NEON_UNARY_OP(rint, vrndnq_f16) inline Simd operator!(Simd v) { auto out = vceqzq_f16(v.value); - return Simd(*(uint16_t*)&out); + return load(reinterpret_cast(&out)); } inline Simd operator-(Simd v) { @@ -94,21 +94,21 @@ DEFINE_NEON_BINARY_OP(operator-, vsubq_f16) DEFINE_NEON_BINARY_OP(operator*, vmulq_f16) DEFINE_NEON_BINARY_OP(operator/, vdivq_f16) -#define DEFINE_NEON_COMPARISON(Op, op) \ - template \ - Simd operator Op(Simd a, T b) { \ - auto out = op(a.value, Simd(b).value); \ - return Simd(*(uint16_t*)(&out)); \ - } \ - template \ - Simd operator Op(T a, Simd b) { \ - auto out = op(Simd(a).value, b.value); \ - return Simd(*(uint16_t*)(&out)); \ - } \ - inline Simd operator Op( \ - Simd a, Simd b) { \ - auto out = op(a.value, b.value); \ - return Simd(*(uint16_t*)(&out)); \ +#define DEFINE_NEON_COMPARISON(Op, op) \ + template \ + Simd operator Op(Simd a, T b) { \ + auto out = op(a.value, Simd(b).value); \ + return load(reinterpret_cast(&out)); \ + } \ + template \ + Simd operator Op(T a, Simd b) { \ + auto out = op(Simd(a).value, b.value); \ + return load(reinterpret_cast(&out)); \ + } \ + inline Simd operator Op( \ + Simd a, Simd b) { \ + auto out = op(a.value, b.value); \ + return load(reinterpret_cast(&out)); \ } DEFINE_NEON_COMPARISON(==, vceqq_f16) diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp index 80101e907..c4084ed4f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp @@ -55,7 +55,7 @@ struct SVDWork< /* ldvt = */ &ldvt, /* work = */ &workspace_dimension, /* lwork = */ &lwork_query, - /* iwork = */ static_cast(buffers[0].buffer.raw_ptr()), + /* iwork = */ static_cast(buffers[0].buffer().raw_ptr()), /* info = */ &info); if (info != 0) { @@ -84,9 +84,9 @@ struct SVDWork< // According to the identity above, lapack will write Uᵀ as Vᵀ. /* vt = */ vt, /* ldvt = */ &ldvt, - /* work = */ static_cast(buffers[1].buffer.raw_ptr()), + /* work = */ static_cast(buffers[1].buffer().raw_ptr()), /* lwork = */ &lwork, - /* iwork = */ static_cast(buffers[0].buffer.raw_ptr()), + /* iwork = */ static_cast(buffers[0].buffer().raw_ptr()), /* info = */ &info); if (info != 0) { @@ -143,8 +143,8 @@ struct SVDWork> { /* ldvt = */ &ldvt, /* work = */ &workspace_dimension, /* lwork = */ &lwork_query, - /* rwork = */ static_cast(buffers[1].buffer.raw_ptr()), - /* iwork = */ static_cast(buffers[0].buffer.raw_ptr()), + /* rwork = */ static_cast(buffers[1].buffer().raw_ptr()), + /* iwork = */ static_cast(buffers[0].buffer().raw_ptr()), /* info = */ &info); if (info != 0) { @@ -173,10 +173,10 @@ struct SVDWork> { // According to the identity above, lapack will write Uᵀ as Vᵀ. /* vt = */ vt, /* ldvt = */ &ldvt, - /* work = */ static_cast(buffers[2].buffer.raw_ptr()), + /* work = */ static_cast(buffers[2].buffer().raw_ptr()), /* lwork = */ &lwork, - /* rwork = */ static_cast(buffers[1].buffer.raw_ptr()), - /* iwork = */ static_cast(buffers[0].buffer.raw_ptr()), + /* rwork = */ static_cast(buffers[1].buffer().raw_ptr()), + /* iwork = */ static_cast(buffers[0].buffer().raw_ptr()), /* info = */ &info); if (info != 0) { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp index 5f4b9f2d5..bb262d0e6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp @@ -396,9 +396,10 @@ CudaAllocator& allocator() { return *allocator_; } -Buffer malloc_async(size_t size, CommandEncoder& encoder) { - return allocator().malloc_async( - size, encoder.device().cuda_device(), encoder.stream()); +allocator::Data malloc_async(size_t size, CommandEncoder& encoder) { + return allocator::Data( + allocator().malloc_async( + size, encoder.device().cuda_device(), encoder.stream())); } } // namespace cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.h b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.h index af76ad905..f0416c083 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.h @@ -89,6 +89,6 @@ class CudaAllocator : public allocator::Allocator { CudaAllocator& allocator(); -Buffer malloc_async(size_t size, CommandEncoder& encoder); +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/binary/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/cuda/binary/CMakeLists.txt index bda289de7..db8326505 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/binary/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/cuda/binary/CMakeLists.txt @@ -5,6 +5,7 @@ target_sources( PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/bitwise_binary.cu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/divide.cu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/equal.cu + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/floor_divide.cu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/greater.cu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/greater_equal.cu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/less.cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh index 6130f63c7..e7445b287 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh @@ -209,9 +209,10 @@ __global__ void binary_g( template constexpr bool supports_binary_op() { if (std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v) { + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v) { return std::is_same_v; } if (std::is_same_v || std::is_same_v || diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/binary/floor_divide.cu b/Source/Cmlx/mlx/mlx/backend/cuda/binary/floor_divide.cu new file mode 100644 index 000000000..337c71af8 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/binary/floor_divide.cu @@ -0,0 +1,7 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/binary/binary.cuh" + +namespace mlx::core { +BINARY_GPU(FloorDivide) +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device.h b/Source/Cmlx/mlx/mlx/backend/cuda/device.h index 198f0b5ad..1ec8f20fa 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device.h @@ -18,7 +18,7 @@ std::pair subgraph_to_key(cudaGraph_t graph); class Worker; -class CommandEncoder { +class MLX_API CommandEncoder { public: struct CaptureContext { CaptureContext(CommandEncoder& enc); @@ -165,7 +165,7 @@ class CommandEncoder { int max_mb_per_graph_; }; -class Device { +class MLX_API Device { public: explicit Device(int device); ~Device(); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/cast_op.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/cast_op.cuh index b85a40c83..c2dfdd4f5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/cast_op.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/cast_op.cuh @@ -25,7 +25,7 @@ struct CastOp, bool> { static constexpr bool is_castable = true; __device__ bool operator()(complex_t x) { - return x.real() != 0 && x.imag() != 0; + return x.real() != 0 || x.imag() != 0; } }; @@ -34,7 +34,7 @@ struct CastOp> { static constexpr bool is_castable = true; __device__ complex_t operator()(bool x) { - return x ? complex_t{1, 1} : complex_t{0, 0}; + return x ? complex_t{1, 0} : complex_t{0, 0}; } }; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh index e6efe0402..476c2339c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh @@ -267,6 +267,9 @@ CUTE_DEVICE void qmm_sm80_mainloop( for (int i = 0; i < size(tCrC_accu); i++) { tCrC(i) = Element(tCrC_accu(i)); } + // The last prefetches may still be landing in the buffers C reuses. + cp_async_wait<0>(); + __syncthreads(); copy(r2s_copy_c, r2s_tCrC, r2s_tCsC); __syncthreads(); copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC); @@ -355,4 +358,116 @@ void qmm_sm80_kernel( thread_idx); } +template +__global__ +__launch_bounds__(decltype(size(make_tiled_mma()))::value) +void qmm_sm80_rhs_kernel( + const Element* A, + const Quant* B, + const Scale* S, + const Element* Z, + const uint32_t* rhs_indices, + Element* C, + int m, + int n, + int k, + int group_count) { + int thread_idx = int(threadIdx.x); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + + auto cta_tiler = CtaTiler{}; + int tile_m = int(size<0>(cta_tiler)); + int y_row = m_coord * tile_m; + int m_max_coord = m - y_row; + int tgp_bm = m_max_coord < tile_m ? m_max_coord : tile_m; + + // Define layouts (mixed). All gathered rows form one matrix; each run of + // equal rhs_indices inside the row tile reuses one expert tile. + auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) + auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) + auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) + auto S_layout = make_scales_layout(n, k, group_count); + + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, 1), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, group_count), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, 1), dC); // (M,N,L) + + Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) + Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) + + Tensor mA = mA_mkl(_,_,0); // (M,K) + Tensor mC = mC_mnl(_,_,0); // (M,N) + uint32_t index_next = rhs_indices[y_row]; + int offset_next = 0; + while (offset_next < tgp_bm) { + int offset = offset_next; + uint32_t b_batch = index_next; + offset_next = tgp_bm; + for (int row = offset + 1; row < tgp_bm; ++row) { + uint32_t index = rhs_indices[y_row + row]; + if (index != b_batch) { + offset_next = row; + index_next = index; + break; + } + } + + int run_rows = offset_next - offset; + Tensor mA_offset = + domain_offset(make_coord(y_row + offset, 0), mA); + Tensor mC_offset = + domain_offset(make_coord(y_row + offset, 0), mC); + Tensor mB = mB_nkl(_,_,b_batch); // (N,K) + Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + + auto run_qmm = [&](auto run_tiler) { + Tensor gA = local_tile(mA_offset, run_tiler, make_coord(0, n_coord, _), Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC_offset, run_tiler, make_coord(0, n_coord, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + Tensor gS = local_tile(mS, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gZ = local_tile(mZ, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + + qmm_sm80_mainloop( + run_tiler, + gA, + gB, + gS, + gZ, + gC, + run_rows, + thread_idx); + }; + + constexpr int tile_rows = decltype(size<0>(cta_tiler))::value; + if constexpr (tile_rows > 32) { + if (run_rows <= 16) { + run_qmm(make_shape( + Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else if (run_rows <= 32) { + run_qmm(make_shape( + Int<32>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else { + run_qmm(cta_tiler); + } + } else if constexpr (tile_rows > 16) { + if (run_rows <= 16) { + run_qmm(make_shape( + Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else { + run_qmm(cta_tiler); + } + } else { + run_qmm(cta_tiler); + } + // The next run's loads must not overwrite C while it is still being stored. + __syncthreads(); + } +} + } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu index bcea00544..aef509626 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu @@ -340,7 +340,8 @@ void Event::wait() { } else { event.atomic->wait(value()); } - CHECK_CUDA_ERROR(cudaPeekAtLastError()); + // Check for errors during kernel execution and reset the error state. + CHECK_CUDA_ERROR(cudaGetLastError()); check_error(); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp index e775d73ec..1b29044df 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp @@ -45,8 +45,7 @@ uint32_t Fence::update(Stream s, const array& a, bool cross_device) { auto& f = cast(); if (cross_device) { // Move to managed memory if there is a device switch - auto& cbuf = - *static_cast(const_cast(a).buffer().ptr()); + auto& cbuf = *static_cast(a.buffer().ptr()); if (cbuf.device != -1) { auto& encoder = cu::get_command_encoder(s); encoder.commit(); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp index 68f5ea139..da461651b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp @@ -459,7 +459,9 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { return; } - array mask_flat = flatten_in_eval(mask, 1, -1, s); + array mask_flat = (mask.ndim() == 1) + ? reshape_in_eval(mask, {mask.shape(0), 1}, s) + : flatten_in_eval(mask, 1, -1, s); if (gpu_ptr(mask_flat) != gpu_ptr(mask)) { encoder.add_temporary(mask_flat); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h b/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h index d208315a0..02bebe673 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h @@ -76,6 +76,7 @@ class LRUCache { auto it = map_.find(key); if (it == map_.end()) return end(); + cache_misses_ = 0; vlist_.splice(vlist_.begin(), vlist_, it->second); return it->second; } @@ -84,6 +85,7 @@ class LRUCache { std::pair emplace(const K& key, U&& value) { auto it = map_.find(key); if (it != map_.end()) { + cache_misses_ = 0; vlist_.splice(vlist_.begin(), vlist_, it->second); return {it->second, false}; } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp index c02357b08..104c7fd3f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp @@ -34,6 +34,15 @@ bool fast::GatedDeltaUpdate::use_fallback( return true; } +bool fast::GatedDeltaUpdateVJP::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s) { + return true; +} + NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) NO_GPU_MULTI(SVD) @@ -43,7 +52,8 @@ NO_GPU_MULTI(Eigh) namespace fast { NO_GPU_MULTI(GatedDeltaUpdate) -} +NO_GPU_MULTI(GatedDeltaUpdateVJP) +} // namespace fast namespace distributed { NO_GPU_MULTI(Send) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu index 41de03e85..4e148744c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/backend/common/quantized.h" #include "mlx/backend/cuda/device/utils.cuh" @@ -64,6 +64,7 @@ __device__ void fp_qmv_impl( const uint8_t* scales_, const T* vec, T* out, + float scale, int rows, int cols) { auto block = cg::this_thread_block(); @@ -150,7 +151,7 @@ __device__ void fp_qmv_impl( sum = cg::reduce(warp, sum, cg::plus{}); if (warp.thread_rank() == 0) { - out[row] = static_cast(sum); + out[row] = static_cast(sum * scale); } } } @@ -170,7 +171,7 @@ __global__ void fp_qmv_single( int rows, int cols) { fp_qmv_impl( - mat, scales, vec, out, rows, cols); + mat, scales, vec, out, 1.0f, rows, cols); } template < @@ -208,7 +209,39 @@ __global__ void fp_qmv_batched( mat_strides, scales_strides); fp_qmv_impl( - mat, scales, vec, out, rows, cols); + mat, scales, vec, out, 1.0f, rows, cols); +} + +template < + typename T, + int rows_per_block, + int n_per_thread, + int bits, + int group_size, + bool use_mx_scale> +__global__ void fp_gather_qmv( + const uint32_t* mat, + const uint8_t* scales, + const float* global_scale, + const T* vec, + T* out, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, + int rows, + int cols) { + constexpr int vals_per_item = bits == 8 ? 4 : 8; + uint32_t idx = cg::this_grid().block_index().z; + int64_t mat_size = static_cast(rows) * (cols / vals_per_item); + int64_t scales_size = static_cast(rows) * (cols / group_size); + mat += rhs_indices[idx] * mat_size; + scales += rhs_indices[idx] * scales_size; + vec += static_cast(lhs_indices[idx]) * cols; + out += static_cast(idx) * rows; + float scale = global_scale + ? global_scale[rhs_indices[idx]] / (F8E4M3_MAX * F4E2M1_MAX) + : 1.0f; + fp_qmv_impl( + mat, scales, vec, out, scale, rows, cols); } } // namespace cu @@ -228,6 +261,24 @@ void dispatch_1_2_4(int n, F&& f) { } } +// Packed words loaded per thread, limited by the alignment of the inputs. +template +int fp_qmv_n_per_thread( + const uint32_t* mat, + const T* vec, + uint32_t K, + int bits) { + if (K % 32 == 0 && cu::is_aligned<4>(mat) && + ((bits == 4 && cu::is_aligned<8>(vec)) || cu::is_aligned<4>(vec))) { + return 4; + } + if (cu::is_aligned<2>(mat) && + ((bits == 4 && cu::is_aligned<4>(vec)) || cu::is_aligned<2>(vec))) { + return 2; + } + return 1; +} + void fp_qmv( const array& x, const array& w, @@ -259,17 +310,7 @@ void fp_qmv( uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block; const uint32_t* mat_ptr = gpu_ptr(mat); const T* vec_ptr = gpu_ptr(vec); - int n = 1; - if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) && - ((bits == 4 && cu::is_aligned<8>(vec_ptr)) || - cu::is_aligned<4>(vec_ptr))) { - n = 4; - } else if ( - cu::is_aligned<2>(mat_ptr) && - ((bits == 4 && cu::is_aligned<4>(vec_ptr)) || - cu::is_aligned<2>(vec_ptr))) { - n = 2; - } + int n = fp_qmv_n_per_thread(mat_ptr, vec_ptr, K, bits); dispatch_1_2_4(n, [&](auto n) { if (B == 1) { auto kernel = @@ -323,4 +364,95 @@ void fp_qmv( }); } +bool supports_fp_gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + constexpr uint32_t max_grid_yz_dim = 65535; + if (mode == QuantizationMode::Affine || !transpose) { + return false; + } + if (out.shape(-2) != 1) { + return false; + } + if (!w.flags().row_contiguous || !scales.flags().row_contiguous) { + return false; + } + int k = x.shape(-1); + if (k < 512 || k % 32 != 0) { + return false; + } + uint32_t n = out.shape(-1); + if (n > rows_per_block * max_grid_yz_dim || + out.size() / n > max_grid_yz_dim) { + return false; + } + return true; +} + +void fp_gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& global_scale, + const array& lhs_indices, + const array& rhs_indices, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder) { + uint32_t N = out.shape(-1); + uint32_t K = x.shape(-1); + uint32_t B = out.size() / N; + + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (global_scale) { + encoder.set_input_array(*global_scale); + } + encoder.set_input_array(x); + encoder.set_input_array(lhs_indices); + encoder.set_input_array(rhs_indices); + encoder.set_output_array(out); + dispatch_float_types(out.dtype(), "gather_qmv", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + dim3 block_dims{WARP_SIZE, rows_per_block}; + uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block; + const uint32_t* mat_ptr = gpu_ptr(w); + const T* vec_ptr = gpu_ptr(x); + int n = fp_qmv_n_per_thread(mat_ptr, vec_ptr, K, bits); + dispatch_1_2_4(n, [&](auto n) { + auto kernel = + cu::fp_gather_qmv; + if (bits == 8) { + kernel = cu::fp_gather_qmv; + } else if (group_size == 16) { + kernel = cu::fp_gather_qmv; + } + encoder.add_kernel_node( + kernel, + {1, blocks_y, B}, + block_dims, + mat_ptr, + gpu_ptr(scales), + global_scale ? gpu_ptr(*global_scale) : nullptr, + vec_ptr, + gpu_ptr(out), + gpu_ptr(lhs_indices), + gpu_ptr(rhs_indices), + N, + K); + }); + } + }); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h index 6991baad1..7c567ed24 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h @@ -57,6 +57,30 @@ void qmm_sm80( QuantizationMode mode, cu::CommandEncoder& encoder); +bool supports_gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + +void gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& rhs_indices, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder); + bool supports_qmm_naive( const array& x, const array& w, @@ -106,6 +130,30 @@ void fp_qmv( cu::CommandEncoder& encoder, Stream s); +bool supports_fp_gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + +void fp_gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& global_scale, + const array& lhs_indices, + const array& rhs_indices, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder); + bool supports_qmv( const array& x, const array& w, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu index 981823956..4ce2817d3 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu @@ -21,6 +21,10 @@ inline auto make_cta_tiler(int m, int group_size) { return cute::make_shape(tile_m, tile_n, tile_k); } +inline int get_quantized_batch_count(const array& w) { + return w.ndim() > 2 ? w.size() / (w.shape(-1) * w.shape(-2)) : 1; +} + } // namespace void qmm_sm80( @@ -113,4 +117,125 @@ void qmm_sm80( broadcast_b); } +bool supports_gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + int out_m = out.ndim() > 1 ? out.shape(-2) : 1; + if (out_m != 1) { + return false; + } + if (!supports_qmm_sm80( + x, + w, + scales, + biases, + out, + transpose, + bits, + group_size, + mode, + device)) { + return false; + } + + int rows = out.size() / out.shape(-1); + int input_rows = x.size() / x.shape(-1); + int group_count = get_quantized_batch_count(w); + return input_rows == rows && rows >= 16 && group_count > 0 && + rows / group_count >= 4; +} + +void gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& rhs_indices, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + int m = x.size() / x.shape(-1); + int n = out.shape(-1); + int k = x.shape(-1); + int group_count = get_quantized_batch_count(w); + auto cta_tiler = make_cta_tiler(m, group_size); + + std::string module_name = fmt::format( + "gather_qmm_rhs_sm80_tn_{}_m{}_b{}_g{}_{}", + dtype_to_string(x.dtype()), + cute::size<0>(cta_tiler), + bits, + group_size, + quantization_mode_to_string(mode)); + + auto [ctype_x, ctype_q, ctype_s] = get_qmm_cutlass_types(x, bits, mode); + std::string kernel_name = fmt::format( + "mlx::core::cu::qmm_sm80_rhs_kernel<{}, {}, {}, {}, {}>", + group_size, + ctype_x, + ctype_q, + ctype_s, + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_qmm_sm80, std::vector{kernel_name}); + }); + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + encoder.set_input_array(rhs_indices); + encoder.set_output_array(out); + + dim3 num_blocks{ + uint32_t(cute::ceil_div(m, cute::size<0>(cta_tiler))), + uint32_t(cute::ceil_div(n, cute::size<1>(cta_tiler))), + 1}; + dim3 block_dims{uint32_t(cute::size(cu::make_tiled_mma()))}; + + auto [sA_layout, sB_layout, sC_layout] = cu::make_smem_layouts(cta_tiler); + size_t smem_bytes = std::max( + cute::cosize(sA_layout) * x.itemsize() + + cute::cosize(sB_layout) * bits / 8, + cute::cosize(sC_layout) * x.itemsize()); + + auto kernel = mod.get_kernel(kernel_name, [&](CUfunction kernel) { + if (smem_bytes > 48000) { + cuFuncSetAttribute( + kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_bytes); + } + }); + + encoder.add_kernel_node_ex( + kernel, + num_blocks, + block_dims, + {}, + smem_bytes, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + gpu_ptr(rhs_indices), + gpu_ptr(out), + m, + n, + k, + group_count); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp index 5c3e25580..a69c49d8f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp @@ -194,8 +194,13 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { // naive kernel until it does. bool can_use_qmm_sm80 = !global_scale.has_value() && supports(supports_qmm_sm80); + bool can_use_gather_qmm_rhs_sm80 = mode_ != QuantizationMode::Affine && + right_sorted_ && !global_scale.has_value() && + supports(supports_gather_qmm_rhs_sm80); bool can_use_qmm_naive = supports(supports_qmm_naive); - bool can_use_qmv = supports(supports_qmv); + bool can_use_fp_gather_qmv = supports(supports_fp_gather_qmv); + bool can_use_qmv = supports(supports_qmv) || can_use_fp_gather_qmv; + bool prefer_qmv = can_use_fp_gather_qmv || (can_use_qmv && M * B < 8); auto call_qmm_sm80 = [&]() { out.set_data(cu::malloc_async(out.nbytes(), encoder)); @@ -231,6 +236,20 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { }; auto call_qmv = [&]() { out.set_data(cu::malloc_async(out.nbytes(), encoder)); + if (can_use_fp_gather_qmv) { + fp_gather_qmv( + x, + w, + scales, + global_scale, + lhs_indices, + rhs_indices, + out, + bits_, + group_size_, + encoder); + return; + } gather_qmv( x, w, @@ -246,8 +265,24 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { encoder); }; + if (can_use_gather_qmm_rhs_sm80) { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + gather_qmm_rhs_sm80( + x, + w, + scales, + biases, + rhs_indices, + out, + bits_, + group_size_, + mode_, + encoder); + return; + } + if (can_use_qmm_sm80) { - if (can_use_qmv && (M * B < 8)) { + if (prefer_qmv) { call_qmv(); } else { call_qmm_sm80(); @@ -256,7 +291,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { } if (can_use_qmm_naive) { - if (can_use_qmv && (M * B < 8)) { + if (prefer_qmv) { call_qmv(); } else { call_qmm_naive(); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu index 8a6327bf3..3ff033f56 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu @@ -375,8 +375,7 @@ void col_reduce_two_pass( cu::malloc_async(intermediate.nbytes(), encoder), data_size, intermediate_strides, - fl, - allocator::free); + fl); encoder.add_temporary(intermediate); encoder.set_input_array(in); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh index fdbe72337..00b4d10f2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh @@ -135,11 +135,7 @@ inline void allocate_same_layout( fl.col_contiguous = cc; fl.contiguous = true; out.set_data( - cu::malloc_async(out.nbytes(), encoder), - data_size, - final_strides, - fl, - allocator::free); + cu::malloc_async(out.nbytes(), encoder), data_size, final_strides, fl); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp index 286d500e2..e6c8abf12 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/backend/cuda/cudnn_utils.h" #include "mlx/backend/cuda/device.h" @@ -583,6 +583,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu index c2ba87527..a3666d46f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu @@ -474,8 +474,9 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { in.flags()); } } else { - in = contiguous_copy_gpu(in, s); - out.copy_shared_buffer(in); + // Output type can be wider than input + copy_gpu(in, out, CopyType::General, s); + in = out; } scan_gpu_inplace(in, out, reduce_type_, axis_, reverse_, inclusive_, s); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp index 525cccc26..6a0b11ba5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp @@ -50,6 +50,14 @@ array compute_dynamic_offset( const Strides& strides, const std::vector& axes, const Stream& s) { + auto& encoder = cu::get_command_encoder(s); + // The kernel reads the indices contiguously + if (!indices.flags().row_contiguous) { + auto copy = contiguous_copy_gpu(indices, s); + encoder.add_temporary(copy); + return compute_dynamic_offset(copy, strides, axes, s); + } + Dtype dtype = indices.dtype(); int nidx = axes.size(); @@ -60,8 +68,6 @@ array compute_dynamic_offset( dtype_to_cuda_type(dtype), nidx); - auto& encoder = cu::get_command_encoder(s); - cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::string source = R"( #include "mlx/backend/cuda/device/utils.cuh" diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu index 5d73681a6..980cab3e3 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu @@ -4,6 +4,7 @@ #include #include +#include "mlx/backend/common/utils.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/device/fp16_math.cuh" #include "mlx/backend/cuda/kernel_utils.cuh" @@ -781,16 +782,24 @@ void single_block_sort( int64_t in_stride_sorted_axis = in.strides()[axis]; int64_t out_stride_sorted_axis = out.strides()[axis]; + // The contiguous kernel walks the rows with a single stride, so the axes + // that are not sorted have to collapse to a single dim. + auto single_dim = + [](const Shape& nc_shape, const Strides& nc_str, int64_t& stride) { + auto [shape, strides] = collapse_contiguous_dims(nc_shape, nc_str); + if (shape.size() > 1) { + return false; + } + if (shape.size() > 0) { + stride = strides[0]; + } + return true; + }; + int64_t in_stride_segment_axis = 0; + int64_t out_stride_segment_axis = 0; bool contiguous = in.flags().contiguous; - auto check_strides = [](const array& x, int64_t sort_stride) { - int64_t min_stride = - *std::min_element(x.strides().begin(), x.strides().end()); - int64_t max_stride = - *std::max_element(x.strides().begin(), x.strides().end()); - return sort_stride == min_stride || sort_stride == max_stride; - }; - contiguous &= check_strides(in, in_stride_sorted_axis); - contiguous &= check_strides(out, out_stride_sorted_axis); + contiguous &= single_dim(nc_shape, in_nc_str, in_stride_segment_axis); + contiguous &= single_dim(nc_shape, out_nc_str, out_stride_segment_axis); auto& encoder = cu::get_command_encoder(s); out.set_data(cu::malloc_async(out.nbytes(), encoder)); @@ -817,19 +826,9 @@ void single_block_sort( ARG_SORT, BLOCK_THREADS, N_PER_THREAD>; - int64_t in_stride_segment_axis = INT64_MAX; - int64_t out_stride_segment_axis = INT64_MAX; - for (int i = 0; i < nc_shape.size(); i++) { - if (nc_shape[i] == 1) { - continue; - } - if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) { - throw std::runtime_error("[Sort::eval_gpu] Stride too large."); - } - in_stride_segment_axis = - std::min(in_stride_segment_axis, in_nc_str[i]); - out_stride_segment_axis = - std::min(out_stride_segment_axis, out_nc_str[i]); + if (in_stride_segment_axis > INT32_MAX || + out_stride_segment_axis > INT32_MAX) { + throw std::runtime_error("[Sort::eval_gpu] Stride too large."); } encoder.add_kernel_node( kernel, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp index 82272b74b..f81f589ff 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp @@ -12,6 +12,8 @@ namespace mlx::core { void check_cuda_error(const char* name, cudaError_t err) { if (err != cudaSuccess) { + // Clear the error so it does not resurface in later checks. + cudaGetLastError(); throw std::runtime_error( fmt::format("{} failed: {}", name, cudaGetErrorString(err))); } diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp index 268d6290b..bf17029ff 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp @@ -246,15 +246,29 @@ void View::eval_gpu(const std::vector& inputs, array& out) { // - type size is the same // - type size is smaller and the last axis is contiguous // - the entire array is row contiguous - if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || - in.flags().row_contiguous) { + // and the offset is a multiple of the output type size + if (in.offset() % obytes == 0 && + (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || + in.flags().row_contiguous)) { auto strides = in.strides(); for (int i = 0; i < static_cast(strides.size()) - 1; ++i) { strides[i] *= ibytes; strides[i] /= obytes; } + // Last axis can have any stride + if (obytes < ibytes) { + strides.back() = 1; + } + auto flags = in.flags(); + if (ibytes != obytes) { + // Last axis size can change contiguity + auto [_, row_contiguous, col_contiguous] = + check_contiguity(out.shape(), strides); + flags.row_contiguous = row_contiguous; + flags.col_contiguous = col_contiguous; + } out.copy_shared_buffer( - in, strides, in.flags(), in.data_size() * ibytes / obytes); + in, strides, flags, in.data_size() * ibytes / obytes); } else { auto tmp = array(in.shape(), in.dtype(), nullptr, {}); tmp.set_data(allocator::malloc(tmp.nbytes())); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt index bc5c2fcf7..cccaad8a5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt @@ -97,13 +97,15 @@ if(MLX_METAL_JIT) 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(steel/gemm/kernels/steel_gemm_thin_nax) make_jit_source(quantized_nax kernels/quantized_utils.h) make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h kernels/fp4.h) make_jit_source(steel/attn/kernels/steel_attention_nax) - make_jit_source(gated_delta_update_nax) + make_jit_source(gated_delta_update_nax kernels/gated_delta_nax_ops.h) + make_jit_source(gated_delta_update_nax_vjp kernels/gated_delta_nax_ops.h) else() message( diff --git a/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp b/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp index 7d6b2a713..a3b3e47d5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp @@ -220,6 +220,7 @@ BINARY_GPU(Divide) BINARY_GPU_MULTI(DivMod) BINARY_GPU(Remainder) BINARY_GPU(Equal) +BINARY_GPU(FloorDivide) BINARY_GPU(Greater) BINARY_GPU(GreaterEqual) BINARY_GPU(Less) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp b/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp index 509552261..59acdc6ae 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp @@ -350,8 +350,13 @@ void implicit_gemm_conv_2D_gpu( wn = 1; } + // The 8 channel loader reads one row per thread, so it needs bm = 64 + if (C_per_group == 8) { + bm = 64; + } + int tn = (implicit_N + bn - 1) / bn; - int tm = (implicit_M + bm - 1) / bm; + int tm = (static_cast(implicit_M) + bm - 1) / bm; int swizzle_log = 0; // Fix small channel specialization @@ -365,6 +370,9 @@ void implicit_gemm_conv_2D_gpu( } else if (C_per_group <= 4) { gemm_k_iters = ((conv_params.wS[0] * conv_params.wS[1] * 4) + bk - 1) / bk; n_channel_specialization = C_per_group; + } else if (C_per_group == 8) { + gemm_k_iters = ((conv_params.wS[0] * conv_params.wS[1] * 8) + bk - 1) / bk; + n_channel_specialization = 8; } bool small_filter = (!n_channel_specialization) && @@ -546,7 +554,7 @@ void implicit_gemm_conv_2D_general_gpu( int bk = 16; int tn = (implicit_N + bn - 1) / bn; - int tm = (adj_implicit_m + bm - 1) / bm; + int tm = (static_cast(adj_implicit_m) + bm - 1) / bm; int swizzle_log = 0; // Get channel iteration info @@ -670,7 +678,7 @@ void implicit_gemm_conv_3D_gpu( } int tn = (implicit_N + bn - 1) / bn; - int tm = (implicit_M + bm - 1) / bm; + int tm = (static_cast(implicit_M) + bm - 1) / bm; int swizzle_log = 0; bool small_filter = @@ -782,6 +790,7 @@ void pad_and_slice_conv_3D_gpu( array x_copy(xshape, x.dtype(), nullptr, {}); array zero(0, x.dtype()); pad_gpu(x, zero, x_copy, {0, -1}, {0, 0}, s); + metal::get_command_encoder(s).add_temporary(zero); metal::get_command_encoder(s).add_temporary(x_copy); return x_copy; @@ -1385,6 +1394,7 @@ void pad_in_channels_conv_2D_gpu( 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(zero); metal::get_command_encoder(s).add_temporary(x_copy); return x_copy; @@ -1421,11 +1431,11 @@ void dispatch_conv_2D_gpu( conv_params.wS[0] <= 7 && conv_params.wS[1] <= 7 && conv_params.str[0] <= 2 && conv_params.str[1] <= 2 && conv_params.wt_strides[1] == conv_params.wS[1] && - conv_params.C % 16 == 0 && conv_params.C == conv_params.O) { + conv_params.C % 8 == 0 && conv_params.C == conv_params.O) { return depthwise_conv_2D_gpu(s, d, in, wt, out, conv_params); } - if ((C_per_group <= 4 || C_per_group % 16 == 0) && + if ((C_per_group <= 4 || C_per_group == 8 || C_per_group % 16 == 0) && (O_per_group <= 16 || O_per_group % 16 == 0)) { return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params); } else { @@ -1442,7 +1452,7 @@ void dispatch_conv_2D_gpu( if (!conv_params.flip && is_stride_one && is_kdil_one && is_idil_one && conv_params.wS[0] == 3 && conv_params.wS[1] == 3 && conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large && - channels_large) { + channels_large && env::get_var("MLX_CONV_WINOGRAD", 1)) { // Only use winograd conv when having enough memory. if (int n_step = winograd_batch_step(d, in, conv_params); n_step > 0) { return winograd_conv_2D_gpu( @@ -1558,7 +1568,8 @@ void conv_1D_gpu( const int O_per_group = O / groups; // Direct to implicit gemm conv - if (is_idil_one && (C_per_group <= 4 || C_per_group % 16 == 0) && + if (is_idil_one && + (C_per_group <= 4 || C_per_group == 8 || C_per_group % 16 == 0) && (O_per_group <= 16 || O_per_group % 16 == 0)) { MLXConvParams<2> conv_params{ /* const int N = */ static_cast(in.shape(0)), diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp index 3e411445f..8e7d10ae5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp @@ -207,7 +207,7 @@ MTL::Library* load_default_library(MTL::Device* device) { return lib; } - NS::Error* error[5]; + NS::Error* error[6]; MTL::Library* lib; // First try the colocated mlx.metallib std::tie(lib, error[0]) = load_colocated_library(device, "mlx"); @@ -220,6 +220,22 @@ MTL::Library* load_default_library(MTL::Device* device) { return lib; } +#ifdef SWIFTPM_BUNDLE + // `swift test` with the default Swift Build build system nests resource + // bundles inside the .xctest and runs tests via the toolchain's + // swiftpm-testing-helper, so the running binary is the toolchain helper: + // neither Bundle.main nor Bundle.allBundles names the test bundle, and the + // SwiftPM bundle search below cannot reach it. In that layout the metallib + // sits at a fixed offset from the binary. + std::tie(lib, error[5]) = load_colocated_library( + device, + std::string("../Resources/") + SWIFTPM_BUNDLE + + ".bundle/Contents/Resources/default"); + if (lib) { + return lib; + } +#endif + // Then try default.metallib in a SwiftPM bundle if we have one std::tie(lib, error[2]) = load_swiftpm_library(device, "default"); if (lib) { @@ -238,7 +254,7 @@ MTL::Library* load_default_library(MTL::Device* device) { if (!lib) { std::ostringstream msg; msg << "Failed to load the default metallib. "; - for (int i = 0; i < 5; i++) { + for (int i = 0; i < 6; i++) { if (error[i] != nullptr) { msg << error[i]->localizedDescription()->utf8String() << " "; } @@ -425,7 +441,7 @@ void CommandEncoder::barrier() { get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); } -void CommandEncoder::end_encoding() { +void CommandEncoder::end_encoding(bool wait_in_flight) { // Each command encoder has a unique fence. We also store a map of // all previous outputs of command encoders to their corresponding fence. // - The command encoder records its inputs and outputs. @@ -462,6 +478,16 @@ void CommandEncoder::end_encoding() { } } } + // A spin wait must start after all earlier encoders, as it might block + // earlier operations to be finished. + // See https://github.com/ml-explore/mlx/pull/4689 + if (wait_in_flight) { + for (auto& [_, fence] : prev_ce_outputs_) { + if (waiting_on.insert(fence).second) { + encoder_->waitForFence(fence.get()); + } + } + } for (auto& out : all_outputs_) { prev_ce_outputs_[out] = fence_; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device.h b/Source/Cmlx/mlx/mlx/backend/metal/device.h index 2d2228335..305a35a77 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/device.h @@ -62,8 +62,13 @@ class MLX_API CommandEncoder { template >> void set_vector_bytes(const Vec& vec, size_t nelems, int idx) { - get_command_encoder()->setBytes( - vec.data(), nelems * sizeof(typename Vec::value_type), idx); + using T = typename Vec::value_type; + if (nelems > 0) { + get_command_encoder()->setBytes(vec.data(), nelems * sizeof(T), idx); + } else { + T val{}; + get_command_encoder()->setBytes(&val, sizeof(T), idx); + } } template >> void set_vector_bytes(const Vec& vec, int idx) { @@ -89,7 +94,7 @@ class MLX_API CommandEncoder { } void barrier(); - void end_encoding(); + void end_encoding(bool wait_in_flight = false); void wait_event(Event event, uint64_t value); void signal_event(Event event, uint64_t value); bool needs_commit() const; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp index 6594f5507..819e2bb1b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp @@ -1,5 +1,8 @@ // Copyright © 2024 Apple Inc. #include "mlx/fence.h" + +#include + #include "mlx/backend/metal/device.h" #include "mlx/scheduler.h" #include "mlx/utils.h" @@ -18,25 +21,22 @@ struct FenceImpl { if (!use_fast) { event = std::make_unique(stream); } else { - auto buf = allocator::malloc(sizeof(uint32_t)).ptr(); - fence = static_cast(buf); + fence.emplace(allocator::malloc(sizeof(uint32_t))); cpu_value()[0] = 0; } } - ~FenceImpl() { - if (use_fast) { - allocator::free(allocator::Buffer{static_cast(fence)}); - } - } bool use_fast{false}; uint32_t count{0}; - void* fence; + std::optional fence; std::unique_ptr event; + MTL::Buffer* buffer() { + return static_cast(fence->buffer().ptr()); + } + std::atomic_uint* cpu_value() { - return static_cast( - static_cast(fence)->contents()); + return static_cast(buffer()->contents()); } }; @@ -66,6 +66,8 @@ void Fence::wait(Stream stream, const array& x, uint32_t value) { auto& d = metal::device(stream.device); auto& compute_encoder = metal::get_command_encoder(stream); + compute_encoder.end_encoding(); + // Register outputs to ensure that no kernels which depends on the // output starts before this one is done compute_encoder.register_output_array(x); @@ -74,10 +76,10 @@ void Fence::wait(Stream stream, const array& x, uint32_t value) { MTL::Size kernel_dims = MTL::Size(1, 1, 1); compute_encoder.set_compute_pipeline_state(kernel); - auto buf = static_cast(f.fence); - compute_encoder.set_buffer(buf, 0); + compute_encoder.set_buffer(f.buffer(), 0); compute_encoder.set_bytes(value, 1); compute_encoder.dispatch_threads(kernel_dims, kernel_dims); + compute_encoder.end_encoding(true); compute_encoder.get_command_buffer()->addCompletedHandler( [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); @@ -125,8 +127,7 @@ uint32_t Fence::update(Stream stream, const array& x, bool cross_device) { MTL::Size kernel_dims = MTL::Size(1, 1, 1); compute_encoder.set_compute_pipeline_state(kernel); - auto buf = static_cast(f.fence); - compute_encoder.set_buffer(buf, 0); + compute_encoder.set_buffer(f.buffer(), 0); compute_encoder.set_bytes(f.count, 1); compute_encoder.dispatch_threads(kernel_dims, kernel_dims); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp b/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp index 3c9195f8e..c3c214ac8 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp @@ -542,7 +542,11 @@ void fft_op( size_t n = out.dtype() == float32 ? out.shape(axis) : in.shape(axis); if (n == 1) { - out.copy_shared_buffer(in); + if (in.dtype() == out.dtype()) { + out.copy_shared_buffer(in); + } else { + copy_gpu(in, out, CopyType::General, s); + } return; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp b/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp index 45f41677c..2a3c3b1f3 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2026 Apple Inc. #include #include "mlx/backend/common/compiled.h" @@ -12,38 +12,41 @@ namespace mlx::core::fast { -bool GatedDeltaUpdate::use_fallback( - const int Hk, - const int Dk, - const int Hv, - const int Dv, - const bool has_mask, - Stream s) { - if (s.device == Device::cpu) { - return true; - } +namespace { - if (has_mask) { - return true; - } +inline int gated_delta_chunk_size(int T) { + const bool nax = metal::is_nax_available(); + const int C = env::get_var("GATED_DELTA_CHUNK", T > 8 ? (nax ? 16 : 8) : 1); + return (C == 16 && !nax) ? 8 : C; +} - if (Dk != 128 || Dv != 128) { - return true; - } +inline int gated_delta_chunk_size_vjp() { + const bool nax = metal::is_nax_available(); + return nax ? env::get_var("GATED_DELTA_CHUNK_VJP", 16) : 1; +} - const bool supported_heads = (Hk == 24 && Hv == 24) || - (Hk == 32 && Hv == 32) || (Hk == 16 && Hv == 32) || - (Hk == 16 && Hv == 48) || (Hk == 16 && Hv == 16) || - (Hk == 16 && Hv == 64); - if (!supported_heads) { - return true; +inline int gated_delta_ckpt(int chunk) { + int c = env::get_var("GATED_DELTA_CKPT", 16); + if (c != 1 && c != 4 && c != 8 && c != 16) { + c = 16; } + return (chunk == 16) ? c : 1; +} - return false; +inline int gated_delta_n_ckpt(int n_chunks, int ckpt) { + return (n_chunks + ckpt - 1) / ckpt; } -inline array -ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { +bool supported_gated_delta_shape(int Hk, int Dk, int Hv, int Dv) { + if (Dk != 128 || Dv != 128) { + return false; + } + return (Hk == 24 && Hv == 24) || (Hk == 32 && Hv == 32) || + (Hk == 16 && Hv == 32) || (Hk == 16 && Hv == 48) || + (Hk == 16 && Hv == 16) || (Hk == 16 && Hv == 64); +} + +array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { if (!x.flags().row_contiguous) { array x_copy = contiguous_copy_gpu(x, s); metal::get_command_encoder(s).add_temporary(x_copy); @@ -53,6 +56,31 @@ ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { } } +array scratch_alloc(Shape shape, Dtype dt, const Stream& s) { + array a(std::move(shape), dt, nullptr, {}); + a.set_data(allocator::malloc(a.nbytes())); + metal::get_command_encoder(s).add_temporary(a); + return a; +} + +} // namespace + +bool GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + if (s.device == Device::cpu) { + return true; + } + if (has_mask) { + return true; + } + return !supported_gated_delta_shape(Hk, Dk, Hv, Dv); +} + void GatedDeltaUpdate::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -76,92 +104,417 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); - int C = 1; - int threshold = env::get_var("GATED_DELTA_THRESH", 16); - if (T >= threshold) - C = metal::is_nax_available() ? 16 : 8; - - C = env::get_var("GATED_DELTA_CHUNK", C); - - if (!metal::is_nax_available()) - C = std::min(C, 8); // override in case nax is not available. - - std::string suffix; - concatenate( - suffix, - get_type_string(q.dtype()), - "_", - std::to_string(Dk), - "_", - std::to_string(Dv), - "_", - std::to_string(Hk), - "_", - std::to_string(Hv)); + int C = gated_delta_chunk_size(T); + const int ckpt = gated_delta_ckpt(C); + + std::string suffix = get_type_string(q.dtype()) + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" + + std::to_string(Hv); auto& compute_encoder = metal::get_command_encoder(s); out.set_data(allocator::malloc(out.nbytes())); hf.set_data(allocator::malloc(hf.nbytes())); + array zero(0, out.dtype()); + compute_encoder.add_temporary(zero); + fill_gpu(zero, out, s); + + // The forward is inference-only here: the vjp runs its own save pass. The + // kernel declares buffers 9, 10 and 11 unconditionally, so bind a placeholder + // and turn the stores off with the function constant. + bool save_state = false; + metal::MTLFCList func_consts = { + {&save_state, MTL::DataType::DataTypeBool, 200}, + }; + + array dummy = scratch_alloc({1}, float32, s); + switch (C) { case 16: { - std::string kernel_name = "gated_delta_fused_nax_"; - std::string base_name; - concatenate(base_name, kernel_name, suffix, "_16"); - std::string hash_name = base_name; - metal::MTLFCList func_consts = {}; + std::string base_name = "gated_delta_fused_nax_" + suffix + "_" + + std::to_string(C) + "_" + std::to_string(ckpt); + + auto delta_kernel = get_gated_delta_nax_kernel( + d, base_name, base_name, func_consts, q, Dk, Dv, Hk, Hv, C, ckpt); - auto delta_kernel = - get_gated_delta_nax_kernel(d, base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); // initial state in + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + compute_encoder.set_output_array(dummy, 9); // state_cache + compute_encoder.set_output_array(dummy, 10); // chunk_mats + compute_encoder.set_output_array(dummy, 11); // chunk_delta + + auto grid = MTL::Size(32, Dv / 16, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); break; } case 8: { - std::string kernel_name = "gated_delta_fused_chunk_"; - std::string base_name; - concatenate(base_name, kernel_name, suffix, "_8"); - std::string hash_name = base_name; - metal::MTLFCList func_consts = {}; + std::string base_name = + "gated_delta_fused_chunk_" + suffix + "_" + std::to_string(C); auto delta_kernel = - get_gated_delta_kernel(d, base_name, hash_name, func_consts); + get_gated_delta_kernel(d, base_name, base_name, func_consts); + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); // initial state in + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + + auto grid = MTL::Size(32, Dv / 8, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); break; } - case 0: - C = 1; // avoiding a div by zero on the grid. - case 1: { - std::string kernel_name = "seq_gated_delta_"; - std::string base_name; - concatenate(base_name, kernel_name, suffix); - std::string hash_name = base_name; - metal::MTLFCList func_consts = {}; + case 1: + case 0: { + // Ckpt is a template parameter now, so even the inference path has to + // name one; it just never stores. + std::string base_name = + "seq_gated_delta_" + suffix + "_" + std::to_string(ckpt); auto delta_kernel = - get_gated_delta_kernel(d, base_name, hash_name, func_consts); + get_gated_delta_kernel(d, base_name, base_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); + + // Order must match the kernel signature: state_in is slot 3, g is 4. + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); + compute_encoder.set_bytes(T, 8); + // Slot 9 is gated off by the function constant, so it stays unbound. + + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + break; + } + default: { + throw std::runtime_error( + "NYI: Only sequential and chunk size 8,16 are supported"); + } + } +} + +///////////////////////// +// BACKWARD PASS STUFF // +///////////////////////// +bool GatedDeltaUpdateVJP::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s) { + if (s.device == Device::cpu) { + return true; + } + if (env::get_var("GATED_DELTA_VJP_FALLBACK", 0) != 0) { + return true; + } + return !supported_gated_delta_shape(Hk, Dk, Hv, Dv); +} + +void GatedDeltaUpdateVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + // Inputs are q, k, v, g, b, h0, cot_o, cot_h + // Outputs are dq, dk, dv, dg, db, dh0 + auto& s = stream(); + auto& d = metal::device(s.device); + + auto q = ensure_row_contiguous(inputs[0], d, s); + auto k = ensure_row_contiguous(inputs[1], d, s); + auto v = ensure_row_contiguous(inputs[2], d, s); + auto g = ensure_row_contiguous(inputs[3], d, s); + auto beta = ensure_row_contiguous(inputs[4], d, s); + auto h0 = ensure_row_contiguous(inputs[5], d, s); + auto cot_o = ensure_row_contiguous(inputs[6], d, s); + auto cot_h = ensure_row_contiguous(inputs[7], d, s); + + int B = q.shape(0); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + // 16 = chunked NAX, anything else = sequential. + int C = gated_delta_chunk_size_vjp(); + const bool chunked = (C == 16); + + const int n_chunks = (T + 15) / 16; + const int ckpt = gated_delta_ckpt(C); + const int n_states = chunked ? gated_delta_n_ckpt(n_chunks, ckpt) : T; + + auto& dq = outputs[0]; + auto& dk = outputs[1]; + auto& dv = outputs[2]; + auto& dg = outputs[3]; + auto& db = outputs[4]; + auto& dh = outputs[5]; + + auto& compute_encoder = metal::get_command_encoder(s); + + dq.set_data(allocator::malloc(dq.nbytes())); + dk.set_data(allocator::malloc(dk.nbytes())); + dv.set_data(allocator::malloc(dv.nbytes())); + dg.set_data(allocator::malloc(dg.nbytes())); + db.set_data(allocator::malloc(db.nbytes())); + dh.set_data(allocator::malloc(dh.nbytes())); + + // The kernels accumulate their gradients as fp32, so we may need a cast. + const bool stage_fp32 = (q.dtype() != float32); + + auto dq_acc = stage_fp32 ? scratch_alloc(dq.shape(), float32, s) : dq; + auto dk_acc = stage_fp32 ? scratch_alloc(dk.shape(), float32, s) : dk; + auto dv_acc = stage_fp32 ? scratch_alloc(dv.shape(), float32, s) : dv; + auto dg_acc = stage_fp32 ? scratch_alloc(dg.shape(), float32, s) : dg; + auto db_acc = stage_fp32 ? scratch_alloc(db.shape(), float32, s) : db; + + array zero(0, float32); + array zero_h(0, dh.dtype()); + compute_encoder.add_temporary(zero); + compute_encoder.add_temporary(zero_h); + fill_gpu(zero, dq_acc, s); + fill_gpu(zero, dk_acc, s); + fill_gpu(zero, dv_acc, s); + fill_gpu(zero, dg_acc, s); + fill_gpu(zero, db_acc, s); + fill_gpu(zero_h, dh, s); + + array state_cache = scratch_alloc({B, Hv, n_states, Dv, Dk}, float32, s); + array seg_states = scratch_alloc( + chunked ? Shape({B, Hv, ckpt, Dv, Dk}) : Shape({1}), float32, s); + array chunk_mats = scratch_alloc( + chunked ? Shape({B, Hv, n_chunks, 3, 16, 16}) : Shape({1}), float32, s); + array chunk_delta = scratch_alloc( + chunked ? Shape({B, Hv, n_chunks, 16, Dv}) : Shape({1}), float32, s); + + if (chunked) { + fill_gpu(zero, state_cache, s); + fill_gpu(zero, chunk_mats, s); + fill_gpu(zero, chunk_delta, s); + } + + array y_scratch = scratch_alloc({B, T, Hv, Dv}, q.dtype(), s); + array hf_scratch = scratch_alloc({B, Hv, Dv, Dk}, float32, s); + + std::string suffix = get_type_string(q.dtype()) + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" + + std::to_string(Hv); + + bool save_state = true; + metal::MTLFCList save_consts = { + {&save_state, MTL::DataType::DataTypeBool, 200}, + }; + metal::MTLFCList no_consts = {}; + + switch (C) { + case 16: { + const std::string ckpt_suffix = + "_" + std::to_string(C) + "_" + std::to_string(ckpt); + + // Forward save pass. + { + std::string base_name = "gated_delta_fused_nax_" + suffix + ckpt_suffix; + + auto delta_kernel = get_gated_delta_nax_kernel( + d, + base_name, + base_name + "_save", + save_consts, + q, + Dk, + Dv, + Hk, + Hv, + C, + ckpt); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(y_scratch, 6); + compute_encoder.set_output_array(hf_scratch, 7); + compute_encoder.set_bytes(T, 8); + compute_encoder.set_output_array(state_cache, 9); + compute_encoder.set_output_array(chunk_mats, 10); + compute_encoder.set_output_array(chunk_delta, 11); + + auto grid = MTL::Size(32, Dv / 16, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + } + + { + std::string base_name = + "gated_delta_vjp_fused_nax_" + suffix + ckpt_suffix; + + auto delta_kernel = get_gated_delta_vjp_nax_kernel( + d, + base_name, + base_name, + no_consts, + q, + Dk, + Dv, + Hk, + Hv, + C, + ckpt, + false); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(cot_o, 5); + compute_encoder.set_input_array(cot_h, 6); + compute_encoder.set_input_array(state_cache, 7); + compute_encoder.set_bytes(T, 8); + compute_encoder.set_output_array(dq_acc, 9); + compute_encoder.set_output_array(dk_acc, 10); + compute_encoder.set_output_array(dv_acc, 11); + compute_encoder.set_output_array(dg_acc, 12); + compute_encoder.set_output_array(db_acc, 13); + compute_encoder.set_output_array(dh, 14); + compute_encoder.set_input_array(chunk_mats, 15); + compute_encoder.set_input_array(chunk_delta, 16); + compute_encoder.set_output_array(seg_states, 17); + + auto grid = MTL::Size(32, Dv / 16, B * Hv); + auto threads = MTL::Size(32, Dv / 16, 1); + compute_encoder.dispatch_threads(grid, threads); + } + + { + int n_total = B * Hv; + + std::string base_name = "gated_delta_dgamma_to_dg_" + + get_type_string(q.dtype()) + "_" + std::to_string(C); + + auto dgamma_kernel = get_gated_delta_vjp_nax_kernel( + d, + base_name, + base_name, + no_consts, + q, + Dk, + Dv, + Hk, + Hv, + C, + ckpt, + true); + + compute_encoder.set_compute_pipeline_state(dgamma_kernel); + compute_encoder.set_input_array(g, 0); + compute_encoder.set_output_array(dg_acc, 1); + compute_encoder.set_bytes(T, 2); + compute_encoder.set_bytes(Hv, 3); + compute_encoder.set_bytes(n_total, 4); + + auto grid = MTL::Size(n_total, n_chunks, 1); + auto threads = MTL::Size(std::min(n_total, 32), 1, 1); + compute_encoder.dispatch_threads(grid, threads); + } + break; + } + case 1: + case 0: { + const std::string ckpt_suffix = "_" + std::to_string(ckpt); + + { + std::string base_name = "seq_gated_delta_" + suffix + ckpt_suffix; + + auto delta_kernel = get_gated_delta_kernel( + d, base_name, base_name + "_save", save_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(y_scratch, 6); + compute_encoder.set_output_array(hf_scratch, 7); + compute_encoder.set_bytes(T, 8); + compute_encoder.set_output_array(state_cache, 9); + + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + } + + { + std::string base_name = "seq_gated_delta_vjp_" + suffix + ckpt_suffix; + + auto delta_kernel = + get_gated_delta_vjp_kernel(d, base_name, base_name, no_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(cot_o, 5); + compute_encoder.set_input_array(cot_h, 6); + compute_encoder.set_input_array(state_cache, 7); + compute_encoder.set_bytes(T, 8); + compute_encoder.set_output_array(dq_acc, 9); + compute_encoder.set_output_array(dk_acc, 10); + compute_encoder.set_output_array(dv_acc, 11); + compute_encoder.set_output_array(dg_acc, 12); + compute_encoder.set_output_array(db_acc, 13); + compute_encoder.set_output_array(dh, 14); + + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + } break; } default: { throw std::runtime_error( - "NYI: Only sequential (C=1) and chunk size 8,16 are supported"); + "NYI: Only sequential and chunk size 16 are supported for vjp"); } } - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_input_array(h0, 3); // initial state in - compute_encoder.set_input_array(g, 4); - compute_encoder.set_input_array(beta, 5); - compute_encoder.set_output_array(out, 6); - compute_encoder.set_output_array(hf, 7); // final state out - compute_encoder.set_bytes(T, 8); - - auto grid = MTL::Size(32, Dv / C, B * Hv); - auto threads = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid, threads); + + if (stage_fp32) { + copy_gpu(dq_acc, dq, CopyType::General, s); + copy_gpu(dk_acc, dk, CopyType::General, s); + copy_gpu(dv_acc, dv, CopyType::General, s); + copy_gpu(dg_acc, dg, CopyType::General, s); + copy_gpu(db_acc, db, CopyType::General, s); + } } } // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp b/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp index 117e402f7..4ee309773 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp @@ -212,9 +212,6 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { compute_encoder.set_vector_bytes(axes_, 6); // Set index info - // - // We don't need to check for empty idx_shapes because gather has a - // idx_ndim == 0 specialization compute_encoder.set_vector_bytes(idx_shapes, 7); compute_encoder.set_vector_bytes(idx_strides, 8); compute_encoder.set_vector_bytes(idx_contigs, 9); @@ -379,42 +376,19 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { idx_contigs.push_back(inputs[i + 1].flags().row_contiguous); } - if (upd_ndim == 0) { - // Need placeholders so Metal doesn't complain - int shape_ = 0; - int64_t stride_ = 0; - compute_encoder.set_bytes(shape_, 3); - compute_encoder.set_bytes(stride_, 4); - } else { - compute_encoder.set_vector_bytes(upd.shape(), 3); - compute_encoder.set_vector_bytes(upd.strides(), 4); - } + compute_encoder.set_vector_bytes(upd.shape(), 3); + compute_encoder.set_vector_bytes(upd.strides(), 4); compute_encoder.set_bytes(upd_ndim, 5); compute_encoder.set_bytes(upd_size, 6); // Set output info size_t out_ndim = out.ndim(); - if (out_ndim == 0) { - // Need placeholders so Metal doesn't complain - int shape_ = 0; - int64_t stride_ = 0; - compute_encoder.set_bytes(shape_, 7); - compute_encoder.set_bytes(stride_, 8); - } else { - compute_encoder.set_vector_bytes(out.shape(), 7); - compute_encoder.set_vector_bytes(out.strides(), 8); - } + compute_encoder.set_vector_bytes(out.shape(), 7); + compute_encoder.set_vector_bytes(out.strides(), 8); compute_encoder.set_bytes(out_ndim, 9); - compute_encoder.set_vector_bytes(axes_, 10); // Set index info - if (idx_ndim == 0) { - // Add a 0 in idx_shapes and strides to avoid the missing buffer binding - // error in the metal API. - idx_shapes.push_back(0); - idx_strides.push_back(0); - idx_contigs.push_back(false); - } + compute_encoder.set_vector_bytes(axes_, 10); compute_encoder.set_vector_bytes(idx_shapes, 11); compute_encoder.set_vector_bytes(idx_strides, 12); compute_encoder.set_vector_bytes(idx_contigs, 13); @@ -668,7 +642,9 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { return; } - array mask_flat = flatten_in_eval(mask, 1, -1, s); + array mask_flat = (mask.ndim() == 1) + ? reshape_in_eval(mask, {mask.shape(0), 1}, s) + : flatten_in_eval(mask, 1, -1, s); if (mask_flat.data() != mask.data()) { compute_encoder.add_temporary(mask_flat); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h index b4f72e660..8ad0a9ca7 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h @@ -54,6 +54,7 @@ const char* steel_gemm_fused_nax(); const char* steel_gemm_gather_nax(); const char* steel_gemm_splitk_nax(); const char* steel_gemm_segmented_nax(); +const char* steel_gemm_thin_nax(); const char* quantized_nax(); const char* fp_quantized_nax(); @@ -62,5 +63,6 @@ const char* steel_attention_nax(); const char* gated_delta_update(); const char* gated_delta_update_nax(); +const char* gated_delta_update_nax_vjp(); } // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp b/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp index d2ce55bbf..aad94eba1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp @@ -30,6 +30,9 @@ const char* steel_gemm_splitk_nax() { const char* steel_gemm_segmented_nax() { return ""; } +const char* steel_gemm_thin_nax() { + return ""; +} const char* quantized_nax() { return ""; } @@ -42,6 +45,9 @@ const char* steel_attention_nax() { const char* gated_delta_update_nax() { return ""; } +const char* gated_delta_update_nax_vjp() { + return ""; +} } // namespace metal #endif // MLX_METAL_NO_NAX @@ -1260,6 +1266,25 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string& template_def) { + auto lib = d.get_library(kernel_name, [&]() { + std::string kernel_source; + concatenate( + kernel_source, + metal::utils(), + metal::gemm_nax(), + metal::steel_gemm_thin_nax(), + template_def); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, @@ -1413,6 +1438,22 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_sdpa_vjp_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_sdpa_vjp_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_gated_delta_kernel( metal::Device& d, const std::string& kernel_name, @@ -1421,15 +1462,78 @@ MTL::ComputePipelineState* get_gated_delta_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } -MTL::ComputePipelineState* get_gated_delta_nax_kernel( +MTL::ComputePipelineState* get_gated_delta_vjp_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, const metal::MTLFCList& func_consts) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& q, + int dk, + int dv, + int hk, + int hv, + int c, + int ckpt) { const auto& lib_name = kernel_name; auto lib = d.get_library(lib_name, [&]() { std::string kernel_source; concatenate(kernel_source, metal::utils(), metal::gated_delta_update_nax()); + kernel_source += get_template_definition( + lib_name, + "gated_delta_fused_nax", + get_type_string(q.dtype()), + dk, + dv, + hk, + hv, + c, + ckpt); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gated_delta_vjp_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& q, + int dk, + int dv, + int hk, + int hv, + int c, + int ckpt, + bool dgamma) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::string kernel_source; + concatenate( + kernel_source, metal::utils(), metal::gated_delta_update_nax_vjp()); + if (dgamma) { + kernel_source += get_template_definition( + lib_name, "gated_delta_dgamma_to_dg", get_type_string(q.dtype()), c); + } else { + kernel_source += get_template_definition( + lib_name, + "gated_delta_vjp_fused_nax", + get_type_string(q.dtype()), + dk, + dv, + hk, + hv, + c, + ckpt); + } return kernel_source; }); return d.get_kernel(kernel_name, lib, hash_name, func_consts); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h index 82c3d3d16..ded6c834a 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h @@ -386,6 +386,13 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( int wm, int wn); +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string& template_def); + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, @@ -444,12 +451,39 @@ MTL::ComputePipelineState* get_gated_delta_kernel( const std::string& hash_name, const metal::MTLFCList& func_consts); -MTL::ComputePipelineState* get_gated_delta_nax_kernel( +MTL::ComputePipelineState* get_gated_delta_vjp_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, const metal::MTLFCList& func_consts); +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& q, + int dk, + int dv, + int hk, + int hv, + int c, + int ckpt); + +MTL::ComputePipelineState* get_gated_delta_vjp_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& q, + int dk, + int dv, + int hk, + int hv, + int c, + int ckpt, + bool dgamma); + // Create a GPU kernel template definition for JIT compilation template std::string get_template_definition( diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt index d5513dd03..355c763a3 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt @@ -55,8 +55,10 @@ build_kernel(layer_norm) build_kernel(random) build_kernel(rms_norm) build_kernel(rope) -build_kernel(scaled_dot_product_attention sdpa_vector.h) +build_kernel(scaled_dot_product_attention sdpa_blocked.h sdpa_vector.h) +build_kernel(scaled_dot_product_attention_vjp) build_kernel(gated_delta_update gated_delta_update.h) +build_kernel(gated_delta_update_vjp gated_delta_update_vjp.h) if(MLX_METAL_VERSION GREATER_EQUAL 320) build_kernel(fence) endif() @@ -113,7 +115,8 @@ set(STEEL_NAX_HEADERS steel/gemm/kernels/steel_gemm_fused_nax.h steel/gemm/kernels/steel_gemm_gather_nax.h steel/gemm/kernels/steel_gemm_splitk_nax.h - steel/gemm/kernels/steel_gemm_segmented_nax.h) + steel/gemm/kernels/steel_gemm_segmented_nax.h + steel/gemm/kernels/steel_gemm_thin_nax.h) set(STEEL_NAX_ATTN_HEADERS steel/defines.h @@ -170,10 +173,13 @@ if(NOT MLX_METAL_JIT) build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_segmented_nax ${STEEL_NAX_HEADERS}) + build_kernel(steel/gemm/kernels/steel_gemm_thin_nax ${STEEL_NAX_HEADERS}) build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS}) build_kernel(gated_delta_update_nax gated_delta_update_nax.h - ${STEEL_NAX_HEADERS}) + gated_delta_nax_ops.h ${STEEL_NAX_HEADERS}) + build_kernel(gated_delta_update_nax_vjp gated_delta_update_nax_vjp.h + gated_delta_nax_ops.h ${STEEL_NAX_HEADERS}) build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h ${STEEL_NAX_HEADERS}) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary.metal index 17ed13c57..475fd45ad 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary.metal @@ -73,6 +73,7 @@ instantiate_binary_types(Add) instantiate_binary_types(Divide) +instantiate_binary_types(FloorDivide) instantiate_binary_types_bool(Equal) instantiate_binary_types_bool(Greater) instantiate_binary_types_bool(GreaterEqual) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h index fe78e0050..71a5fa439 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h @@ -264,10 +264,20 @@ constexpr complex_t operator*(complex_t a, complex_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}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -275,11 +285,7 @@ template < 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 = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal index 6f3ce0cd4..a16ae8a27 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal @@ -418,8 +418,8 @@ winograd_conv_2d_weight_transform( // wt_out is stored transposed (A x A x C x O) short ohw_0 = sm * 8 + sn; short ohw_1 = sm * 8 + sn + 1; - device T* wt_out_0 = wt_out + ohw_0 * C * O + ko; - device T* wt_out_1 = wt_out + ohw_1 * C * O + ko; + device T* wt_out_0 = wt_out + size_t(ohw_0) * C * O + ko; + device T* wt_out_1 = wt_out + size_t(ohw_1) * C * O + ko; // Prepare shared memory threadgroup T Ws[BO][R][R][BC]; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h index b137a6d47..bd001c335 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h @@ -334,6 +334,7 @@ template < typename T, int group_size, int bits, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( @@ -367,6 +368,15 @@ METAL_FUNC void fp_qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // Reuse the last valid row for reads beyond the output tile. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; @@ -376,8 +386,12 @@ METAL_FUNC void fp_qmv_fast_impl( load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device auto* sl = scales + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device auto* sl = scales + src * in_vec_size_g; U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); @@ -395,7 +409,7 @@ METAL_FUNC void fp_qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { if constexpr (has_global_scale) { y[row] = static_cast(result[row] * inv_scale_enc); } else { @@ -796,6 +810,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const bool has_global_scale = false, const int BM = 32, const int BK = 32, @@ -805,7 +820,7 @@ METAL_FUNC void fp_qmm_t_impl( const device uint8_t* scales, const device float* global_scale, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -830,7 +845,7 @@ METAL_FUNC void fp_qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1243,6 +1258,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( @@ -1284,6 +1300,7 @@ template < T, group_size, bits, + partial_rows, has_global_scale, results_per_simdgroup>( w, @@ -1303,6 +1320,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv( @@ -1511,14 +1529,14 @@ template template < typename T, - const int group_size, - const int bits, - const bool aligned_N, - const bool batched, - const bool has_global_scale = false, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool aligned_N, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, @@ -1562,7 +1580,16 @@ template < s_strides, tid); } - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -1684,7 +1711,7 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( + fp_qmv_fast_impl( w, scales, global_scale, @@ -1881,7 +1908,16 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( w, scales, global_scale, @@ -1911,7 +1947,7 @@ template < const device uint32_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], const device T* x [[buffer(2)]], - device T* y [[buffer(3)]], + device float* y [[buffer(3)]], const constant int& K [[buffer(4)]], const constant int& N [[buffer(5)]], const constant int& M [[buffer(6)]], @@ -1936,7 +1972,7 @@ template < scales += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - fp_qmm_t_impl( + fp_qmm_t_impl( (const device uint32_t*)wl, scales, nullptr, diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal index fe508a6aa..ffd8b5b9c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal @@ -21,7 +21,7 @@ bits, \ true) -#define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \ +#define instantiate_quantized_batched(mode, name, type, group_size, bits, batched) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \ fp_ ## name, \ @@ -38,6 +38,18 @@ batched, \ true) +#define instantiate_quantized_qmv(mode, name, type, group_size, bits, batched, partial_rows, has_global_scale, results) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_pr_" #partial_rows "_hgs_" #has_global_scale "_r_" #results, \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + batched, \ + partial_rows, \ + has_global_scale, \ + results) + #define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \ @@ -79,30 +91,6 @@ batched, \ true) -#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, \ @@ -161,13 +149,20 @@ true) #define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \ - instantiate_quantized_batched(mode, name, type, 1, group_size, bits) \ - instantiate_quantized_batched(mode, name, type, 0, group_size, bits) + instantiate_quantized_batched(mode, name, type, group_size, bits, 1) \ + instantiate_quantized_batched(mode, name, type, group_size, bits, 0) + +#define instantiate_quantized_qmv_wrap(name, type, mode, group_size, bits, partial_rows, results) \ + instantiate_quantized_qmv(mode, name, type, group_size, bits, 1, partial_rows, 1, results) \ + instantiate_quantized_qmv(mode, name, type, group_size, bits, 1, partial_rows, 0, results) \ + instantiate_quantized_qmv(mode, name, type, group_size, bits, 0, partial_rows, 1, results) \ + instantiate_quantized_qmv(mode, name, type, group_size, bits, 0, partial_rows, 0, results) -#define instantiate_quantized_all_batched(type, mode, group_size, bits) \ - instantiate_quantized_batched_wrap(qmv_fast, type, mode, group_size, bits) \ - instantiate_quantized_batched_wrap(qmv, type, mode, group_size, bits) \ - instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \ +#define instantiate_quantized_all_batched(type, mode, group_size, bits) \ + instantiate_quantized_qmv_wrap(qmv_fast, type, mode, group_size, bits, 1, 4) \ + instantiate_quantized_qmv_wrap(qmv_fast, type, mode, group_size, bits, 0, 4) \ + instantiate_quantized_qmv_wrap(qmv, type, mode, group_size, bits, 0, 4) \ + instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \ instantiate_quantized_batched_wrap(qmm_n, type, mode, group_size, bits) #define instantiate_quantized_all_single(type, mode, group_size, bits) \ @@ -248,14 +243,15 @@ instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_modes(type, nvfp4, 16, 4) \ - instantiate_quantized_modes(type, mxfp8, 32, 8) \ - instantiate_quantized_modes(type, mxfp4, 32, 4) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ - instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ - instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ - instantiate_quantized_qmv_fast_r2(nvfp4, type, 16, 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_wrap(qmv_fast, type, nvfp4, 16, 4, 1, 2) \ + instantiate_quantized_qmv_wrap(qmv_fast, type, nvfp4, 16, 4, 0, 2) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_nax_ops.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_nax_ops.h new file mode 100644 index 000000000..e8e4069f6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_nax_ops.h @@ -0,0 +1,588 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include + +#include +#include + +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } \ No newline at end of file diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h index 9553356c7..78171e4b4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h @@ -3,6 +3,8 @@ #include #include "mlx/backend/metal/kernels/utils.h" +constant bool save_state [[function_constant(200)]]; + #define AT(TILE, IDX) TILE.thread_elements()[IDX] #define SUB(TILE0, TILE1, TILE2) \ { \ @@ -263,25 +265,24 @@ template } /* - auto grid = MTL::Size(32, Dv, B * Hv); + auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); */ -template +template [[kernel]] void gated_delta_seq( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], const device float* state_in [[buffer(3)]], - const device InT* g [[buffer(4)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* g [[buffer(4)]], // [B, T, Hv] const device InT* beta [[buffer(5)]], // [B, T, Hv] - // [B, Hv, Dv, Dk] device InT* y [[buffer(6)]], // [B, T, Hv, Dv] device float* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9), function_constant(save_state)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - // kernel implementation auto n = thread_position_in_grid.z; auto b_idx = n / Hv; auto hv_idx = n % Hv; @@ -313,7 +314,23 @@ template auto g_ = g + b_idx * T * Hv; auto beta_ = beta + b_idx * T * Hv; + // Checkpoint every Ckpt-th step rather than every step: at T = 4096 with + // B = 4, Hv = 32 a dense cache is 34 GB, which the backward's replay makes + // unnecessary. + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + auto c_state_base = state_cache + n * n_ckpt * Dv * Dk + dv_idx * Dk; + for (int t = 0; t < T; ++t) { + // Stored before the decay, so slot t/Ckpt holds the state at *entry* to + // step t. The backward replays forward from here, so it must be the entry + // state and not the exit state. + if (save_state && (t % Ckpt) == 0) { + auto c_state = c_state_base + (t / Ckpt) * Dv * Dk; + for (int i = 0; i < n_per_t; ++i) { + c_state[n_per_t * dk_idx + i] = state[i]; + } + } + float kv_mem = 0.0f; for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal index 21238df43..f0797e80b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal @@ -3,15 +3,17 @@ using namespace metal; -#define instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ - instantiate_kernel( \ - "seq_gated_delta_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv, \ - gated_delta_seq, \ - in_type, \ - dk, \ - dv, \ - hk, \ - hv) +#define instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) #define instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ @@ -25,18 +27,24 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ - instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ - instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) +// Ckpt is a template parameter of the sequential kernel, so the inference path +// has to name one even though it never stores. These must cover every value +// gated_delta_ckpt() can return, since the host puts it in the kernel name. +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 16) \ + instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) -#define instantiate_gated_delta(in_type) \ - instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ - instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ - instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ - instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) \ - instantiate_gated_delta_dims(in_type, 128, 128, 16, 16) \ - instantiate_gated_delta_dims(in_type, 128, 128, 16, 64) +#define instantiate_gated_delta(in_type) \ + instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 64) instantiate_gated_delta(float); instantiate_gated_delta(bfloat16_t); -instantiate_gated_delta(float16_t); \ No newline at end of file +instantiate_gated_delta(float16_t); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h index 4e65fe1e5..09986fe80 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h @@ -1,360 +1,19 @@ +// Copyright © 2026 Apple Inc. #pragma once -#include - -#include -#include - -#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/gated_delta_nax_ops.h" using namespace metal; using namespace mpp; using namespace mpp::tensor_ops; -// NAX MACROS I can probably do a nice template instead of doing this -// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for -// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns -#define AT_NAX(TILE, IDX) TILE.elems()[IDX] - -#define SUB_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ - } \ - } - -#define ADD_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ - } \ - } - -#define FMA_NAX(TILE0, S, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ - (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ - } \ - } - -#define SCALE_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - AT_NAX(TILE0, _i) *= (S); \ - } \ - } - -#define SCALE_ROW_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= \ - metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ - } \ - } - -#define SCALE_BETA_NAX(TILE0, BETA2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ - } \ - } - -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ - AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ - } \ - } - -#define SCALE_TRI_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ - ? 0.f \ - : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ - } \ - } - -#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ - } \ - } - -namespace mlx { -namespace steel { -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A0, - const thread BaseNAXFrag::dtype_frag_t& A1, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B0, - const thread BaseNAXFrag::dtype_frag_t& B1, - metal::bool_constant) { - // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 16, 32, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A0[i]; - ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; - ct_b[i] = B0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; - ct_c[i] = C[i]; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B, - metal::bool_constant) { - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = B[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - ct_c[i] = C[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - } +/////////////////////////////////////////////////////////////////////////////// +// Function constants +/////////////////////////////////////////////////////////////////////////////// - gemm_op.run(ct_a, ct_b, ct_c); +constant bool save_state_cache [[function_constant(200)]]; - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mman( - thread BaseNAXFrag::dtype_frag_t& Cn0, - thread BaseNAXFrag::dtype_frag_t& Cn1, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& Bn0, - const thread BaseNAXFrag::dtype_frag_t& Bn1, - metal::bool_constant) { - // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - // Create matmul op - mpp::tensor_ops::matmul2d gemm_op; - - // Create matmul operands in registers - using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< - AType, - BType, - CType>()); - using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< - AType, - BType, - CType>()); - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - - // Create matmul output in register - auto ct_c = - gemm_op - .template get_destination_cooperative_tensor(); - - // Load A in to left operand registers - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = Bn0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; - ct_c[i] = Cn0[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; - } - - // Do matmul - gemm_op.run(ct_a, ct_b, ct_c); - - // Copy out results - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - Cn0[i] = ct_c[i]; - Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; - } -} - -} // namespace steel -} // namespace mlx - -#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - A.frag_at(0, (AO) + 1), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -template +template [[kernel]] void gated_delta_fused_nax( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], @@ -365,6 +24,9 @@ template device InT* y [[buffer(6)]], device float* state_out [[buffer(7)]], constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9)]], // [B, Hv, n_ckpt, Dv, Dk] + device float* chunk_mats [[buffer(10)]], // [B, Hv, n_chunks, 3, 16, 16] + device float* chunk_delta [[buffer(11)]], // [B, Hv, n_chunks, C, Dv] uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { @@ -397,6 +59,13 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; + const int n_chunks = (T + C - 1) / C; + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + auto c_state = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + + auto o_chunk_mats = chunk_mats + n * n_chunks * 3 * 256; + auto o_chunk_delta = chunk_delta + n * n_chunks * C * Dv + dv_idx; + threadgroup float gamma_all[C * 4]; threadgroup float* gamma = gamma_all + sg_id * C; @@ -429,6 +98,10 @@ template } mlx::steel::NAXTile TMP_tile; + // Which chunk process_chunk is on, so it knows whether this is a segment + // boundary. + int chunk_idx = 0; + auto process_chunk = [&](const short valid_rows, auto bounded_tag) __attribute__((always_inline)) { constexpr bool B = decltype(bounded_tag)::value; @@ -441,19 +114,23 @@ template } }; + // Checkpoint the state entering this chunk before anything mutates it, on + // segment boundaries only. + if (save_state_cache && (chunk_idx % Ckpt) == 0) { + S_tile.store(c_state, Dk); + c_state += Dv * Dk; + } + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? metal::fast::log( metal::max( - static_cast( - g_[thread_index_in_simdgroup * Hv + hv_idx]), - 1e-6f)) + float(g_[thread_index_in_simdgroup * Hv + hv_idx]), 1e-6)) : 0.0f; auto gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = static_cast(gamma_val); } - simdgroup_barrier(mem_flags::mem_threadgroup); beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; @@ -465,6 +142,12 @@ template MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); } + // KKt_raw is identical for every Dv tile of this chunk; only the first + // tile needs to write it out for the vjp kernel to reuse. + if (save_state_cache && dv_idx == 0) { + KKt_tile.store(o_chunk_mats, 16); + } + KKtK_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); @@ -475,6 +158,12 @@ template SUB_NAX(Tinv_tile, I_tile, TMP_tile); } + // Same as KKt_raw: cache the pre-decay inverse before it is overwritten + // below by the gamma-scaled (TUinv) version. + if (save_state_cache && dv_idx == 0) { + Tinv_tile.store(o_chunk_mats + 256, 16); + } + STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { load_seq(K_tile, k_ + nn * 16, Dk * Hk); @@ -496,20 +185,29 @@ template SUB_NAX(delta_tile, U_tile, WS_tile) - tmp_tile.clear(); QKt_tile.clear(); + tmp_tile.clear(); for (int kk = 0; kk < Dk; kk += 32) { load_seq(Q_tile, q_ + kk, Hk * Dk); load_seq(K_tile, k_ + kk, Hk * Dk); - MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); - SCALE_ROW_NAX(Q_tile, gamma); MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); } - SCALE_TRI_NAX(QKt_tile, gamma) + if (save_state_cache) { + if (dv_idx == 0) { + QKt_tile.store(o_chunk_mats + 512, 16); + } + // Every Dv tile stores its own slice. + delta_tile.store(o_chunk_delta, Dv); + o_chunk_mats += 3 * 256; + o_chunk_delta += C * Dv; + } + // Output path runs unconditionally: the forward serves both inference and + // the training save, so y and state_out are always needed. + SCALE_TRI_NAX(QKt_tile, gamma) out_tile = tmp_tile; MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); @@ -531,6 +229,8 @@ template SCALE2_NAX(K_tile, gamma); MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); } + + chunk_idx++; }; int t = 0; @@ -548,4 +248,4 @@ template } S_tile.store(o_state, Dk); -} \ No newline at end of file +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal index 35790fa49..de84d866c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal @@ -3,20 +3,25 @@ using namespace metal; -#define instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ - "_" #c, \ - gated_delta_fused_nax, \ - in_type, \ - dk, \ - dv, \ - hk, \ - hv, \ - c) +#define instantiate_gated_delta_update_fused_nax( \ + in_type, dk, dv, hk, hv, c, ckpt) \ + instantiate_kernel( \ + "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv \ + "_" #hk "_" #hv "_" #c "_" #ckpt, \ + gated_delta_fused_nax, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c, \ + ckpt) -#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ - instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16) +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16, 1) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16, 4) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16, 8) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16, 16) #define instantiate_gated_delta(in_type) \ instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h new file mode 100644 index 000000000..3c2190f04 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h @@ -0,0 +1,1737 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "mlx/backend/metal/kernels/gated_delta_nax_ops.h" + +#include +#include "mlx/backend/metal/kernels/atomic.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +template +METAL_FUNC void reduce_tile_tg( + thread TileT& acc, + threadgroup float* scratch, + device mlx_atomic* out, + int row_stride, + int col_off, + short valid_rows, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(sg_id * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _s = 1; _s < kNSG; _s++) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(_s * 32 + simd_lane_id) * kE + _i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Simdgroup 0 alone holds the total. + if (sg_id != 0) { + return; + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * row_stride + col_off + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(out, AT_NAX(acc, _i), idx); + } else { + mlx_atomic_store_explicit(out, AT_NAX(acc, _i), idx); + } + } + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + device float* seg_states [[buffer(17)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + const int n_chunks = (T + C - 1) / C; + // Ceiling: the floor form gives 0 when n_chunks < Ckpt, which would leave the + // allocation empty and fault on the first read. + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + const short tail = short(T - (n_chunks - 1) * C); + + // Bases, indexed by chunk rather than walked. With two nested loops the + // incremental decrements are no longer expressible. + auto q_base = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_base = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dq_base = dq + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dk_base = dk + b_idx * T * Hk * Dk + hk_idx * Dk; + + auto v_base = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto dv_base = dv + b_idx * T * Hv * Dv + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto g_base = g + b_idx * T * Hv; + auto beta_base = beta + b_idx * T * Hv; + auto dg_base = dg + b_idx * T * Hv; + auto db_base = db + b_idx * T * Hv; + + auto cm_base = chunk_mats + n * n_chunks * 3 * 256; + auto cd_base = chunk_delta + n * n_chunks * C * Dv + dv_idx; + auto ck_base = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + // Each simdgroup owns a distinct dv slice, so these regions are disjoint and + // the replay needs no cross-simdgroup synchronisation. + auto seg_base = seg_states + (n * Ckpt * Dv + dv_idx) * Dk; + + // Per-chunk pointers, assigned from the bases before each process_chunk call. + auto q_ = q_base; + auto k_ = k_base; + auto dq_ = dq_base; + auto dk_ = dk_base; + auto v_ = v_base; + auto dv_ = dv_base; + auto co_ = co_base; + auto g_ = g_base; + auto beta_ = beta_base; + auto dg_ = dg_base; + auto db_ = db_base; + auto i_chunk_mats = cm_base; + auto i_chunk_delta = cd_base; + auto i_seg = seg_base; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, taken from the replayed segment. + _M16xDk S_tile; + + _M16x32 K_tile, Q_tile; + _M16x16 V_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // From the replayed segment rather than a per-chunk checkpoint. + S_tile.load(i_seg, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + if (fm < valid_rows) { + db_st[fm] = _b0; + } + if (fm1 < valid_rows) { + db_st[fm1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (fm == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (fm < valid_rows) { + dg_st[fm] = _d0; + } + if (fm1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (fm1 < valid_rows) { + dg_st[fm1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk segments in reverse. Each is replayed forward from its checkpoint to + // recover the entry states, then walked backwards. + // Backward over one segment: replay its entry states from the checkpoint, + // then walk it in reverse. HasTail is a compile-time tag so valid_rows is the + // constant C on every chunk of an interior segment, which folds away the + // row guards in the dq/dk stores and the load_rows switch in the replay. + auto process_segment = [&](const int seg, + auto tail_tag) __attribute__((always_inline)) { + constexpr bool HasTail = decltype(tail_tag)::value; + + const int seg_start = seg * Ckpt; + const int rem = n_chunks - seg_start; + const int n_in_seg = (rem < Ckpt) ? rem : Ckpt; + const int n_full = HasTail ? n_in_seg - 1 : n_in_seg; + + // delta is cached per chunk, so the replay is the state recurrence alone: + // S <- exp(gamma_C) S + K_dec^T delta + // No T^-1, no W, no U -- roughly a tenth of a full chunk's work, which is + // what makes trading Ckpt-fold memory for it worthwhile. + // + // Scoped so S is dead before process_chunk runs: it is a _M16xDk, and + // having it live alongside S_tile and dS_tile would put three of them in + // flight. + { + _M16xDk S; + S.load(ck_base + seg * Dv * Dk, Dk); + + auto replay_one = [&](const int j, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + S.store(seg_base + j * Dv * Dk, Dk); + + const int c = seg_start + j; + auto k_c = k_base + c * C * Hk * Dk; + auto g_c = g_base + c * C * Hv; + auto cd_c = cd_base + c * C * Dv; + + float gv = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + g_c[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + auto gs = simd_prefix_inclusive_sum(gv); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gs); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float g_last = gamma[C - 1]; + const float dec[2] = { + metal::fast::exp(g_last - gamma[fm]), + metal::fast::exp(g_last - gamma[fm1])}; + + _M16x16 delta_tile; + delta_tile.load(cd_c, Dv); + + S = S * metal::fast::exp(g_last); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 Kd; + if constexpr (B) { + Kd.load_rows(k_c + kk, Dk * Hk, valid_rows); + } else { + Kd.load(k_c + kk, Dk * Hk); + } + Kd = scale_rows(Kd, dec); + MMA16x32x16(S, kk / 16, delta_tile, true, 0, Kd, false, 0); + } + }; + + for (int j = 0; j < n_full; j++) { + replay_one(j, C, metal::false_type{}); + } + if constexpr (HasTail) { + replay_one(n_in_seg - 1, tail, metal::true_type{}); + } + } + + auto set_chunk_ptrs = [&](const int j) __attribute__((always_inline)) { + const int c = seg_start + j; + q_ = q_base + c * C * Hk * Dk; + k_ = k_base + c * C * Hk * Dk; + dq_ = dq_base + c * C * Hk * Dk; + dk_ = dk_base + c * C * Hk * Dk; + v_ = v_base + c * C * Hv * Dv; + dv_ = dv_base + c * C * Hv * Dv; + co_ = co_base + c * C * Hv * Dv; + g_ = g_base + c * C * Hv; + beta_ = beta_base + c * C * Hv; + dg_ = dg_base + c * C * Hv; + db_ = db_base + c * C * Hv; + i_chunk_mats = cm_base + c * 3 * 256; + i_chunk_delta = cd_base + c * C * Dv; + i_seg = seg_base + j * Dv * Dk; + }; + + if constexpr (HasTail) { + set_chunk_ptrs(n_in_seg - 1); + process_chunk(tail, metal::true_type{}); + } + for (int j = n_full - 1; j >= 0; --j) { + set_chunk_ptrs(j); + process_chunk(C, metal::false_type{}); + } + }; + + // Only the last segment can hold a short chunk + int seg = n_ckpt - 1; + if (tail != C) { + process_segment(seg, metal::true_type{}); + --seg; + } + for (; seg >= 0; --seg) { + process_segment(seg, metal::false_type{}); + } + + dS_tile.store(o_dh, Dk); +} + +template +[[kernel]] void gated_delta_vjp_fused_nax1( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; // KdKb_tmp folded away + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_r0 < valid_rows) { + db_st[_r0] = _b0; + } + if (_r1 < valid_rows) { + db_st[_r1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (_r0 == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (_r0 < valid_rows) { + dg_st[_r0] = _d0; + } + if (_r1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (_r1 < valid_rows) { + dg_st[_r1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +template +METAL_FUNC void reduce_tile_tg0( + thread TileT& acc, + threadgroup float* scratch, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short lo = kNSG / 2; lo > 0; lo >>= 1) { + if (sg_id >= lo && sg_id < 2 * lo) { + const short dst = sg_id - lo; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(dst * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id < lo) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(sg_id * 32 + simd_lane_id) * kE + _i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax0( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 2) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 NI_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(NI_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(NI_tile, _i) = (_c.x == _c.y) ? -1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw; + MUL_NAX(QKt_tile, QKt_tile, tri_decay); + + TUinv_tile = TWinv_tile; + MUL_NAX(TUinv_tile, TUinv_tile, tri_decay); + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + MUL_NAX(dODT, dODT, tri_decay); + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + SCALE2_P(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + _M16x32 Kb_tile; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + + SCALE_BETA_NAX_O(Kb_tile, K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg(dq_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dq_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + const short _fn = _c.x + _f * 16; + const short _fm = _c.y; + if (_fm < valid_rows) { + const int idx = _fm * Hk * Dk + kk + _fn; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } else { + // Dv reduction done in threadgroup memory: single writer. + mlx_atomic_store_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } + } + } + } + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + SCALE_BETA_NAX(dVb_tile, beta_fm); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dVb_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.y < valid_rows) { + dv_[_c.y * Hv * Dv + dv_idx + _c.x] = AT_NAX(dVb_tile, _i); + } + } + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + MUL_NAX(dTU_tile, dTU_tile, tri_decay); // dTinv = dT_W + (dT_U * D). + + ADD_NAX(dTinv_tile, dTW_tile, dTU_tile); + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + SCALE_NAX(dA_tile, -1.0f); + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = dA_tile; + SCALE_TRIEQ_NAX1(G_tile, beta_fm); + _M16x16 GGt_tile; + GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile; + _M16x16 R_tile; + + MUL_NAX(P_tile, QKt_raw, dODT); + MUL_NAX(R_tile, TWinv_tile, dTU_tile); + ADD_NAX(dgam_pair, P_tile, R_tile); + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, dKb_tile, t1, KdKb_tmp; + _M16x32 QgdOS, TbK, KdKC_blk; + + load_seq(Q_tile, q_ + kk, Hk * Dk); // [16 x 32] + load_seq(K_tile, k_ + kk, Dk * Hk); // [16 x 32] + + // gamma: + rowsum(Q * gamma*(dO @ S)) + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(QgdOS, row_exp); + MULA_NAX(dgam_row32, QgdOS, Q_tile); + + // dODT.T @ Q and (G + G.T) @ K + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + // B (T.T @ dW_raw), dW_raw = -gamma * (ddelta @ S) + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + + // dbeta term uses dK_b before the beta scaling, with raw K + MUL_NAX(KdKb_tmp, dKb_tile, K_tile); + ADD_NAX(KdKb, KdKb, KdKb_tmp); + + SCALE_BETA_NAX(dKb_tile, beta_fm); + ADD_NAX(dk_acc, dk_acc, dKb_tile); + + // D_C * (delta @ dS) + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + SCALE2_P(t1, dec_exp); + ADD_NAX(dk_acc, dk_acc, t1); + + // gamma stuff + MUL_NAX(KdKC_blk, K_tile, t1); + SUB_NAX(dgam_row32, dgam_row32, KdKC_blk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC_blk)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC_blk, _i); + } + + // gamma stuff: gamma_C * sum(S * dS) + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16)[_i] * + dS_tile.frag_at(0, kk / 16)[_i]; + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]; + } + + // gamma stuff + SCALE_BETA_NAX(K_tile, beta_fm); + SCALE_NAX(dW_raw, -1.0f); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + MULS_NAX(dgam_row32, TbK, dW_raw); + + // dk reduces over Dv the same way dq does. + reduce_tile_tg(dk_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dk_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * Hk * Dk + kk + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } else { + mlx_atomic_store_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } + } + } + } + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + _M16x16 VdVb; + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + MUL_NAX(VdVb, V_tile, TUddelta); + + _M16x16 AKKt; + TRIL_NAX(AKKt, dA_tile); + MUL_NAX(AKKt, AKKt, KKt_tile); + + ADD_NAX(VdVb, VdVb, AKKt); + ROWSUM1_NAX(dbeta_acc, VdVb); + ROWSUM1_NAX(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _r = 0; _r < 2; _r++) { + float _b = dbeta_acc[_r]; + _b += simd_shuffle_xor(_b, ushort(1)); + _b += simd_shuffle_xor(_b, ushort(8)); + + const short _row = fm + _r * mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_row < valid_rows) { + db_st[_row] = _b; + } + } + } + + // Fold the two fragments: only the row sum of dgam_row is used downstream, + // and row sums are additive across fragments. + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + AT_NAX(dgam_row, _i) = AT_NAX(dgam_row32, _i) + + AT_NAX(dgam_row32, mlx::steel::BaseNAXFrag::kElemsPerFrag + _i); + } + + // reduce and store dgamma: rowsum(row + pair) - colsum(pair) + ADD_NAX(dgam_row, dgam_row, dgam_pair); + _M16x16 dgam_tile, cs_tile; + MM16x16x16(dgam_tile, 0, dgam_row, false, 0, Ones_tile, false, 0); + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + SUB_NAX(dgam_tile, dgam_tile, cs_tile); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dgam_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.x == 0) { + dg_st[_c.y] = AT_NAX(dgam_tile, _i); + } + } + // The gamma_{C-1}-only term lands on the same slot the loop just wrote, + // so the lanes must be ordered before lane 0 accumulates into it. + simdgroup_barrier(mem_flags::mem_threadgroup); + const float dgam_last_red = simd_sum(dgam_last); + if (simd_lane_id == 0 && valid_rows > 0) { + dg_st[valid_rows - 1] += dgam_last_red; + } + } + + // One barrier serves both stages; simdgroup 0 commits. + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + SCALE_ROW_P(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + NSCALE_ROW_P(W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +// Postprocessing dg. Similar idea to what is done in the triton kernel. +// The kernel computes the gradient with respect to the log. +template +[[kernel]] void gated_delta_dgamma_to_dg( + const device InT* g [[buffer(0)]], + device float* dg [[buffer(1)]], // in: dL/dgamma, out: dL/dg + constant int& T [[buffer(2)]], + constant int& Hv [[buffer(3)]], + constant int& n_total [[buffer(4)]], // B * Hv + uint2 pos [[thread_position_in_grid]]) { + const int n = int(pos.x); // b * Hv + hv + const int c = int(pos.y); // chunk index + + if (n >= n_total) { + return; + } + + const int t0 = c * C; + if (t0 >= T) { + return; + } + const int len = min(C, T - t0); + + const int b_idx = n / Hv; + const int hv_idx = n % Hv; + const int base = b_idx * T * Hv + hv_idx; + + float acc = 0.0f; + for (int j = len - 1; j >= 0; --j) { + const int idx = base + (t0 + j) * Hv; + acc += static_cast(dg[idx]); + float gv = metal::max(static_cast(g[idx]), 1e-6f); + dg[idx] = acc / gv; + } +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.metal new file mode 100644 index 000000000..954478d42 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax_vjp.metal @@ -0,0 +1,44 @@ +#include "mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h" +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +#define instantiate_gdu_vjp_nax(in_type, dk, dv, hk, hv, c, ckpt) \ + instantiate_kernel( \ + "gated_delta_vjp_fused_nax_" #in_type "_" #dk "_" #dv \ + "_" #hk "_" #hv "_" #c "_" #ckpt, \ + gated_delta_vjp_fused_nax, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c, \ + ckpt) + +#define instantiate_gdu_vjp_dims(in_type, dk, dv, hk, hv) \ + instantiate_gdu_vjp_nax(in_type, dk, dv, hk, hv, 16, 1) \ + instantiate_gdu_vjp_nax(in_type, dk, dv, hk, hv, 16, 4) \ + instantiate_gdu_vjp_nax(in_type, dk, dv, hk, hv, 16, 8) \ + instantiate_gdu_vjp_nax(in_type, dk, dv, hk, hv, 16, 16) + +#define instantiate_gdu_vjp(in_type) \ + instantiate_gdu_vjp_dims(in_type, 128, 128, 24, 24) \ + instantiate_gdu_vjp_dims(in_type, 128, 128, 32, 32) \ + instantiate_gdu_vjp_dims(in_type, 128, 128, 16, 32) \ + instantiate_gdu_vjp_dims(in_type, 128, 128, 16, 16) \ + instantiate_gdu_vjp_dims(in_type, 128, 128, 16, 48) + +instantiate_gdu_vjp(float); +instantiate_gdu_vjp(bfloat16_t); + +// Postprocessing: converts dL/dgamma to dL/dg. Not templated on ckpt. +#define instantiate_gdu_dgamma(in_type, c) \ + instantiate_kernel( \ + "gated_delta_dgamma_to_dg_" #in_type "_" #c, \ + gated_delta_dgamma_to_dg, \ + in_type, \ + c) + +instantiate_gdu_dgamma(float, 16); +instantiate_gdu_dgamma(bfloat16_t, 16); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.h new file mode 100644 index 000000000..13b96416c --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.h @@ -0,0 +1,184 @@ +// Copyright © 2024 Apple Inc. +#pragma once + +#include +#include + +#include "mlx/backend/metal/kernels/atomic.h" +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +template +[[kernel]] void gated_delta_vjp_seq( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* b [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B*Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device mlx_atomic* dg [[buffer(12)]], + device mlx_atomic* db [[buffer(13)]], + device float* dh [[buffer(14)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + const int qk_stride = Hk * Dk; + const int v_stride = Hv * Dv; + const int g_stride = Hv; + + auto q_base = q + b_idx * T * qk_stride + hk_idx * Dk; + auto k_base = k + b_idx * T * qk_stride + hk_idx * Dk; + auto dq_base = dq + b_idx * T * qk_stride + hk_idx * Dk; + auto dk_base = dk + b_idx * T * qk_stride + hk_idx * Dk; + + auto v_base = v + b_idx * T * v_stride + hv_idx * Dv; + auto dv_base = dv + b_idx * T * v_stride + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * v_stride + hv_idx * Dv; + + auto g_base = g + b_idx * T * g_stride + hv_idx; + auto b_base = b + b_idx * T * g_stride + hv_idx; + auto dg_base = dg + b_idx * T * g_stride + hv_idx; + auto db_base = db + b_idx * T * g_stride + hv_idx; + + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + + float s_hat[n_per_t]; + auto base_state = cot_h + (n * Dv + dv_idx) * Dk; + for (int i = 0; i < n_per_t; i++) { + s_hat[i] = base_state[n_per_t * dk_idx + i]; + } + + // Only this thread's n_per_t slice of the state is ever needed, so a whole + // segment of entry states costs Ckpt * n_per_t registers and the replay needs + // no device buffer. + float seg[Ckpt][n_per_t]; + float s_prev[n_per_t]; + float s_dec[n_per_t]; + + for (int seg_idx = n_ckpt - 1; seg_idx >= 0; --seg_idx) { + const int t0 = seg_idx * Ckpt; + const int seg_len = metal::min(Ckpt, T - t0); + + // Replay forward from the checkpoint, recording the entry state of each + // step in the segment. + auto c_state = + state_cache + n * n_ckpt * Dv * Dk + seg_idx * Dv * Dk + dv_idx * Dk; + + float s[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + s[i] = c_state[n_per_t * dk_idx + i]; + } + + for (int j = 0; j < seg_len; ++j) { + const int t = t0 + j; + + for (int i = 0; i < n_per_t; ++i) { + seg[j][i] = s[i]; + } + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto k_t = k_base + t * qk_stride; + + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + const int s_idx = n_per_t * dk_idx + i; + s[i] *= gamma; + kv_mem += s[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + + float delta = + beta * (static_cast(v_base[t * v_stride + dv_idx]) - kv_mem); + + for (int i = 0; i < n_per_t; ++i) { + s[i] += delta * static_cast(k_t[n_per_t * dk_idx + i]); + } + } + + // Backward over the same steps, newest first. + for (int j = seg_len - 1; j >= 0; --j) { + const int t = t0 + j; + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto q_t = q_base + t * qk_stride; + auto k_t = k_base + t * qk_stride; + auto dq_t = dq_base + t * qk_stride; + auto dk_t = dk_base + t * qk_stride; + + float kv_mem = 0.0f; + float co = static_cast(co_base[t * v_stride + dv_idx]); + float w = 0.0f; + for (int i = 0; i < n_per_t; i++) { + const int s_idx = n_per_t * dk_idx + i; + s_prev[i] = seg[j][i]; + s_dec[i] = s_prev[i] * gamma; + kv_mem += s_dec[i] * static_cast(k_t[s_idx]); + + s_hat[i] += co * static_cast(q_t[s_idx]); + + w += s_hat[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + w = simd_sum(w); + + if (thread_index_in_simdgroup == 0) { + dv_base[t * v_stride + dv_idx] = beta * w; + } + + float u = static_cast(v_base[t * v_stride + dv_idx]) - kv_mem; + float delta = beta * u; + + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(db_base + t * g_stride, w * u, 0); + } + + float dgamma = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + + float s_t = s_dec[i] + delta * static_cast(k_t[s_idx]); + mlx_atomic_fetch_add_explicit(dq_t, co * s_t, s_idx); + + float contrib = beta * (u * s_hat[i] - w * s_dec[i]); + mlx_atomic_fetch_add_explicit(dk_t, contrib, s_idx); + + s_hat[i] -= beta * w * static_cast(k_t[s_idx]); + + dgamma += s_hat[i] * s_prev[i]; + } + dgamma = simd_sum(dgamma); + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(dg_base + t * g_stride, dgamma, 0); + } + + for (int i = 0; i < n_per_t; ++i) { + s_hat[i] *= gamma; + } + } + } + + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + dh[(n * Dv + dv_idx) * Dk + s_idx] = s_hat[i]; + } +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.metal new file mode 100644 index 000000000..7112750d4 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_vjp.metal @@ -0,0 +1,35 @@ +#include "mlx/backend/metal/kernels/gated_delta_update_vjp.h" +#include "mlx/backend/metal/kernels/gated_delta_update.h" +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +#define instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_vjp_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_vjp_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) + +#define instantiate_gated_delta_vjp_dims(in_type, dk, dv, hk, hv) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 16) + +#define instantiate_gated_delta_vjp(in_type) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 64) + +instantiate_gated_delta_vjp(float); +instantiate_gated_delta_vjp(bfloat16_t); +instantiate_gated_delta_vjp(float16_t); \ No newline at end of file diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv_masked.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv_masked.h index 58ab74405..18c269906 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv_masked.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv_masked.h @@ -452,7 +452,7 @@ struct GEMVTKernel { const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x); const int out_mask_offset = - !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0]; + !has_output_mask ? 0 : n_block_idx * out_mask_strides[0]; int mat_mask_offset = !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0]; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h index 2ba547400..fe2e2b2fe 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h @@ -32,7 +32,7 @@ template if (src_contiguous) { out[idx] = src[batch_idx * src_batch_size + src_index]; } else { - out[idx] = src[elem_to_loc( + out[idx] = src[elem_to_loc( batch_idx * src_batch_size + src_index, src_shapes, src_strides, diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h index 3604f20f6..56fead768 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h @@ -38,8 +38,8 @@ METAL_FUNC void scatter_impl( for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) { LocT out_idx = out_offset; for (int i = 0; i < NIDX; ++i) { - auto idx_loc = indices.row_contiguous[i] - ? ind_idx + LocT idx_loc = indices.row_contiguous[i] + ? LocT(ind_idx) : elem_to_loc( ind_idx, &indices.shapes[indices.ndim * i], @@ -50,7 +50,7 @@ METAL_FUNC void scatter_impl( out_idx += static_cast(idx_val) * static_cast(out_strides[ax]); } - auto upd_idx = ind_idx * static_cast(upd_size) + gid.x; + LocT upd_idx = ind_idx * static_cast(upd_size) + gid.x; if constexpr (!UPD_ROW_CONTIG) { upd_idx = elem_to_loc(upd_idx, upd_shape, upd_strides, upd_ndim); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h index be1ebff59..48fe55cd9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h @@ -28,15 +28,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -46,8 +61,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -61,7 +76,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -71,8 +86,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -86,7 +101,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -107,15 +122,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -125,8 +155,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -141,7 +171,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -151,8 +181,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -166,7 +196,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -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,14 +557,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -585,9 +679,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(); @@ -753,7 +847,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -765,7 +859,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(); @@ -787,19 +882,37 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; y += tid.x * out_vec_size + out_row; - for (int k = 0; k < in_vec_size; k += block_size) { + const int aligned_end = (in_vec_size / block_size) * block_size; + + for (int k = 0; k < aligned_end; k += block_size) { U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -812,9 +925,37 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } + // Tail block: handles K that is a multiple of 512 (dispatch gate) but not + // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for + // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; + // their result[row] stays at 0 and contributes 0 to the simd_sum below. + if (aligned_end < in_vec_size) { + bool in_bounds = + (aligned_end + static_cast(simd_lid) * values_per_thread) < + in_vec_size; + if (in_bounds) { + U sum = load_vector(x, x_thread); + + for (int row = 0; row < results_per_simdgroup; row++) { + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; + + U s = sl[0]; + U b = bl[0]; + result[row] += + qdot(wl, x_thread, s, b, sum); + } + } + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1193,6 +1334,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1201,7 +1343,7 @@ METAL_FUNC void qmm_t_impl( const device T* scales, const device T* biases, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -1226,7 +1368,7 @@ METAL_FUNC void qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1602,6 +1744,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( @@ -1642,7 +1785,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1660,6 +1803,7 @@ template < int group_size, const int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv( @@ -1946,7 +2090,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -1977,7 +2121,7 @@ template < const device T* scales [[buffer(1)]], const device T* biases [[buffer(2)]], const device T* x [[buffer(3)]], - device T* y [[buffer(4)]], + device float* y [[buffer(4)]], const constant int& K [[buffer(5)]], const constant int& N [[buffer(6)]], const constant int& M [[buffer(7)]], @@ -2005,7 +2149,7 @@ template < biases += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2336,7 +2480,7 @@ template < s_strides, b_strides, tid); - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2617,14 +2761,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; @@ -2728,7 +2881,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/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal index 069482cba..8e46ef97b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal @@ -23,6 +23,16 @@ bits, \ batched) +#define instantiate_quantized_qmv(name, type, group_size, bits, batched, partial_rows) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_pr_" #partial_rows "_hgs_0_r_4", \ + name, \ + type, \ + group_size, \ + bits, \ + batched, \ + partial_rows) + #define instantiate_quantized_aligned(name, type, group_size, bits, aligned) \ instantiate_kernel( \ #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \ @@ -90,9 +100,14 @@ instantiate_quantized_batched(name, type, group_size, bits, 1) \ instantiate_quantized_batched(name, type, group_size, bits, 0) -#define instantiate_quantized_all_batched(type, group_size, bits) \ - instantiate_quantized_batched_wrap(affine_qmv_fast, type, group_size, bits) \ - instantiate_quantized_batched_wrap(affine_qmv, type, group_size, bits) \ +#define instantiate_quantized_qmv_wrap(name, type, group_size, bits, partial_rows) \ + instantiate_quantized_qmv(name, type, group_size, bits, 1, partial_rows) \ + instantiate_quantized_qmv(name, type, group_size, bits, 0, partial_rows) + +#define instantiate_quantized_all_batched(type, group_size, bits) \ + instantiate_quantized_qmv_wrap(affine_qmv_fast, type, group_size, bits, 1) \ + instantiate_quantized_qmv_wrap(affine_qmv_fast, type, group_size, bits, 0) \ + instantiate_quantized_qmv_wrap(affine_qmv, type, group_size, bits, 0) \ instantiate_quantized_batched_wrap(affine_qvm, type, group_size, bits) \ instantiate_quantized_batched_wrap(affine_qmm_n, type, group_size, bits) @@ -172,6 +187,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/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h index a34774992..95f26868f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h @@ -31,15 +31,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -49,8 +64,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -64,7 +79,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -74,8 +89,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -89,7 +104,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -110,15 +125,30 @@ 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -128,8 +158,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -144,7 +174,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -154,8 +184,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { 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]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -169,7 +199,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -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 == 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); + } + } - 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); @@ -487,14 +567,36 @@ 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}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -587,9 +689,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -715,9 +817,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal index 9557fd838..4d61232ea 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal @@ -102,6 +102,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/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal index 3f5f91be4..0cd290ab8 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal @@ -19,6 +19,8 @@ instantiate_init_reduce(or, bool_, bool, Or) #define instantiate_init_sum_prod(name, op) \ instantiate_init_reduce(name, int32, int32_t, op) \ instantiate_init_reduce(name, int64, int64_t, op) \ + instantiate_init_reduce(name, uint32, uint32_t, op) \ + instantiate_init_reduce(name, uint64, uint64_t, op) \ instantiate_init_reduce(name, float16, float16_t, op) \ instantiate_init_reduce(name, bfloat16, bfloat16_t, op) \ instantiate_init_reduce(name, float32, float, op) \ @@ -131,7 +133,8 @@ instantiate_init_min_max(max, Max) 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_reduce_functions(name, float32, float, bool, op) \ + instantiate_reduce_functions(name, complex64, complex64_t, bool, op) instantiate_and_or(and, And) instantiate_and_or(or, Or) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h index b1546adb5..fce36a4cd 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h @@ -117,7 +117,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); @@ -172,6 +172,7 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], + const constant int64_t& column_offset [[buffer(11)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -193,7 +194,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT out_idx = gid.y + gsize.y * IdxT(gid.z); @@ -243,7 +244,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += out_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -311,7 +312,8 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], - const constant size_t& out_size [[buffer(11)]], + const constant int64_t& column_offset [[buffer(11)]], + const constant size_t& out_size [[buffer(12)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -336,7 +338,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT full_idx = gid.y + gsize.y * IdxT(gid.z); @@ -383,7 +385,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += full_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h index 604efa785..9a1e363e1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h @@ -3,6 +3,7 @@ template [[kernel]] void init_reduce( device T* out [[buffer(0)]], - uint tid [[thread_position_in_grid]]) { - out[tid] = Op::init; + uint2 index [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + out[index.x + grid_dim.x * int64_t(index.y)] = Op::init; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 04206c940..d7862d256 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -2,10 +2,28 @@ // clang-format off #include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/sdpa_blocked.h" #include "mlx/backend/metal/kernels/sdpa_vector.h" using namespace metal; +instantiate_kernel( + "sdpa_blocked_scale_copy_bfloat16", + sdpa_blocked_scale_copy, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_scale_copy_float16", + sdpa_blocked_scale_copy, + float16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_bfloat16", + sdpa_blocked_causal_mask, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_float16", + sdpa_blocked_causal_mask, + float16_t) + // SDPA vector instantiations #define instantiate_sdpa_vector_aggregation(type, value_dim) \ instantiate_kernel( \ @@ -61,4 +79,17 @@ using namespace metal; instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) instantiate_sdpa_vector_heads(float16_t) - // clang-format on + +#define instantiate_sdpa_vector_tk1(type, dim) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_" #type "_" #dim "_" #dim "_tk1", \ + sdpa_vector_2pass_1, \ + type, \ + dim, \ + dim, \ + 1) + +instantiate_sdpa_vector_tk1(bfloat16_t, 256) +instantiate_sdpa_vector_tk1(float, 192) +instantiate_sdpa_vector_tk1(float, 256) +// clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention_vjp.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention_vjp.metal new file mode 100644 index 000000000..ddd9f4bcd --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention_vjp.metal @@ -0,0 +1,154 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +template +[[kernel]] void sdpa_vjp_odo( + const device InT* o [[buffer(0)]], // [B, H, qL, Dv] + const device InT* cot_o [[buffer(1)]], // [B, H, qL, Dv] + device float* odo [[buffer(2)]], // [B, H, qL] + constant int& qL [[buffer(3)]], + uint3 tpg [[thread_position_in_grid]], + uint simd_lane_id [[thread_index_in_simdgroup]]) { + const uint bh = tpg.z; + const int i = int(tpg.y); + + if (i >= qL) { + return; + } + + const size_t row = (size_t)bh * qL + i; + auto o_i = o + row * Dv; + auto co_i = cot_o + row * Dv; + + float acc = 0.0f; + for (int d = int(simd_lane_id); d < Dv; d += 32) { + acc += float(o_i[d]) * float(co_i[d]); + } + acc = simd_sum(acc); + + if (simd_lane_id == 0) { + odo[row] = acc; + } +} + +struct SDPAVJPTileParams { + int bq; + int bk; + int qL; + int i0; + int j0; + float scale; + int diag_off; + int causal; +}; + +template +[[kernel]] void sdpa_vjp_ds( + const device T* S [[buffer(0)]], // [BH, bq, bk] + const device T* dP [[buffer(1)]], // [BH, bq, bk] + const device float* lse [[buffer(2)]], // [B, H, qL] + const device float* odo [[buffer(3)]], // [B, H, qL] + device T* dS [[buffer(4)]], // [BH, bq, bk] + device T* P [[buffer(5)]], // [BH, bk, bq] + const constant SDPAVJPTileParams& p [[buffer(6)]], + uint3 gid [[thread_position_in_grid]]) { + int col = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (row >= p.bq || col >= p.bk) { + return; + } + + size_t sbase = (size_t(bh) * size_t(p.bq) + size_t(row)) * size_t(p.bk); + size_t qi = size_t(bh) * size_t(p.qL) + size_t(p.i0 + row); + + float l = lse[qi]; + float dlt = odo[qi]; + // Fully masked rows carry lse == -inf; exp(-inf - -inf) would be NaN. + bool dead = (l < 0.0f) && metal::isinf(l); + int lim = p.i0 + row + p.diag_off; + + float pv; + if (dead || (p.causal != 0 && (p.j0 + col) > lim)) { + pv = 0.0f; + } else { + pv = metal::fast::exp(static_cast(S[sbase + col]) * p.scale - l); + } + float dsv = pv * (static_cast(dP[sbase + col]) - dlt) * p.scale; + + dS[sbase + col] = static_cast(dsv); + P[sbase + col] = static_cast(pv); +} + +template +[[kernel]] void sdpa_vjp_reduce( + const device T* src [[buffer(0)]], + device float* acc [[buffer(1)]], + device T* out [[buffer(2)]], + const constant int& rows [[buffer(3)]], + const constant int& dim [[buffer(4)]], + const constant int& group [[buffer(5)]], + const constant int& acc_rows [[buffer(6)]], + const constant int& row_off [[buffer(7)]], + uint3 gid [[thread_position_in_grid]]) { + int c = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (c >= dim || row >= rows) { + return; + } + + size_t sbase = size_t(bh) * size_t(group) * size_t(rows) * size_t(dim); + float sum = 0.0f; + for (int g = 0; g < group; ++g) { + size_t si = sbase + + (size_t(g) * size_t(rows) + size_t(row)) * size_t(dim) + size_t(c); + sum += static_cast(src[si]); + } + + size_t ai = + (size_t(bh) * size_t(acc_rows) + size_t(row_off + row)) * size_t(dim) + + size_t(c); + float total = Accum ? acc[ai] + sum : sum; + acc[ai] = total; + out[ai] = static_cast(total); +} + +// Instantiations +#define instantiate_odo(in_type, dv) \ + instantiate_kernel( \ + "sdpa_vjp_odo_" #in_type "_" #dv, sdpa_vjp_odo, in_type, dv) + +#define instantiate_odo_shapes(in_type) \ + instantiate_odo(in_type, 64); \ + instantiate_odo(in_type, 72); \ + instantiate_odo(in_type, 80); \ + instantiate_odo(in_type, 96); \ + instantiate_odo(in_type, 128); \ + instantiate_odo(in_type, 192); \ + instantiate_odo(in_type, 256); \ + instantiate_odo(in_type, 512); \ + +instantiate_odo_shapes(bfloat16_t); +instantiate_odo_shapes(float16_t); +instantiate_odo_shapes(float); + +#define instantiate_sdpa_vjp_ds(tname, type) \ + instantiate_kernel("sdpa_vjp_ds_" #tname, sdpa_vjp_ds, type) + +instantiate_sdpa_vjp_ds(float, float); +instantiate_sdpa_vjp_ds(float16_t, float16_t); +instantiate_sdpa_vjp_ds(bfloat16_t, bfloat16_t); + +#define instantiate_sdpa_vjp_reduce(tname, type) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_add_" #tname, sdpa_vjp_reduce, type, true) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_set_" #tname, sdpa_vjp_reduce, type, false) + +instantiate_sdpa_vjp_reduce(float, float); +instantiate_sdpa_vjp_reduce(float16_t, float16_t); +instantiate_sdpa_vjp_reduce(bfloat16_t, bfloat16_t); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_blocked.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_blocked.h new file mode 100644 index 000000000..be4e900c6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_blocked.h @@ -0,0 +1,45 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "mlx/backend/metal/kernels/defines.h" + +using namespace metal; + +template +[[kernel]] void sdpa_blocked_scale_copy( + const device T* q [[buffer(0)]], + device T* q_block [[buffer(1)]], + constant float& scale [[buffer(2)]], + constant int64_t* q_strides [[buffer(3)]], + constant int& heads [[buffer(4)]], + constant int& rows [[buffer(5)]], + constant int& head_dim [[buffer(6)]], + constant int& q_offset [[buffer(7)]], + uint3 elem [[thread_position_in_grid]]) { + int d = elem.x; + int row = elem.y; + int head = elem.z % heads; + int batch = elem.z / heads; + int64_t q_index = int64_t(batch) * q_strides[0] + + int64_t(head) * q_strides[1] + int64_t(q_offset + row) * q_strides[2] + d; + size_t q_block_index = (size_t(elem.z) * rows + row) * head_dim + d; + q_block[q_block_index] = q[q_index] * static_cast(scale); +} + +template +[[kernel]] void sdpa_blocked_causal_mask( + device T* scores [[buffer(0)]], + constant int& tile_rows [[buffer(1)]], + constant int& columns [[buffer(2)]], + uint2 elem [[thread_position_in_grid]]) { + int column = elem.x; + int score_row = elem.y; + int row = score_row % tile_rows; + if (column > row) { + size_t index = size_t(score_row) * columns + columns - tile_rows + column; + scores[index] = Limits::finite_min; + } +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h index 4594f0389..a38265f55 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h @@ -216,6 +216,7 @@ template const int q_batch_head_idx = tid.x; const int q_seq_idx = tid.y; const int kv_head_idx = q_batch_head_idx / gqa_factor; + const int kv_seq_idx = simd_gid; const int o_offset = q_batch_head_idx * tpg.y + q_seq_idx; const int q_offset = query_transposed ? tpg.x * q_seq_idx + q_batch_head_idx : o_offset; @@ -226,11 +227,11 @@ template simd_lid * v_per_thread; if (bool_mask) { bmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } if (float_mask) { fmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } out += o_offset * V + simd_gid * v_per_thread; @@ -331,7 +332,7 @@ template } } -template +template [[kernel]] void sdpa_vector_2pass_1( const device T* queries [[buffer(0)]], const device T* keys [[buffer(1)]], @@ -414,8 +415,73 @@ template sum_exp_score = 1; } + const int stride_k_tok = blocks * int(k_seq_stride); + const int stride_v_tok = blocks * int(v_seq_stride); + const int stride_mask_tok = + (bool_mask || float_mask) ? blocks * mask_kv_seq_stride : 0; + // For each key - for (int i = block_idx; i < N; i += blocks) { + int i = block_idx; + for (; i + (TK - 1) * blocks < N; i += blocks * TK) { + // Prefetch: issue all TK K/V loads first so the HW can overlap them. + T k_reg[TK][qk_per_thread]; + T v_reg[TK][v_per_thread]; + for (int t = 0; t < TK; t++) { + const device T* kp = keys + t * stride_k_tok; + for (int j = 0; j < qk_per_thread; j++) { + k_reg[t][j] = kp[j]; + } + } + for (int t = 0; t < TK; t++) { + const device T* vp = values + t * stride_v_tok; + for (int j = 0; j < v_per_thread; j++) { + v_reg[t][j] = vp[j]; + } + } + + for (int t = 0; t < TK; t++) { + int ii = i + t * blocks; + bool use_key = true; + if (do_causal) { + use_key = ii <= (N - q_seq_len + int(q_seq_idx)); + } else if (bool_mask) { + use_key = bmask[t * stride_mask_tok]; + } else if (float_mask) { + use_key = (fmask[t * stride_mask_tok] >= Limits::finite_min); + } + if (use_key) { + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * static_cast(k_reg[t][j]); + } + score = simd_sum(score); + if (float_mask) { + score += static_cast(fmask[t * stride_mask_tok]); + } + + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * static_cast(v_reg[t][j]); + } + } + } + + keys += TK * stride_k_tok; + values += TK * stride_v_tok; + if (bool_mask) { + bmask += TK * stride_mask_tok; + } + if (float_mask) { + fmask += TK * stride_mask_tok; + } + } + + // Tail (fewer than TK remaining keys) + for (; i < N; i += blocks) { bool use_key = true; if (do_causal) { use_key = i <= (N - q_seq_len + int(q_seq_idx)); @@ -427,8 +493,8 @@ template if (use_key) { // Compute the i-th score U score = 0; - for (int i = 0; i < qk_per_thread; i++) { - score += q[i] * keys[i]; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * keys[j]; } score = simd_sum(score); @@ -445,19 +511,19 @@ template sum_exp_score = sum_exp_score * factor + exp_score; // Update the output accumulator - for (int i = 0; i < v_per_thread; i++) { - o[i] = o[i] * factor + exp_score * values[i]; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * values[j]; } } // Move the pointers to the next kv - keys += blocks * int(k_seq_stride); - values += blocks * int(v_seq_stride); + keys += stride_k_tok; + values += stride_v_tok; if (bool_mask) { - bmask += blocks * mask_kv_seq_stride; + bmask += stride_mask_tok; } if (float_mask) { - fmask += blocks * mask_kv_seq_stride; + fmask += stride_mask_tok; } } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h index d5cb00e42..d8f5efdc1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h @@ -14,6 +14,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; struct MaxOp { template @@ -77,6 +78,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -103,6 +105,12 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -532,6 +540,20 @@ template < Otile.template row_bin_op(sum_score); threadgroup_barrier(mem_flags::mem_none); + // Store the logsumexp for the backward pass + if (save_lse && sn == 0) { + using stile_t = decltype(Stile); + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + sm + i * stile_t::kFragRows; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = M_LN2_F * (max_score[i] + metal::log2(sum_score[i])); + } + } + } + // Store results O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; @@ -545,4 +567,4 @@ template < } else { Otile.template store(O, params->O_strides[2]); } -} +} \ No newline at end of file diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index 7a7099b8f..13cf24f78 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -17,6 +17,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; template struct TransformScale { @@ -90,6 +91,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -117,6 +119,13 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + // [B, H, qL], contiguous. + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -475,6 +484,16 @@ template < Otile.template row_bin_op(rcp); + if (save_lse && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + i * otile_t::kFragRowsJump + sm; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = max_score[i] * M_LN2_F + metal::log(sum_score[i]); + } + } + } + // Store results O += tm * int(params->O_strides[2]); @@ -516,6 +535,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* LSE [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -903,6 +923,21 @@ template < rcp[i] = 1.f / sum_score[i]; } + if (save_lse && d_group == 0 && sn == 0) { + const int lse_row_base = int(tid.x) * BQ + tm; + const int lse_head_off = (int(tid.z) * params->H + int(tid.y)) * params->qL; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = lse_row_base + i * stile_t::kFragRowsJump + sm; + if (row < params->qL) { + LSE[lse_head_off + row] = (sum_score[i] == 0) + ? -INFINITY + : float(M_LN2_F * (max_score[i] + metal::log2(sum_score[i]))); + } + } + } + Otile.template row_bin_op(rcp); if (!align_Q && is_last_q) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index 2296187fe..eba254155 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal @@ -22,11 +22,13 @@ instantiate_attn_dsplit(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 128, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 64, 64, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 256, 256, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 128, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 96, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 96, 64, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 64, 64, 4, 1, mname, mtype) + #define instantiate_attn_mask_helper(iname, itype) \ instantiate_attn_shapes_helper(iname, itype, iname, itype) \ instantiate_attn_shapes_helper(iname, itype, bool_, bool) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h index f559596b7..f948ec069 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h @@ -46,7 +46,7 @@ implicit_gemm_conv_2d( using loader_a_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DInputBlockLoaderSmallChannels< @@ -84,7 +84,7 @@ implicit_gemm_conv_2d( // Weight loader using loader_b_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DWeightBlockLoaderSmallChannels< diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal index 7851e5ef1..475d28e53 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal @@ -43,13 +43,17 @@ instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 3, 3, l, false) \ instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 4, 4, l, false) -#define instantiate_implicit_2d_blocks(name, itype) \ - instantiate_implicit_2d_filter(name, itype, 32, 8, 16, 4, 1) \ - instantiate_implicit_2d_filter(name, itype, 64, 8, 16, 4, 1) \ - instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \ - instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \ - instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \ - instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2) +// The 8 channel loader reads one row per thread and needs bm = 64 +#define instantiate_implicit_2d_blocks(name, itype) \ + instantiate_implicit_2d_filter(name, itype, 32, 8, 16, 4, 1) \ + instantiate_implicit_2d_filter(name, itype, 64, 8, 16, 4, 1) \ + instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2) \ + instantiate_implicit_conv_2d(name, itype, 64, 8, 16, 4, 1, 8, 8, l, false) \ + instantiate_implicit_conv_2d(name, itype, 64, 32, 16, 2, 2, 8, 8, l, false) \ + instantiate_implicit_conv_2d(name, itype, 64, 64, 16, 2, 2, 8, 8, l, false) instantiate_implicit_2d_blocks(float32, float); instantiate_implicit_2d_blocks(float16, half); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h index 38250f9b8..bb896711c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h @@ -186,7 +186,7 @@ implicit_gemm_conv_2d_general( if (offset_n >= gemm_params->N) return; - short diff = gemm_params->N - offset_n; + int diff = gemm_params->N - offset_n; STEEL_PRAGMA_UNROLL for (int i = 0; i < mma_t::TM; i++) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h index 15dc13ccc..62881c755 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h @@ -47,8 +47,8 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -197,8 +197,8 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -384,6 +384,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -405,6 +406,7 @@ struct Conv2DWeightBlockLoader { weight_hw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -419,7 +421,7 @@ struct Conv2DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -481,9 +483,9 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; short kdil_d; short kdil_h; @@ -665,9 +667,9 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -887,6 +889,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -908,6 +911,7 @@ struct Conv3DWeightBlockLoader { weight_dhw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -922,7 +926,7 @@ struct Conv3DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h index 6a6c16daf..452f344af 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h @@ -71,6 +71,7 @@ struct Conv2DInputBlockLoaderSmallChannels { // Rows / strided reads within the block STEEL_CONST short n_rows = BROWS / TROWS; + static_assert(n_rows > 0, "BM must be at least tgp_size * vec_size / BK"); // Thread location indices const short thread_idx; @@ -233,6 +234,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -253,6 +255,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { params(params_), weight_hw(thread_idx % TCOLS), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -260,7 +263,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { if (bi >= BROWS || bj >= BCOLS) return; - if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -289,7 +292,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } } else { for (short i = 0; i < BROWS; i += TROWS) { - if (((read_n + i) < params->O)) { + if (((read_n + i) < N)) { STEEL_PRAGMA_UNROLL for (short j = 0; j < n_channels; j++) { dst[i * dst_ld + j] = curr_src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h index cf7794235..6bc90e869 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h @@ -48,8 +48,8 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -248,8 +248,8 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const int start_row; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h index 5a43e2233..6814ef50d 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h @@ -79,8 +79,8 @@ template < k_start = segments[tid.z]; k_end = segments[tid.z + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid.z * params->batch_stride_d; // Prepare threadgroup mma operation diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h index 9cfece5c5..5731cb8c1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h @@ -51,8 +51,8 @@ void segmented_mm_nax( k_start = segments[tid_s]; k_end = segments[tid_s + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid_s * params->batch_stride_d; constexpr short SM = BM / WM; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h new file mode 100644 index 000000000..28229a03c --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal new file mode 100644 index 000000000..e3f87e8dc --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal @@ -0,0 +1,27 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "mlx/backend/metal/kernels/utils.h" + +#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" + +// clang-format off +#define instantiate_gemm_thin(tname, trans_b, iname, itype, um, un, sm, sn, ks) \ + instantiate_kernel( \ + "steel_gemm_thin_nax_" #tname "_" #iname \ + "_um" #um "_un" #un "_sm" #sm "_sn" #sn "_ks" #ks, \ + gemm_thin_nax, itype, um, un, sm, sn, ks, trans_b) + +#define instantiate_gemm_thin_transpose_helper(iname, itype, um, un, sm, sn, ks) \ + instantiate_gemm_thin(nn, false, iname, itype, um, un, sm, sn, ks) \ + instantiate_gemm_thin(nt, true , iname, itype, um, un, sm, sn, ks) + +#define instantiate_gemm_thin_shapes_helper(iname, itype) \ + instantiate_gemm_thin_transpose_helper(iname, itype, 2, 2, 1, 1, 4) \ + instantiate_gemm_thin_transpose_helper(iname, itype, 2, 2, 1, 2, 2) + +instantiate_gemm_thin_shapes_helper(float16, half); +instantiate_gemm_thin_shapes_helper(bfloat16, bfloat); +// clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp index 5fc2c1936..3de3ffd7e 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp @@ -7,6 +7,7 @@ #include #include "mlx/backend/common/broadcasting.h" +#include "mlx/backend/common/compiled.h" #include "mlx/backend/common/matmul.h" #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/binary.h" @@ -990,7 +991,62 @@ void steel_matmul_axpby( /* float beta = */ beta); } - // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K + // Case 2: Few output rows or columns, small tiles with A in registers + if (use_nax && batch_size_out == 1 && !transpose_a && + (N <= 64 || (M <= 64 && transpose_b)) && out.dtype() != float32 && + (transpose_b || int64_t(K) * ldb <= INT_MAX)) { + int sn = M < 2048 ? 1 : 2, ks = 4 / sn; + std::ostringstream kname; + kname << "steel_gemm_thin_nax_n" << (transpose_b ? 't' : 'n') << "_" + << type_to_name(out) << "_um2_un2_sm1_sn" << sn << "_ks" << ks; + auto base_name = kname.str(); + bool use_out_source = CHECK_AB && (alpha != 0.0f || beta != 1.0f); + bool do_axpby = use_out_source && (alpha != 1.0f || beta != 1.0f); + bool align_M = M % 32 == 0, align_N = N % (32 * sn) == 0; + bool align_K = K % 64 == 0; + kname << "_" << use_out_source << do_axpby << align_M << align_N << align_K; + auto kernel = get_steel_gemm_thin_nax_kernel( + d, + base_name, + kname.str(), + {{&use_out_source, MTL::DataType::DataTypeBool, 100}, + {&do_axpby, MTL::DataType::DataTypeBool, 110}, + {&align_M, MTL::DataType::DataTypeBool, 200}, + {&align_N, MTL::DataType::DataTypeBool, 201}, + {&align_K, MTL::DataType::DataTypeBool, 202}}, + get_template_definition( + base_name, + "gemm_thin_nax", + get_type_string(out.dtype()), + 2, + 2, + 1, + sn, + ks, + transpose_b)); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + int tm = (M + 31) / 32, tn = (N + 32 * sn - 1) / (32 * sn); + int k_iters = ((K + 63) / 64 + ks - 1) / ks; + steel::GEMMParams params{ + M, N, K, lda, ldb, N, tn, tm, 0, 0, 0, 0, k_iters, 0}; + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(b, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(params, 4); + if (use_out_source) { + int ldc = c.strides()[c.ndim() - 2], fdc = c.strides()[c.ndim() - 1]; + steel::GEMMAddMMParams addmm_params{ldc, fdc, 0, alpha, beta}; + compute_encoder.set_input_array(c, 2); + compute_encoder.set_bytes(addmm_params, 5); + } + compute_encoder.dispatch_threadgroups( + MTL::Size(tm * tn, 1, 1), MTL::Size(32 * sn * ks, 1, 1)); + compute_encoder.add_temporaries(std::move(copies)); + return; + } + + // Case 3: Large K with sufficient M, N, and NAX is available, use NAX split-K if (use_nax && batch_size_out == 1 && (K >= 3 * std::max(M, N) || (std::max(M, N) <= 1024 && K > 2 * std::max(M, N)))) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp index 07e635def..93f8cf311 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp @@ -459,6 +459,15 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string&) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, @@ -527,12 +536,42 @@ MTL::ComputePipelineState* get_gated_delta_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } -MTL::ComputePipelineState* get_gated_delta_nax_kernel( +MTL::ComputePipelineState* get_gated_delta_vjp_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, const metal::MTLFCList& func_consts) { return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int, + int, + int, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gated_delta_vjp_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int, + int, + int, + int, + int, + bool) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp index 97c418e87..052dd761c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp @@ -121,7 +121,7 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } else if (D <= 4096 && O <= 4096) { return 10; } else { - return 6; + return (arch_gen == 13 && arch_size == 's') ? 7 : 6; } } } else { @@ -485,7 +485,8 @@ void qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + bool fast = (K % qmv_fast_k_alignment(bits) == 0); + bool partial_rows = fast && (N % bn != 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 && @@ -503,9 +504,11 @@ void qmv( group_size, "_b_", bits, - use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); + partial_rows ? "_pr_1" : "_pr_0", + global_scale ? "_hgs_1" : "_hgs_0", + "_r_", + results_per_simdgroup); auto kernel = get_quantized_kernel_wrapped( d, kname, @@ -515,6 +518,7 @@ void qmv( group_size, bits, B > 1, + partial_rows, global_scale.has_value(), results_per_simdgroup); @@ -629,11 +633,11 @@ void qmv_wide( } void qvm_split_k( - const array& x, + array x, const array& w, const array& scales, const std::optional& biases, - array& out, + array out, int group_size, int bits, int M, @@ -647,6 +651,16 @@ void qvm_split_k( int split_k = K > 8192 ? 32 : 8; int split_D = (K + split_k - 1) / split_k; int B = out.size() / M / N; + + // When weights are 2D, collapse leading batch dimensions of x and out into M + // so the kernel and reduction operate on standard 2D matrices. + if (B == 1 && x.ndim() != 2) { + x = reshape_in_eval(x, {M, K}, s); + } + if (B == 1 && out.ndim() != 2) { + out = reshape_in_eval(out, {M, N}, s); + } + B *= split_k; constexpr int num_simdgroups = 2; @@ -1203,15 +1217,11 @@ void qmm_splitk( int k_partition_size = K / split_k; int split_k_partition_stride = M * N; - // Allocate intermediate buffer: insert split_k at the front so that - // partition_stride = M * N matches the leading stride of the buffer. + // Allocate a float32 buffer for the partial sums, so that each partition is + // rounded to the output type only once, after they're summed, as in + // steel_matmul_splitk. auto& compute_encoder = metal::get_command_encoder(s); - auto temp_shape = out.shape(); - if (temp_shape.size() == 1) { - temp_shape.insert(temp_shape.begin(), 1); - } - temp_shape.insert(temp_shape.begin(), split_k); - array intermediate(temp_shape, x.dtype(), nullptr, {}); + array intermediate({split_k, M, N}, float32, nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); compute_encoder.add_temporary(intermediate); @@ -1253,13 +1263,18 @@ void qmm_splitk( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - // Sum across split_k dimension (axis 0) - ReductionPlan plan( - ReductionOpType::ContiguousStridedReduce, - {intermediate.shape(0)}, - {intermediate.strides(0)}); - strided_reduce_general_dispatch( - intermediate, out, "sum", plan, {0}, compute_encoder, d, s); + // Sum the partitions and cast to the output type + auto accum_kname = "steel_gemm_splitk_accum_" + type_to_name(out) + "_" + + type_to_name(intermediate); + auto accum_kernel = get_steel_gemm_splitk_accum_kernel( + d, accum_kname, intermediate, out, false); + compute_encoder.set_compute_pipeline_state(accum_kernel); + compute_encoder.set_input_array(intermediate, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(split_k, 2); + compute_encoder.set_bytes(split_k_partition_stride, 3); + compute_encoder.set_bytes(N, 4); + compute_encoder.dispatch_threads(MTL::Size(N, M, 1), get_block_dims(N, M, 1)); } void gather_qmm( @@ -1675,7 +1690,7 @@ void gather_qmm_rhs( metal::Device& d, const Stream& s, const std::string mode) { - if (metal::is_nax_available() && transpose && + if (metal::is_nax_available() && (env::enable_tf32() || x_.dtype() != float32)) { return gather_qmm_rhs_nax( /* const array& x_ = */ x_, @@ -1838,7 +1853,8 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { + if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && + !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp b/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp index ac11ac635..5d9f6d43c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp @@ -242,6 +242,27 @@ inline auto output_grid_for_col_reduce( return get_2d_grid_dims(out_shape, out_strides); } +// Metal silently truncates a grid dimension past UINT32_MAX threads, so a +// column reduction too wide for one dispatch is split into several, each +// covering the BN-wide column blocks starting at column_offset. +inline void dispatch_col_reduce( + CommandEncoder& compute_encoder, + int64_t n_columns, + int BN, + int threadgroup_size, + size_t grid_y, + size_t grid_z) { + const int64_t n_blocks = (n_columns + BN - 1) / BN; + const int64_t max_blocks = UINT32_MAX / threadgroup_size; + const MTL::Size group_dims(threadgroup_size, 1, 1); + for (int64_t start = 0; start < n_blocks; start += max_blocks) { + const int64_t blocks = std::min(max_blocks, n_blocks - start); + compute_encoder.set_bytes(start * BN, 11); + compute_encoder.dispatch_threads( + MTL::Size(threadgroup_size * blocks, grid_y, grid_z), group_dims); + } +} + std::pair remap_reduce_types( const array& in, const std::string& op_name) { @@ -275,9 +296,7 @@ std::pair remap_reduce_types( } 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. + // bit set, so it has to be tested as a float. switch (in.dtype()) { case float16: return {float16, bool_}; @@ -285,6 +304,8 @@ std::pair remap_reduce_types( return {bfloat16, bool_}; case float32: return {float32, bool_}; + case complex64: + return {complex64, bool_}; default: break; } @@ -315,11 +336,10 @@ void init_reduce( std::string kname = func_name; concatenate(kname, "_", op_name, type_to_name(out_type)); auto kernel = get_reduce_init_kernel(d, kname, func_name, op_name, out_type); - size_t nthreads = out.size(); - MTL::Size grid_dims = MTL::Size(nthreads, 1, 1); + MTL::Size grid_dims = get_2d_grid_dims(out.shape(), out.strides()); NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - if (thread_group_size > nthreads) { - thread_group_size = nthreads; + if (thread_group_size > grid_dims.width) { + thread_group_size = grid_dims.width; } MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1); compute_encoder.set_compute_pipeline_state(kernel); @@ -421,7 +441,7 @@ void row_reduce_small( auto [in_type, out_type] = remap_reduce_types(in, op_name); const std::string func_name = "row_reduce_small"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -518,7 +538,7 @@ void row_reduce_looped( int n = get_kernel_reduce_ndim(args.reduce_ndim); const std::string func_name = "row_reduce_looped"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -602,7 +622,7 @@ void strided_reduce_small( int n = get_kernel_reduce_ndim(args.reduce_ndim); const std::string func_name = "col_reduce_small"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -693,7 +713,7 @@ void strided_reduce_longcolumn( int n = get_kernel_reduce_ndim(args.reduce_ndim); std::string func_name = "col_reduce_longcolumn"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -728,8 +748,7 @@ void strided_reduce_longcolumn( second_args.reduce_strides.push_back(out.size()); second_args.reduce_ndim++; int BN = 32; - grid_dims = MTL::Size(256 * ((out.size() + BN - 1) / BN), 1, 1); - group_dims = MTL::Size(256, 1, 1); + int threadgroup_size = 256; // Set the 2nd kernel func_name = "col_reduce_looped"; @@ -755,7 +774,7 @@ void strided_reduce_longcolumn( compute_encoder.set_input_array(intermediate, 0); compute_encoder.set_output_array(out, 1); second_args.encode(compute_encoder); - compute_encoder.dispatch_threads(grid_dims, group_dims); + dispatch_col_reduce(compute_encoder, out.size(), BN, threadgroup_size, 1, 1); } void strided_reduce_looped( @@ -778,17 +797,12 @@ void strided_reduce_looped( int BN = 32; int BM = 1024 / BN; int threadgroup_size = 8 * 32; - MTL::Size grid_dims( - threadgroup_size * ((args.reduction_stride + BN - 1) / BN), - out_grid_size.width, - out_grid_size.height); - MTL::Size group_dims(threadgroup_size, 1, 1); // Set the kernel int n = get_kernel_reduce_ndim(args.reduce_ndim); std::string func_name = "col_reduce_looped"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -820,7 +834,13 @@ void strided_reduce_looped( compute_encoder.set_input_array(in, 0); compute_encoder.set_output_array(out, 1); args.encode(compute_encoder); - compute_encoder.dispatch_threads(grid_dims, group_dims); + dispatch_col_reduce( + compute_encoder, + args.reduction_stride, + BN, + threadgroup_size, + out_grid_size.width, + out_grid_size.height); } void strided_reduce_2pass( @@ -855,17 +875,12 @@ void strided_reduce_2pass( int BN = 32; int BM = 1024 / BN; int threadgroup_size = 8 * 32; - MTL::Size grid_dims( - threadgroup_size * ((args.reduction_stride + BN - 1) / BN), - out_grid_size.width * outer_blocks, - out_grid_size.height); - MTL::Size group_dims(threadgroup_size, 1, 1); // Set the kernel int n = get_kernel_reduce_ndim(args.reduce_ndim); std::string func_name = "col_reduce_2pass"; std::string kname = func_name; - bool large = in.size() > INT32_MAX; + bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; if (large) { kname += "_large"; } @@ -897,15 +912,20 @@ void strided_reduce_2pass( compute_encoder.set_input_array(in, 0); compute_encoder.set_output_array(intermediate, 1); args.encode(compute_encoder); - compute_encoder.set_bytes(out_size, 11); - compute_encoder.dispatch_threads(grid_dims, group_dims); + compute_encoder.set_bytes(out_size, 12); + dispatch_col_reduce( + compute_encoder, + args.reduction_stride, + BN, + threadgroup_size, + out_grid_size.width * outer_blocks, + out_grid_size.height); // Make the 2nd pass arguments and grid_dims ColReduceArgs second_args(intermediate); second_args.reduce_shape.push_back(outer_blocks); second_args.reduce_strides.push_back(out.size()); second_args.reduce_ndim++; - grid_dims = MTL::Size(threadgroup_size * ((out.size() + BN - 1) / BN), 1, 1); // Set the 2nd kernel func_name = "col_reduce_looped"; @@ -931,7 +951,7 @@ void strided_reduce_2pass( compute_encoder.set_input_array(intermediate, 0); compute_encoder.set_output_array(out, 1); second_args.encode(compute_encoder); - compute_encoder.dispatch_threads(grid_dims, group_dims); + dispatch_col_reduce(compute_encoder, out.size(), BN, threadgroup_size, 1, 1); } void strided_reduce_general_dispatch( diff --git a/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp b/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp index 243699363..26ab718e3 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp @@ -1,10 +1,12 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #include #include +#include #include #include #include +#include #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/resident.h" @@ -33,21 +35,49 @@ ResidencySets::ResidencySets(MTL::Device* d) { 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; + std::string error; if (!add_set_locked(&error)) { std::ostringstream msg; msg << "[metal::Device] Unable to construct residency set.\n"; - if (error) { - msg << error->localizedDescription()->utf8String() << "\n"; + if (!error.empty()) { + msg << error << "\n"; } throw std::runtime_error(msg.str()); } + if (int interval = + env::get_var("MLX_METAL_RESIDENCY_REFRESH_INTERVAL_MS", 0); + interval > 0) { + refresh_thread_ = std::thread([this, interval] { + std::unique_lock lock(mtx_); + while (!refresh_cv_.wait_for( + lock, std::chrono::milliseconds(interval), [this] { + return stop_refresh_; + })) { + auto pool = new_scoped_memory_pool(); + // macOS can drop residency after GPU idle despite a standing request. + for (auto& s : sets_) { + if (s.size != 0) { + s.set->requestResidency(); + } + } + } + }); + } } } -ResidencySets::~ResidencySets() = default; +ResidencySets::~ResidencySets() { + if (refresh_thread_.joinable()) { + { + std::lock_guard lock(mtx_); + stop_refresh_ = true; + } + refresh_cv_.notify_one(); + refresh_thread_.join(); + } +} -bool ResidencySets::add_set_locked(NS::Error** error_out) { +bool ResidencySets::add_set_locked(std::string* error_out) { NS::SharedPtr set; if (__builtin_available(macOS 15, iOS 18, *)) { auto pool = new_scoped_memory_pool(); @@ -58,8 +88,9 @@ bool ResidencySets::add_set_locked(NS::Error** error_out) { // 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; + } else if (error_out && error) { + // The error is autoreleased into this pool, so copy its text out now. + *error_out = error->localizedDescription()->utf8String(); } } if (!set) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/resident.h b/Source/Cmlx/mlx/mlx/backend/metal/resident.h index 23960ad13..3339307b6 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/resident.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/resident.h @@ -1,10 +1,13 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #pragma once #include +#include #include #include +#include +#include #include #include @@ -96,7 +99,7 @@ class ResidencySets { // 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); + bool add_set_locked(std::string* 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); @@ -118,6 +121,9 @@ class ResidencySets { // taking the lock. std::atomic num_sets_{0}; mutable std::mutex mtx_; + std::condition_variable refresh_cv_; + bool stop_refresh_{false}; + std::thread refresh_thread_; }; } // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp b/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp index 3f5ed8336..5b157c6c5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -1,4 +1,5 @@ // Copyright © 2024-26 Apple Inc. +#include #include #include "mlx/backend/common/compiled.h" @@ -7,6 +8,7 @@ #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/kernels/steel/attn/params.h" +#include "mlx/backend/metal/matmul.h" #include "mlx/backend/metal/utils.h" #include "mlx/fast_primitives.h" #include "mlx/utils.h" @@ -15,6 +17,190 @@ namespace mlx::core::fast { namespace { +bool use_blocked_d512( + const array& q, + const array& k, + const array& v, + bool has_arr_mask, + bool do_causal, + bool has_sinks, + bool output_logsumexp, + Stream s) { + return s.device == Device::gpu && !metal::is_nax_available() && + (q.dtype() == float16 || q.dtype() == bfloat16) && q.shape(-1) == 512 && + v.shape(-1) == 512 && q.shape(2) >= 1024 && q.shape(2) <= k.shape(2) && + do_causal && !has_arr_mask && !has_sinks && !output_logsumexp; +} + +void sdpa_blocked_softmax(const Stream& s, metal::Device& d, array& scores) { + constexpr int looped_limit = 4096; + constexpr int simd_size = 32; + + int axis_size = scores.shape().back(); + int n_rows = scores.size() / axis_size; + std::string kernel_name = + axis_size > looped_limit ? "looped_softmax_" : "block_softmax_"; + if (scores.dtype() != float32) { + kernel_name += "precise_"; + } + kernel_name += type_to_name(scores); + + auto kernel = get_softmax_kernel(d, kernel_name, true, scores); + auto& compute_encoder = metal::get_command_encoder(s); + + size_t threadgroup_size; + if (axis_size <= looped_limit) { + size_t reads = (axis_size + SOFTMAX_N_READS - 1) / SOFTMAX_N_READS; + size_t simds = (reads + simd_size - 1) / simd_size; + threadgroup_size = simd_size * simds; + assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup()); + } else { + threadgroup_size = kernel->maxTotalThreadsPerThreadgroup(); + } + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(scores, 0); + compute_encoder.set_output_array(scores, 1); + compute_encoder.set_bytes(axis_size, 2); + compute_encoder.dispatch_threads( + MTL::Size(n_rows * threadgroup_size, 1, 1), + MTL::Size(threadgroup_size, 1, 1)); +} + +void sdpa_full_self_attention_blocked( + const Stream& s, + metal::Device& d, + const array& q, + const array& k, + const array& v, + float scale, + array& o) { + // This block size balances score-buffer memory and GEMM efficiency. + constexpr int block_rows = 256; + + int B = q.shape(0); + int H = q.shape(1); + int qL = q.shape(2); + int kL = k.shape(2); + int D = q.shape(3); + int kvH = k.shape(1); + int gqa = H / kvH; + + array score_buffer({B, H, block_rows, kL}, q.dtype(), nullptr, {}); + score_buffer.set_data(allocator::malloc(score_buffer.nbytes())); + array q_buffer({B, H, block_rows, D}, q.dtype(), nullptr, {}); + q_buffer.set_data(allocator::malloc(q_buffer.nbytes())); + array out_buffer({B, H, block_rows, D}, o.dtype(), nullptr, {}); + out_buffer.set_data(allocator::malloc(out_buffer.nbytes())); + + auto& compute_encoder = metal::get_command_encoder(s); + auto scale_kernel = + d.get_kernel("sdpa_blocked_scale_copy_" + type_to_name(q)); + auto mask_kernel = + d.get_kernel("sdpa_blocked_causal_mask_" + type_to_name(q)); + Shape batch_shape = {B, kvH}; + Strides k_batch_strides = {k.strides(0), k.strides(1)}; + Strides v_batch_strides = {v.strides(0), v.strides(1)}; + + std::vector copies; + for (int q_offset = 0; q_offset < qL; q_offset += block_rows) { + int rows = std::min(block_rows, qL - q_offset); + // Skip columns that are masked for every row in this query block. + int columns = std::min(kL, kL - qL + q_offset + rows); + + array q_block({B, H, rows, D}, q.dtype(), nullptr, {}); + q_block.copy_shared_buffer( + q_buffer, q_block.strides(), q_block.flags(), q_block.size()); + compute_encoder.set_compute_pipeline_state(scale_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_output_array(q_block, 1); + compute_encoder.set_bytes(scale, 2); + compute_encoder.set_vector_bytes(q.strides(), 3); + compute_encoder.set_bytes(H, 4); + compute_encoder.set_bytes(rows, 5); + compute_encoder.set_bytes(D, 6); + compute_encoder.set_bytes(q_offset, 7); + auto scale_group_dims = get_block_dims(D, rows, B * H); + compute_encoder.dispatch_threads( + MTL::Size(D, rows, B * H), scale_group_dims); + + array scores({B, H, rows, columns}, q.dtype(), nullptr, {}); + scores.copy_shared_buffer( + score_buffer, scores.strides(), scores.flags(), scores.size()); + + Strides q_batch_strides = {int64_t(H) * rows * D, int64_t(gqa) * rows * D}; + steel_matmul( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ q_block, + /* const array& b = */ k, + /* array& out = */ scores, + /* int M = */ gqa * rows, + /* int N = */ columns, + /* int K = */ D, + /* int batch_size_out = */ B * kvH, + /* int lda = */ D, + /* int ldb = */ k.strides(2), + /* bool transpose_a = */ false, + /* bool transpose_b = */ true, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ batch_shape, + /* Strides A_batch_stride = */ q_batch_strides, + /* Strides B_batch_stride = */ k_batch_strides); + + int total_rows = B * H * rows; + compute_encoder.set_compute_pipeline_state(mask_kernel); + compute_encoder.set_output_array(scores, 0); + compute_encoder.set_bytes(rows, 1); + compute_encoder.set_bytes(columns, 2); + auto mask_group_dims = get_block_dims(rows, total_rows, 1); + compute_encoder.dispatch_threads( + MTL::Size(rows, total_rows, 1), mask_group_dims); + + sdpa_blocked_softmax(s, d, scores); + + array out_block({B, H, rows, D}, o.dtype(), nullptr, {}); + out_block.copy_shared_buffer( + out_buffer, out_block.strides(), out_block.flags(), out_block.size()); + + Strides scores_batch_strides = { + int64_t(H) * rows * columns, int64_t(gqa) * rows * columns}; + steel_matmul( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ scores, + /* const array& b = */ v, + /* array& out = */ out_block, + /* int M = */ gqa * rows, + /* int N = */ D, + /* int K = */ columns, + /* int batch_size_out = */ B * kvH, + /* int lda = */ columns, + /* int ldb = */ v.strides(2), + /* bool transpose_a = */ false, + /* bool transpose_b = */ false, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ batch_shape, + /* Strides A_batch_stride = */ scores_batch_strides, + /* Strides B_batch_stride = */ v_batch_strides); + copy_gpu_inplace( + /* const array& in = */ out_block, + /* array& out = */ o, + /* const Shape& data_shape = */ out_block.shape(), + /* const Strides& i_strides = */ out_block.strides(), + /* const Strides& o_strides = */ o.strides(), + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ int64_t(q_offset) * o.strides(2), + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ s); + } + + compute_encoder.add_temporaries(std::move(copies)); + compute_encoder.add_temporary(std::move(score_buffer)); + compute_encoder.add_temporary(std::move(q_buffer)); + compute_encoder.add_temporary(std::move(out_buffer)); +} + void sdpa_full_self_attention_nax( const Stream& s, metal::Device& d, @@ -25,7 +211,8 @@ void sdpa_full_self_attention_nax( array& o, bool do_causal_, const std::optional& mask, - const std::optional& sinks) { + const std::optional& sinks, + array* lse) { using namespace mlx::steel; int bd = q.shape(-1); @@ -33,7 +220,7 @@ void sdpa_full_self_attention_nax( int bq = bd == 512 ? 32 : 64; int bk = 32; - bool split_d = bd == 256 || bd == 512; + bool split_d = (bd == 256 || bd == 512); int wm = bd == 512 ? 2 : 4; int wn = split_d ? bd / 128 : 1; int B = q.shape(0); @@ -79,13 +266,15 @@ void sdpa_full_self_attention_nax( const bool has_mask = mask.has_value(); const bool do_causal = do_causal_; const bool has_sinks = sinks.has_value(); + const bool save_lse = (lse != nullptr); metal::MTLFCList func_consts = { {&align_Q, MTL::DataType::DataTypeBool, 200}, {&align_K, MTL::DataType::DataTypeBool, 201}, {&has_mask, MTL::DataType::DataTypeBool, 300}, {&do_causal, MTL::DataType::DataTypeBool, 301}, - {&has_sinks, MTL::DataType::DataTypeBool, 302}}; + {&has_sinks, MTL::DataType::DataTypeBool, 302}, + {&save_lse, MTL::DataType::DataTypeBool, 303}}; std::string base_name; concatenate( @@ -120,7 +309,9 @@ void sdpa_full_self_attention_nax( "_do_causal_", (do_causal ? 't' : 'n'), "_has_sinks_", - (has_sinks ? 't' : 'n')); + (has_sinks ? 't' : 'n'), + "_save_lse_", + (save_lse ? 't' : 'n')); auto& compute_encoder = metal::get_command_encoder(s); @@ -191,6 +382,9 @@ void sdpa_full_self_attention_nax( if (has_sinks) { compute_encoder.set_input_array(*sinks, 7); } + if (save_lse) { + compute_encoder.set_output_array(*lse, 8); + } MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); @@ -209,15 +403,12 @@ void sdpa_full_self_attention_metal( array& o, bool do_causal_, const std::optional& mask, - const std::optional& sinks) { - int B = q.shape(0); - int H = q.shape(1); + const std::optional& sinks, + array* lse /* = nullptr */) { int D = q.shape(3); - int gqa_factor = q.shape(1) / k.shape(1); int qL = q.shape(2); int kL = k.shape(2); - if (metal::is_nax_available() && (D == 64 || D == 96 || D == 128 || D == 256 || D == 512) && (env::enable_tf32() || q.dtype() != float32)) { @@ -231,15 +422,18 @@ void sdpa_full_self_attention_metal( /* array& o = */ o, /* bool do_causal_ = */ do_causal_, /* const std::optional& mask = */ mask, - /* const std::optional& sinks = */ sinks); + /* const std::optional& sinks = */ sinks, + /* array* lse = */ lse); } // Pad head dims 72 and 80 to 96 to reach the NAX kernel. The added lanes // are zero and the caller's scale is retained. Enable by default only for // long, unmasked half-precision attention, where the attention work can // amortize the padding and output copy. The override is read per call. + // + // Skipped when the logsumexp is needed bool pad_default = qL >= 512 && kL >= 512 && !do_causal_ && !mask && !sinks; - if ((D == 72 || D == 80) && metal::is_nax_available() && + if ((D == 72 || D == 80) && lse == nullptr && metal::is_nax_available() && (q.dtype() == float16 || q.dtype() == bfloat16) && env::get_var("MLX_SDPA_PAD_HEAD_DIM", pad_default ? 1 : 0) == 1) { constexpr int pad_to = 96; @@ -286,7 +480,8 @@ void sdpa_full_self_attention_metal( /* array& o = */ op, /* bool do_causal_ = */ do_causal_, /* const std::optional& mask = */ mask, - /* const std::optional& sinks = */ sinks); + /* const std::optional& sinks = */ sinks, + /* array* lse = */ nullptr); // Slice the padded lanes back out into o's caller-chosen strides. copy_gpu_inplace( @@ -306,6 +501,9 @@ void sdpa_full_self_attention_metal( using namespace mlx::steel; + int B = q.shape(0); + int H = q.shape(1); + int gqa_factor = q.shape(1) / k.shape(1); char devc = d.get_architecture().back(); int bd = q.shape(-1); int bv = v.shape(-1); @@ -322,13 +520,15 @@ void sdpa_full_self_attention_metal( const bool has_mask = mask.has_value(); const bool do_causal = do_causal_; const bool has_sinks = sinks.has_value(); + const bool save_lse = (lse != nullptr); metal::MTLFCList func_consts = { {&align_Q, MTL::DataType::DataTypeBool, 200}, {&align_K, MTL::DataType::DataTypeBool, 201}, {&has_mask, MTL::DataType::DataTypeBool, 300}, {&do_causal, MTL::DataType::DataTypeBool, 301}, - {&has_sinks, MTL::DataType::DataTypeBool, 302}}; + {&has_sinks, MTL::DataType::DataTypeBool, 302}, + {&save_lse, MTL::DataType::DataTypeBool, 303}}; std::string base_name; concatenate( @@ -363,7 +563,9 @@ void sdpa_full_self_attention_metal( "_do_causal_", (do_causal ? 't' : 'n'), "_has_sinks_", - (has_sinks ? 't' : 'n')); + (has_sinks ? 't' : 'n'), + "_save_lse_", + (save_lse ? 't' : 'n')); auto& compute_encoder = metal::get_command_encoder(s); @@ -433,6 +635,9 @@ void sdpa_full_self_attention_metal( if (has_sinks) { compute_encoder.set_input_array(*sinks, 7); } + if (save_lse) { + compute_encoder.set_output_array(*lse, 8); + } MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); @@ -654,6 +859,18 @@ void sdpa_vector_2pass( // Get the kernel auto kernel = d.get_kernel(kname, hash_name, func_consts); + // M1/M2 do not have enough threads per threadgroup for default TK=4 kernel + bool has_tk1_variant = + (q.dtype() == bfloat16 && q.shape(-1) == 256 && v.shape(-1) == 256) || + (q.dtype() == float32 && q.shape(-1) == v.shape(-1) && + (q.shape(-1) == 192 || q.shape(-1) == 256)); + if (has_tk1_variant && + group_dims.width * group_dims.height * group_dims.depth > + kernel->maxTotalThreadsPerThreadgroup()) { + kname += "_tk1"; + hash_name += "_tk1"; + 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); @@ -726,11 +943,12 @@ std::tuple has_fused_kernel( if (s.device != Device::gpu) { return {false, "the fused kernels require a GPU (Metal) stream."}; } - if (output_logsumexp) { + + if (output_logsumexp && has_arr_mask) { return { false, - "the fused forward does not produce the logsumexp required for " - "the fused VJP; use default routing when training."}; + "the backward pass does not support an array mask; only causal masking " + "is implemented."}; } const int value_head_dim = v.shape(-1); @@ -832,6 +1050,341 @@ std::tuple has_fused_kernel( return {true, ""}; } +/////////////////////////////////////////////////////////////////////////////// +// Backward pass +/////////////////////////////////////////////////////////////////////////////// + +struct SDPAVJPTileParams { + int bq; + int bk; + int qL; + int i0; + int j0; + float scale; + int diag_off; + int causal; +}; + +array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { + if (!x.flags().row_contiguous) { + array x_copy = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(x_copy); + return x_copy; + } else { + return x; + } +} + +array vjp_alloc(Shape shape, Dtype dt, const Stream& s) { + array a(std::move(shape), dt, nullptr, {}); + a.set_data(allocator::malloc(a.nbytes())); + metal::get_command_encoder(s).add_temporary(a); + return a; +} + +// Takes a view of the buffer +array vjp_view(const array& base, Shape shape) { + array v(shape, base.dtype(), nullptr, {}); + Strides st(shape.size()); + int64_t acc = 1; + for (int i = static_cast(shape.size()) - 1; i >= 0; --i) { + st[i] = acc; + acc *= shape[i]; + } + array::Flags f{1, 1, shape.size() <= 1}; + v.copy_shared_buffer(base, st, f, static_cast(acc), 0); + return v; +} + +// A [B, H, len, D] window of a row-contiguous [B, H, T, D] input, starting at +// sequence position t0. +array vjp_row_slice(const array& x, int t0, int len) { + const auto& shp = x.shape(); + int64_t Dx = shp[3]; + int64_t Tx = shp[2]; + Shape ns = {shp[0], shp[1], len, static_cast(Dx)}; + Strides st = {shp[1] * Tx * Dx, Tx * Dx, Dx, 1}; + array v(ns, x.dtype(), nullptr, {}); + array::Flags f{1, len == shp[2], false}; + size_t base = static_cast(x.offset()) / x.itemsize(); + v.copy_shared_buffer( + x, + st, + f, + x.data_size(), + base + static_cast(t0) * static_cast(Dx)); + return v; +} + +// Block sizes for the score tile. The best results are usually for 1024x1024 so +// I set those as default. This probably depends on the hardware so it may need +// to be tuned in the future. +std::pair sdpa_vjp_blocks(int B, int H, int qL, int kL, Dtype ctype) { + constexpr int kDefaultBlock = 1024; + + int blk = env::get_var("MLX_SDPA_VJP_BLOCK", kDefaultBlock); + + int bq = std::min(blk, qL); + int bk = std::min(blk, kL); + if (int e = env::get_var("MLX_SDPA_VJP_BQ", 0); e > 0) { + bq = std::min(e, qL); + } + if (int e = env::get_var("MLX_SDPA_VJP_BK", 0); e > 0) { + bk = std::min(e, kL); + } + return {bq, bk}; +} + +void sdpa_vjp_blocked( + const Stream& s, + metal::Device& d, + const array& q, + const array& k, + const array& v, + const array& odo, + const array& lse, + const array& cot_o, + array& dq, + array& dk, + array& dv, + float scale, + bool causal) { + auto& compute_encoder = metal::get_command_encoder(s); + + const int B = q.shape(0); + const int H = q.shape(1); + const int qL = q.shape(2); + const int D = q.shape(3); + const int Hk = k.shape(1); + const int kL = k.shape(2); + const int Dv = v.shape(3); + const int G = H / Hk; + const int BH = B * H; + const int BHk = B * Hk; + // Under a causal mask the diagonal sits at j == i + (kL - qL), which is + // nonzero whenever the queries are a suffix of the keys. + const int diag_off = kL - qL; + + auto [BQ, BK] = sdpa_vjp_blocks(B, H, qL, kL, q.dtype()); + + Dtype ctype = q.dtype(); + std::string tname = get_type_string(ctype); + + auto ds_kernel = d.get_kernel("sdpa_vjp_ds_" + tname); + auto red_add = d.get_kernel("sdpa_vjp_reduce_add_" + tname); + auto red_set = d.get_kernel("sdpa_vjp_reduce_set_" + tname); + + // Temporary tiles + array s_buf = vjp_alloc({BH, BQ, BK}, ctype, s); + array dp_buf = vjp_alloc({BH, BQ, BK}, ctype, s); + array p_buf = vjp_alloc({BH, BQ, BK}, ctype, s); + + // Output of the gemms + array dq_tile = vjp_alloc({BH, BQ, D}, ctype, s); + array dk_tile = vjp_alloc({BH, BK, D}, ctype, s); + array dv_tile = vjp_alloc({BH, BK, Dv}, ctype, s); + + // Accumulators for the result + array dq_acc = vjp_alloc({B, H, qL, D}, float32, s); + array dk_acc = vjp_alloc({B, Hk, kL, D}, float32, s); + array dv_acc = vjp_alloc({B, Hk, kL, Dv}, float32, s); + + std::vector copies; + + auto run_reduce = [&](const array& src, + array& acc, + array& out, + int rows, + int dim, + int group, + int acc_rows, + int row_off, + int nbh, + bool accum) { + compute_encoder.set_compute_pipeline_state(accum ? red_add : red_set); + compute_encoder.set_input_array(src, 0); + compute_encoder.set_input_array(acc, 1); + compute_encoder.set_output_array(acc, 1); + compute_encoder.set_output_array(out, 2); + compute_encoder.set_bytes(rows, 3); + compute_encoder.set_bytes(dim, 4); + compute_encoder.set_bytes(group, 5); + compute_encoder.set_bytes(acc_rows, 6); + compute_encoder.set_bytes(row_off, 7); + compute_encoder.dispatch_threads( + MTL::Size(dim, rows, nbh), MTL::Size(std::min(dim, 32), 8, 1)); + }; + + const int n_i = (qL + BQ - 1) / BQ; + + const Strides q_bs = { + H * int64_t(qL) * D, G * int64_t(qL) * D, int64_t(qL) * D}; + const Strides o_bs = { + H * int64_t(qL) * Dv, G * int64_t(qL) * Dv, int64_t(qL) * Dv}; + // A zero innermost batch stride broadcasts one kv head over its query group. + const Strides k_bs = {Hk * int64_t(kL) * D, int64_t(kL) * D, 0}; + const Strides v_bs = {Hk * int64_t(kL) * Dv, int64_t(kL) * Dv, 0}; + const Shape bshape = {B, Hk, G}; + + for (int j0 = 0; j0 < kL; j0 += BK) { + const int bk_len = std::min(BK, kL - j0); + + array k_sl = vjp_row_slice(k, j0, bk_len); + array v_sl = vjp_row_slice(v, j0, bk_len); + + bool kv_accum = false; + + for (int ii = 0; ii < n_i; ++ii) { + const int i0 = ii * BQ; + const int bq_len = std::min(BQ, qL - i0); + // Skip tiles that lie entirely above the causal diagonal. + if (causal && j0 > i0 + bq_len - 1 + diag_off) { + continue; + } + + array q_sl = vjp_row_slice(q, i0, bq_len); + array o_sl = vjp_row_slice(cot_o, i0, bq_len); + + array s_v = vjp_view(s_buf, {BH, bq_len, bk_len}); + // S = Q @ K.T + steel_matmul( + s, + d, + q_sl, + k_sl, + s_v, + bq_len, + bk_len, + D, + BH, + D, + D, + false, + true, + copies, + bshape, + q_bs, + k_bs); + + array dp_v = vjp_view(dp_buf, {BH, bq_len, bk_len}); + // dP = dO @ V.T + steel_matmul( + s, + d, + o_sl, + v_sl, + dp_v, + bq_len, + bk_len, + Dv, + BH, + Dv, + Dv, + false, + true, + copies, + bshape, + o_bs, + v_bs); + + array ds_v = vjp_view(s_buf, {BH, bq_len, bk_len}); + array p_v = vjp_view(p_buf, {BH, bq_len, bk_len}); + + SDPAVJPTileParams params{ + bq_len, bk_len, qL, i0, j0, scale, diag_off, causal ? 1 : 0}; + + // P = exp(scale * S - lse), dS = P * (dP - delta) * scale + compute_encoder.set_compute_pipeline_state(ds_kernel); + compute_encoder.set_input_array(s_v, 0); + compute_encoder.set_input_array(dp_v, 1); + compute_encoder.set_input_array(lse, 2); + compute_encoder.set_input_array(odo, 3); + compute_encoder.set_output_array(s_v, 4); + compute_encoder.set_output_array(p_v, 5); + compute_encoder.set_bytes(params, 6); + compute_encoder.dispatch_threads( + MTL::Size(bk_len, bq_len, BH), MTL::Size(32, 8, 1)); + + int64_t ss2 = static_cast(bq_len) * bk_len; + Strides sc_bs = {H * ss2, G * ss2, ss2}; + + array dq_v = vjp_view(dq_tile, {BH, bq_len, D}); + // dQ = dS @ K + steel_matmul( + s, + d, + ds_v, + k_sl, + dq_v, + bq_len, + D, + bk_len, + BH, + bk_len, + D, + false, + false, + copies, + bshape, + sc_bs, + k_bs); + + array dk_v = vjp_view(dk_tile, {BH, bk_len, D}); + // dK = dS.T @ Q + steel_matmul( + s, + d, + ds_v, + q_sl, + dk_v, + bk_len, + D, + bq_len, + BH, + bk_len, + D, + true, + false, + copies, + bshape, + sc_bs, + q_bs); + + array dv_v = vjp_view(dv_tile, {BH, bk_len, Dv}); + // dV = P.T @ dO + steel_matmul( + s, + d, + p_v, + o_sl, + dv_v, + bk_len, + Dv, + bq_len, + BH, + bk_len, + Dv, + true, + false, + copies, + bshape, + sc_bs, + o_bs); + + // dQ[i0] += dQ_tile, dK[j0] += sum_G dK_tile, dV[j0] += sum_G dV_tile + run_reduce(dq_v, dq_acc, dq, bq_len, D, 1, qL, i0, BH, j0 != 0); + run_reduce(dk_v, dk_acc, dk, bk_len, D, G, kL, j0, BHk, kv_accum); + run_reduce(dv_v, dv_acc, dv, bk_len, Dv, G, kL, j0, BHk, kv_accum); + + kv_accum = true; + } + } + + for (auto& c : copies) { + compute_encoder.add_temporary(c); + } +} + } // namespace bool ScaledDotProductAttention::use_fallback( @@ -841,6 +1394,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, @@ -857,16 +1411,19 @@ bool ScaledDotProductAttention::use_fallback( } return false; } - - if (is_training) { - // It's faster for training on Metal to use the unfused SDPA for both - // forward and backward. - return true; + if (use_blocked_d512( + q, k, v, has_arr_mask, do_causal, has_sinks, output_logsumexp, s)) { + return false; } if (!has_fused) { return true; } + // The logsumexp only comes out of the fused forward + if (output_logsumexp) { + return false; + } + const int query_sequence_length = q.shape(2); const int query_head_dim = q.shape(-1); const int value_head_dim = v.shape(-1); @@ -922,7 +1479,6 @@ void ScaledDotProductAttention::eval_gpu( auto& o = outputs[0]; std::vector copies; - // Define some copy functions to ensure the layout of the inputs is as // expected. copies.reserve(inputs.size()); @@ -950,6 +1506,12 @@ void ScaledDotProductAttention::eval_gpu( // We are in vector mode ie single query if (q_pre.shape(2) <= 8) { + if (outputs.size() > 1) { + throw std::runtime_error( + "[scaled_dot_product_attention] the vector kernels do not produce a " + "logsumexp; the VJP path requires the full-attention kernel."); + } + auto q_copy_unless = [](const array& arr) { if (arr.flags().row_contiguous) { return true; @@ -1048,21 +1610,103 @@ void ScaledDotProductAttention::eval_gpu( ? std::optional{copy_unless(is_matrix_contiguous, inputs[3])} : std::nullopt; - sdpa_full_self_attention_metal( - s, d, q, k, v, scale_, o, do_causal_, mask, sinks); + array* lse = nullptr; + if (outputs.size() > 1) { + outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); + lse = &outputs[1]; + } + + if (use_blocked_d512( + q, + k, + v, + has_arr_mask, + do_causal_, + sinks.has_value(), + lse != nullptr, + s)) { + sdpa_full_self_attention_blocked(s, d, q, k, v, scale_, o); + } else { + sdpa_full_self_attention_metal( + s, d, q, k, v, scale_, o, do_causal_, mask, sinks, lse); + } } metal::get_command_encoder(s).add_temporaries(std::move(copies)); } +// The vjp uses matmuls directly so it should be fine to call it every time as +// it uses a lot less memory. It can be overriden by setting +// MLX_SDPA_VJP_FALLBACK=1 bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) { - return true; + if (s.device != Device::gpu) { + return true; + } + if (env::get_var("MLX_SDPA_VJP_FALLBACK", 0) != 0) { + return true; + } + auto dt = q.dtype(); + return !(dt == float32 || dt == float16 || dt == bfloat16); } void ScaledDotProductAttentionVJP::eval_gpu( const std::vector& inputs, std::vector& outputs) { - throw std::runtime_error("NYI"); + auto& s = stream(); + auto& d = metal::device(s.device); + + if (has_sinks_) { + throw std::runtime_error( + "[ScaledDotProductAttentionVJP] NYI: attention sinks"); + } + + auto q = ensure_row_contiguous(inputs[0], d, s); + auto k = ensure_row_contiguous(inputs[1], d, s); + auto v = ensure_row_contiguous(inputs[2], d, s); + const int n_in = static_cast(inputs.size()); + auto o = ensure_row_contiguous(inputs[n_in - 3], d, s); + auto lse = ensure_row_contiguous(inputs[n_in - 2], d, s); + auto cot_o = ensure_row_contiguous(inputs[n_in - 1], d, s); + + const int B = q.shape(0); + const int H = q.shape(1); + const int qL = q.shape(2); + const int Dv = v.shape(3); + + auto& dq = outputs[0]; + auto& dk = outputs[1]; + auto& dv = outputs[2]; + + auto& compute_encoder = metal::get_command_encoder(s); + + dq.set_data(allocator::malloc(dq.nbytes())); + dk.set_data(allocator::malloc(dk.nbytes())); + dv.set_data(allocator::malloc(dv.nbytes())); + + array odo({B, H, qL}, float32, nullptr, {}); + odo.set_data(allocator::malloc(odo.nbytes())); + compute_encoder.add_temporary(odo); + + std::string odo_name; + concatenate( + odo_name, + "sdpa_vjp_odo_", + get_type_string(q.dtype()), + "_", + std::to_string(Dv)); + + auto odo_kernel = d.get_kernel(odo_name); + + compute_encoder.set_compute_pipeline_state(odo_kernel); + compute_encoder.set_input_array(o, 0); + compute_encoder.set_input_array(cot_o, 1); + compute_encoder.set_output_array(odo, 2); + compute_encoder.set_bytes(qL, 3); + compute_encoder.dispatch_threads( + MTL::Size(32, qL, B * H), MTL::Size(32, 1, 1)); + + sdpa_vjp_blocked( + s, d, q, k, v, odo, lse, cot_o, dq, dk, dv, scale_, do_causal_); } } // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp index 02d707f2b..7253d3c12 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp @@ -133,8 +133,9 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { in.flags()); } } else { - in = contiguous_copy_gpu(in, stream()); - out.copy_shared_buffer(in); + // Output type can be wider than input + copy_gpu(in, out, CopyType::General, stream()); + in = out; } scan_gpu_inplace( diff --git a/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp b/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp index f6ff41abc..d5f8ef92f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp @@ -50,6 +50,13 @@ array compute_dynamic_offset( auto& d = metal::device(s.device); auto& compute_encoder = metal::get_command_encoder(s); + // The kernel reads the indices contiguously + if (!indices.flags().row_contiguous) { + auto copy = contiguous_copy_gpu(indices, s); + compute_encoder.add_temporary(copy); + return compute_dynamic_offset(copy, strides, axes, s); + } + // Kernel to compute offset here. array offset({1}, int64, nullptr, {}); bool donate = indices.is_donatable() && diff --git a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp index 15d6bdfb4..b38d3ac2b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp @@ -2,6 +2,7 @@ #include +#include "mlx/backend/common/utils.h" #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" @@ -39,17 +40,24 @@ void single_block_sort( int in_stride_sorted_axis = in.strides()[axis]; int out_stride_sorted_axis = out.strides()[axis]; - // We can only use the contiguous kernel if the sorted axis - // has the largest or smallest stride. - // We also need the input to be contiguous + // The contiguous kernel walks the rows with a single stride, so the axes + // that are not sorted have to collapse to a single dim. + auto single_dim = + [](const Shape& nc_shape, const Strides& nc_str, int64_t& stride) { + auto [shape, strides] = collapse_contiguous_dims(nc_shape, nc_str); + if (shape.size() > 1) { + return false; + } + if (shape.size() > 0) { + stride = strides[0]; + } + return true; + }; + int64_t in_stride_segment_axis = 0; + int64_t out_stride_segment_axis = 0; bool contiguous = in.flags().contiguous; - auto check_strides = [](array x, int sort_stride) { - int min_stride = *std::min_element(x.strides().begin(), x.strides().end()); - int max_stride = *std::max_element(x.strides().begin(), x.strides().end()); - return sort_stride == min_stride || sort_stride == max_stride; - }; - contiguous &= check_strides(in, in_stride_sorted_axis); - contiguous &= check_strides(out, out_stride_sorted_axis); + contiguous &= single_dim(nc_shape, in_nc_str, in_stride_segment_axis); + contiguous &= single_dim(nc_shape, out_nc_str, out_stride_segment_axis); // Prepare kernel name std::ostringstream kname; @@ -74,22 +82,12 @@ void single_block_sort( compute_encoder.set_bytes(out_stride_sorted_axis, 4); if (contiguous) { - int in_stride_segment_axis = INT32_MAX; - int out_stride_segment_axis = INT32_MAX; - for (int i = 0; i < in_nc_str.size(); i++) { - if (nc_shape[i] == 1) { - continue; - } - if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) { - throw std::runtime_error("[Sort::eval_gpu] Stride too large."); - } - in_stride_segment_axis = - std::min(in_stride_segment_axis, static_cast(in_nc_str[i])); - out_stride_segment_axis = - std::min(out_stride_segment_axis, static_cast(out_nc_str[i])); + if (in_stride_segment_axis > INT32_MAX || + out_stride_segment_axis > INT32_MAX) { + throw std::runtime_error("[Sort::eval_gpu] Stride too large."); } - compute_encoder.set_bytes(in_stride_segment_axis, 5); - compute_encoder.set_bytes(out_stride_segment_axis, 6); + compute_encoder.set_bytes(static_cast(in_stride_segment_axis), 5); + compute_encoder.set_bytes(static_cast(out_stride_segment_axis), 6); } else { compute_encoder.set_bytes(nc_dim, 5); if (nc_shape.empty()) { diff --git a/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp index e52230715..2ad646118 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp @@ -66,6 +66,7 @@ NO_CPU(Expm1) NO_CPU(FFT) NO_CPU(Flatten) NO_CPU(Floor) +NO_CPU(FloorDivide) NO_CPU(Full) NO_CPU(Gather) NO_CPU(GatherAxis) diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp index c37515f9e..e3ee598ea 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/primitives.h" #include "mlx/distributed/primitives.h" @@ -30,6 +30,7 @@ bool fast::ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, @@ -62,6 +63,15 @@ bool fast::GatedDeltaUpdate::use_fallback( return true; } +bool fast::GatedDeltaUpdateVJP::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s) { + return true; +} + NO_GPU(Abs) NO_GPU(Add) NO_GPU(AddMM) @@ -109,6 +119,7 @@ NO_GPU(Expm1) NO_GPU(FFT) NO_GPU(Flatten) NO_GPU(Floor) +NO_GPU(FloorDivide) NO_GPU(Full) NO_GPU(Gather) NO_GPU(GatherAxis) @@ -190,6 +201,7 @@ NO_GPU_USE_FALLBACK(RoPE) NO_GPU_MULTI(ScaledDotProductAttention) NO_GPU_MULTI(ScaledDotProductAttentionVJP) NO_GPU_MULTI(GatedDeltaUpdate) +NO_GPU_MULTI(GatedDeltaUpdateVJP) NO_GPU_MULTI(ConvertFP8) NO_GPU_MULTI(Quantize) NO_GPU_MULTI(CustomKernel) diff --git a/Source/Cmlx/mlx/mlx/compile.cpp b/Source/Cmlx/mlx/mlx/compile.cpp index 3cab4c05d..e6bb2953b 100644 --- a/Source/Cmlx/mlx/mlx/compile.cpp +++ b/Source/Cmlx/mlx/mlx/compile.cpp @@ -47,17 +47,16 @@ bool is_unary(const Primitive& p) { } bool is_binary(const Primitive& p) { - return ( - typeid(p) == typeid(Add) || typeid(p) == typeid(Divide) || - typeid(p) == typeid(Equal) || typeid(p) == typeid(Greater) || - typeid(p) == typeid(GreaterEqual) || typeid(p) == typeid(Less) || - typeid(p) == typeid(LessEqual) || typeid(p) == typeid(LogicalNot) || - typeid(p) == typeid(LogicalAnd) || typeid(p) == typeid(LogicalOr) || - typeid(p) == typeid(LogAddExp) || typeid(p) == typeid(Maximum) || - typeid(p) == typeid(Minimum) || typeid(p) == typeid(Multiply) || - typeid(p) == typeid(NotEqual) || typeid(p) == typeid(Power) || - typeid(p) == typeid(Subtract) || typeid(p) == typeid(BitwiseBinary) || - typeid(p) == typeid(ArcTan2)); + return typeid(p) == typeid(Add) || typeid(p) == typeid(ArcTan2) || + typeid(p) == typeid(BitwiseBinary) || typeid(p) == typeid(Divide) || + typeid(p) == typeid(Equal) || typeid(p) == typeid(FloorDivide) || + typeid(p) == typeid(Greater) || typeid(p) == typeid(GreaterEqual) || + typeid(p) == typeid(Less) || typeid(p) == typeid(LessEqual) || + typeid(p) == typeid(LogicalNot) || typeid(p) == typeid(LogicalAnd) || + typeid(p) == typeid(LogicalOr) || typeid(p) == typeid(LogAddExp) || + typeid(p) == typeid(Maximum) || typeid(p) == typeid(Minimum) || + typeid(p) == typeid(Multiply) || typeid(p) == typeid(NotEqual) || + typeid(p) == typeid(Power) || typeid(p) == typeid(Subtract); } bool is_ternary(const Primitive& p) { @@ -651,11 +650,16 @@ void compile_simplify( return pa.is_equivalent(pb); }; - // Merge scalars + std::unordered_map output_map; + for (auto& o : outputs) { + output_map.insert({o.id(), o}); + } + + // Merge scalars, except outputs std::vector new_tape; for (auto& arr : tape) { // Check if we can merge scalars - if (is_scalar(arr)) { + if (is_scalar(arr) && output_map.find(arr.id()) == output_map.end()) { auto scalar = scalars.find(get_scalar_rep(arr)); if (scalar->second.id() != arr.id()) { merge(scalar->second, arr, parents_map); @@ -669,10 +673,6 @@ void compile_simplify( // Remove no-ops { - std::unordered_map output_map; - for (auto& o : outputs) { - output_map.insert({o.id(), o}); - } for (auto& arr : tape) { if (!arr.has_primitive() || !is_noop(arr.primitive())) { new_tape.push_back(std::move(arr)); @@ -694,11 +694,6 @@ void compile_simplify( tape_order.insert({tape[i].id(), i}); } - std::unordered_set output_set; - for (auto& o : outputs) { - output_set.insert(o.id()); - } - // Multi-pass merge only keeping non-orphaned arrays in the tape for (int pass = 0; pass < passes; ++pass) { for (auto& arr : tape) { @@ -718,7 +713,7 @@ void compile_simplify( auto& src = parents->second[src_idx].first; auto& dst = parents->second[dst_idx].first; if (src.id() != dst.id() && array_equivalent(src, dst) && - output_set.find(src.id()) == output_set.end()) { + output_map.find(src.id()) == output_map.end()) { merge(dst, src, parents_map); mask[src_idx] = true; } @@ -784,7 +779,7 @@ void compile_simplify( } return false; } else { - return output_set.find(a.id()) == output_set.end(); + return output_map.find(a.id()) == output_map.end(); } }; bool discard = maybe_merge_parents(arr); diff --git a/Source/Cmlx/mlx/mlx/distributed/distributed.cpp b/Source/Cmlx/mlx/mlx/distributed/distributed.cpp index 34be2ec60..d8fca8b16 100644 --- a/Source/Cmlx/mlx/mlx/distributed/distributed.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/distributed.cpp @@ -148,12 +148,11 @@ get_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}); + return Group(std::make_shared()); } + auto& backends = get_backends(); + backends.insert({"any", group}); backends.insert({std::move(bk), group}); return Group(group); } diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp index 7715022fc..17bc350a1 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp @@ -180,14 +180,25 @@ bool Config::is_valid() const { (is_valid_mesh() || is_valid_ring()); } -std::vector Config::get_mesh_connectivity() const { +std::vector> Config::get_mesh_connectivity() const { if (!is_valid_mesh()) { throw std::runtime_error("[jaccl] The devices do not form a valid mesh."); } - std::vector devices(size_); - for (int dst = 0; dst < size_; dst++) { - if (dst != rank_) { - devices[dst] = devices_[rank_][dst][0]; + size_t n_wires = devices_[0][1].size(); + for (int src = 0; src < size_; src++) { + for (int dst = 0; dst < size_; dst++) { + if (src != dst) { + n_wires = std::min(n_wires, devices_[src][dst].size()); + } + } + } + std::vector> devices( + n_wires, std::vector(size_)); + for (int w = 0; w < devices.size(); w++) { + for (int dst = 0; dst < size_; dst++) { + if (dst != rank_) { + devices[w][dst] = devices_[rank_][dst][w]; + } } } return devices; diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h index 39feb34b4..dc430c336 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h @@ -51,7 +51,7 @@ class Config { friend std::shared_ptr init(const Config& cfg, bool strict); private: - std::vector get_mesh_connectivity() const; + std::vector> get_mesh_connectivity() const; std::pair, std::vector> get_ring_connectivity() const; SideChannel get_side_channel() const; diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 5b749f2bd..fe856faf9 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -8,12 +8,13 @@ namespace jaccl { MeshGroup::MeshGroup( int rank, - const std::vector& device_names, + const std::vector>& device_names, SideChannel sc) : rank_(rank), - size_(device_names.size()), + size_(device_names[0].size()), + n_wires_(device_names.size()), side_channel_(std::move(sc)), - connections_(create_connections(device_names)) { + pool_(n_wires_ - 1) { if (size_ > MESH_MAX_PEERS) { std::ostringstream msg; msg << "[jaccl] The JACCL mesh supports up to " << MESH_MAX_PEERS @@ -21,54 +22,69 @@ MeshGroup::MeshGroup( throw std::runtime_error(msg.str()); } + for (auto& names : device_names) { + connections_.push_back(create_connections(names)); + } + // Initialize all the connections and allocate buffers initialize(); // Make sure every node has reached here before continuing side_channel_.barrier(); - // Create the mesh implementation object - mesh_ = MeshImpl(rank_, size_, connections_, buffers_, scatter_buffers_); + // Create the mesh implementation objects + for (int w = 0; w < n_wires_; w++) { + meshes_.emplace_back( + rank_, size_, connections_[w], buffers_[w], scatter_buffers_[w]); + } } void MeshGroup::initialize() { // Create the queue pairs - for (auto& conn : connections_) { - if (conn.ctx == nullptr) { - continue; + for (auto& wire : connections_) { + for (auto& conn : wire) { + if (conn.ctx == nullptr) { + continue; + } + conn.allocate_protection_domain(); + conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_queue_pair(); } - conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); - conn.create_queue_pair(); } allocate_buffers(); // First init all connections - for (int peer = 0; peer < size_; peer++) { - if (peer == rank_) { - continue; + for (auto& wire : connections_) { + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + wire[peer].queue_pair_init(); } - connections_[peer].queue_pair_init(); } // Gather the information to be exchanged, this also serves as a barrier // so that all peers have initialized their connections before attempting // to transition to RTS. std::vector info; - for (auto& conn : connections_) { - info.emplace_back(conn.info()); + for (auto& wire : connections_) { + for (auto& conn : wire) { + info.emplace_back(conn.info()); + } } auto all_infos = side_channel_.all_gather(info); // Transition queue pairs to RTS - for (int peer = 0; peer < size_; peer++) { - if (peer == rank_) { - continue; + for (int w = 0; w < n_wires_; w++) { + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + auto peer_info = all_infos[peer][w * size_ + rank_]; + connections_[w][peer].queue_pair_rtr(peer_info); + connections_[w][peer].queue_pair_rts(); } - auto peer_info = all_infos[peer][rank_]; - connections_[peer].queue_pair_rtr(peer_info); - connections_[peer].queue_pair_rts(); } } @@ -76,54 +92,62 @@ void MeshGroup::allocate_buffers() { // Deregister any buffers and free the memory buffers_.clear(); scatter_buffers_.clear(); + buffers_.resize(n_wires_); + scatter_buffers_.resize(n_wires_); - // Allocate the memory - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - // Mesh buffers - for (int j = 0; j < size_; j++) { - 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)); + for (int w = 0; w < n_wires_; w++) { + auto& conns = connections_[w]; + auto& buffers = buffers_[w]; + auto& scatter_buffers = scatter_buffers_[w]; + + // Allocate the memory + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + // Mesh buffers + for (int j = 0; j < size_; j++) { + 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)); + } } } - } - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - // Mesh buffers - for (int j = 0; j < size_; j++) { - if (j == rank_) { - // This is our send buffer so register it with all pds so we can - // send it to all connected devices. - for (auto& conn : connections_) { - if (conn.ctx != nullptr) { - buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] - .register_to_protection_domain(conn.protection_domain); + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + // Mesh buffers + for (int j = 0; j < size_; j++) { + if (j == rank_) { + // This is our send buffer so register it with all pds so we can + // send it to all connected devices. + for (auto& conn : conns) { + if (conn.ctx != nullptr) { + buffers[k * NUM_BUFFERS * size_ + i * size_ + j] + .register_to_protection_domain(conn.protection_domain); + } } + } else { + // This is the recv buffer from rank j so register it to rank j's + // protection domain. + buffers[k * NUM_BUFFERS * size_ + i * size_ + j] + .register_to_protection_domain(conns[j].protection_domain); } - } else { - // This is the recv buffer from rank j so register it to rank j's - // protection domain. - buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] - .register_to_protection_domain(connections_[j].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. 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( + conns[j].protection_domain); + scatter_buffers[scatter_base + size_ + j] + .register_to_protection_domain(conns[j].protection_domain); } - 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); } } } @@ -163,8 +187,11 @@ void MeshGroup::all_min( } void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) { - mesh_.all_gather( - static_cast(input), static_cast(output), n_bytes); + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + split_wires(n_bytes, n_bytes, [&](int w, int64_t offset, int64_t n) { + meshes_[w].all_gather(in_ptr + offset, out_ptr + offset, n, n_bytes); + }); } void MeshGroup::sum_scatter( @@ -179,11 +206,17 @@ void MeshGroup::sum_scatter( } void MeshGroup::send(const void* input, size_t n_bytes, int dst) { - mesh_.send(static_cast(input), n_bytes, dst); + auto in_ptr = static_cast(input); + split_wires(n_bytes, n_bytes, [&](int w, int64_t offset, int64_t n) { + meshes_[w].send(in_ptr + offset, n, dst); + }); } void MeshGroup::recv(void* output, size_t n_bytes, int src) { - mesh_.recv(static_cast(output), n_bytes, src); + auto out_ptr = static_cast(output); + split_wires(n_bytes, n_bytes, [&](int w, int64_t offset, int64_t n) { + meshes_[w].recv(out_ptr + offset, n, src); + }); } void MeshGroup::barrier() { @@ -191,6 +224,16 @@ void MeshGroup::barrier() { all_sum(&b, &b, sizeof(b), Dtype::UInt8); } +template +void MeshGroup::split_wires(size_t n_bytes, int64_t total, Fn&& fn) { + int n_wires = n_bytes >= MESH_MULTI_WIRE_MIN_BYTES ? n_wires_ : 1; + int64_t per_wire = (total + n_wires - 1) / n_wires; + dispatch_wires(&pool_, n_wires, [&](int w) { + int64_t offset = std::min(total, w * per_wire); + fn(w, offset, std::min(per_wire, total - offset)); + }); +} + template void MeshGroup::all_reduce( const void* input, @@ -200,16 +243,19 @@ 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 && 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); - } + split_wires(n_bytes, count, [&](int w, int64_t offset, int64_t n) { + if (size_ > 2 && n * sizeof(T) > 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. + meshes_[w].all_reduce_scatter_gather( + in_ptr + offset, out_ptr + offset, n, reduce_op); + } else { + // Small messages are latency bound so use the single phase fully + // connected all_reduce cause it is a bit better. + meshes_[w].all_reduce(in_ptr + offset, out_ptr + offset, n, reduce_op); + } + }); } template @@ -223,7 +269,10 @@ void MeshGroup::reduce_scatter( 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); + split_wires(n_bytes, count, [&](int w, int64_t offset, int64_t n) { + meshes_[w].sum_scatter( + in_ptr + offset, out_ptr + offset, n, count, reduce_op); + }); } } // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h index 8e5164aee..aa19fddb8 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h @@ -5,13 +5,14 @@ #include "jaccl/group.h" #include "jaccl/mesh_impl.h" #include "jaccl/rdma.h" +#include "jaccl/threadpool.h" namespace jaccl { /** - * The JACCL communication group for a fully connected mesh. We expect one - * connection per peer and it should be the lowest latency communication group - * for small to medium size messages. + * The JACCL communication group for a fully connected mesh. We expect one or + * more connections per peer and it should be the lowest latency communication + * group for small to medium size messages. * * Like all JACCL groups it uses a side channel to exchange the necessary * information and then configure the connections to be ready for RDMA @@ -21,7 +22,7 @@ class MeshGroup : public Group { public: MeshGroup( int rank, - const std::vector& device_names, + const std::vector>& device_names, SideChannel sc); int rank() override { @@ -66,6 +67,9 @@ class MeshGroup : public Group { size_t n_bytes, ReduceOp reduce_op); + template + void split_wires(size_t n_bytes, int64_t total, Fn&& fn); + /** * Performs the connection initialization. Namely, after this call all * Connection objects should have a queue pair in RTS state and all buffers @@ -80,12 +84,13 @@ class MeshGroup : public Group { int rank_; int size_; + int n_wires_; SideChannel side_channel_; - std::vector connections_; - std::vector buffers_; - std::vector scatter_buffers_; - - MeshImpl mesh_; + std::vector> connections_; + std::vector> buffers_; + std::vector> scatter_buffers_; + ThreadPool pool_; + std::vector meshes_; }; } // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 75c540e91..7beed48e6 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -10,6 +10,7 @@ constexpr int MESH_MAX_PEERS = 8; constexpr int MESH_PIPELINE = 2; constexpr int64_t MAX_BUFFER_SIZE = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); +constexpr size_t MESH_MULTI_WIRE_MIN_BYTES = 512 * 1024; namespace jaccl { @@ -201,17 +202,21 @@ class MeshImpl { } } - void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) { + void all_gather( + const char* in_ptr, + char* out_ptr, + int64_t n_bytes, + int64_t stride) { // 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); + // out_ptr + rank_ * stride). + char* data = out_ptr; + char* our_data = out_ptr + rank_ * stride; + if (in_ptr != our_data) { + std::memcpy(our_data, in_ptr, n_bytes); } // Fully connected all gather - char* data = out_ptr; - char* our_data = out_ptr + rank_ * n_bytes; auto [sz, N] = buffer_size_from_message(n_bytes); constexpr int PIPELINE = 2; constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; @@ -274,7 +279,7 @@ class MeshImpl { recv_buffer(sz, buff, rank).begin(), recv_buffer(sz, buff, rank).begin() + std::min(N, total - write_offset[rank]), - data + rank * n_bytes + write_offset[rank]); + data + rank * stride + write_offset[rank]); write_offset[rank] += N; if (write_offset[rank] + N * (PIPELINE - 1) < total) { recv_from(sz, rank, buff); @@ -286,20 +291,26 @@ class MeshImpl { } template - void sum_scatter(const T* in, T* out, int64_t count, ReduceOp reduce_op) { + void sum_scatter( + const T* in, + T* out, + int64_t count, + int64_t stride, + 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. + // The input holds size_ chunks of `count` elements each, `stride` elements + // apart. 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; + const T* our_chunk = in + static_cast(rank_) * stride; auto [sz, buffer_size] = buffer_size_from_message(count * sizeof(T)); int64_t N = buffer_size / sizeof(T); @@ -328,7 +339,7 @@ class MeshImpl { if (p == rank_) { continue; } - const T* src = in + static_cast(p) * count + read_offset; + const T* src = in + static_cast(p) * stride + read_offset; std::copy( src, src + elems, scatter_send_buffer(sz, buff, p).begin()); } @@ -362,7 +373,8 @@ class MeshImpl { if (p == rank_) { continue; } - const T* src = in + static_cast(p) * count + read_offset; + const T* src = + in + static_cast(p) * stride + read_offset; std::copy( src, src + elems, @@ -424,7 +436,7 @@ class MeshImpl { // 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); + sum_scatter(in, shard, chunk, 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 @@ -432,6 +444,7 @@ class MeshImpl { all_gather( reinterpret_cast(shard), reinterpret_cast(out), + chunk * sizeof(T), chunk * sizeof(T)); } diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index d6305d9d4..4fd46547b 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -2,7 +2,6 @@ #pragma once -#include #include #include "jaccl/rdma.h" @@ -90,7 +89,7 @@ class RingImpl { // 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) { + dispatch_wires(pool_, n_wires, [&](int lw) { all_reduce_wire( in_ptr, out_ptr, @@ -196,7 +195,7 @@ class RingImpl { // 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) { + dispatch_wires(pool_, n_wires, [&](int lw) { all_gather_wire( out_ptr, n_bytes, @@ -273,7 +272,7 @@ class RingImpl { // 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) { + dispatch_wires(pool_, n_wires, [&](int lw) { reduce_scatter_wire( in_ptr, out_ptr, chunk, size_per_wire, n_wires, lw, reduce_op); }); @@ -595,7 +594,7 @@ class RingImpl { // 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) { + dispatch_wires(pool_, n_wires, [&](int lw) { send_wire(in_ptr, n_bytes, dir, bytes_per_wire, lw); }); } @@ -674,7 +673,7 @@ class RingImpl { // 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) { + dispatch_wires(pool_, n_wires, [&](int lw) { recv_wire(out_ptr, n_bytes, dir, bytes_per_wire, lw); }); } @@ -780,38 +779,6 @@ class RingImpl { 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++) { - fn(lw); - } - return; - } - - std::vector> futures; - futures.reserve(n_wires - 1); - for (int lw = 0; lw < n_wires - 1; lw++) { - futures.emplace_back(pool_->enqueue(fn, 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(); - } - } - int rank_; int size_; int n_conns_; diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h index b5d05153e..6e96152cb 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h @@ -134,4 +134,36 @@ inline void ThreadPool::start_threads(size_t threads) { } } +// 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(ThreadPool* pool, int n_wires, Fn&& fn) { + if (n_wires <= 1 || pool == nullptr) { + for (int lw = 0; lw < n_wires; lw++) { + fn(lw); + } + return; + } + + std::vector> futures; + futures.reserve(n_wires - 1); + for (int lw = 0; lw < n_wires - 1; lw++) { + futures.emplace_back(pool->enqueue(fn, 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(); + } +} + } // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/export.cpp b/Source/Cmlx/mlx/mlx/export.cpp index 55ccbe187..c10934d1b 100644 --- a/Source/Cmlx/mlx/mlx/export.cpp +++ b/Source/Cmlx/mlx/mlx/export.cpp @@ -384,6 +384,7 @@ struct PrimitiveFactory { SERIALIZE_PRIMITIVE(FFT), SERIALIZE_PRIMITIVE(Flatten), SERIALIZE_PRIMITIVE(Floor), + SERIALIZE_PRIMITIVE(FloorDivide), SERIALIZE_PRIMITIVE(Full), SERIALIZE_PRIMITIVE(Gather), SERIALIZE_PRIMITIVE(GatherAxis), diff --git a/Source/Cmlx/mlx/mlx/fast.cpp b/Source/Cmlx/mlx/mlx/fast.cpp index 101fe127c..6f9f3db73 100644 --- a/Source/Cmlx/mlx/mlx/fast.cpp +++ b/Source/Cmlx/mlx/mlx/fast.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include #include @@ -352,7 +352,7 @@ array layer_norm( auto xc = subtract(x, mu, s); auto v = mean(square(xc, s), /* axis= */ -1, /* keepdims= */ true, s); - x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s)); + x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s), s); x = astype(x, out_type, s); // If the LN is affine then transform x according to the weight and bias @@ -934,6 +934,7 @@ array scaled_dot_product_attention( has_mask, has_arr_mask, do_causal, + has_sinks, is_training, output_logsumexp, force_fused, @@ -944,8 +945,9 @@ array scaled_dot_product_attention( array& mask = inputs[3]; mask = where( mask, - full_like(mask, 0, final_type, s), - full_like(mask, -inf, final_type, s)); + full_like(mask, 0, final_type, stream), + full_like(mask, -inf, final_type, stream), + stream); } Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)}; auto primitive = std::make_shared( @@ -977,13 +979,16 @@ std::vector ScaledDotProductAttention::vjp( const std::vector& outputs) { assert(primals.size() >= 3); assert(cotangents.size() == outputs.size()); - auto s = stream(); if (ScaledDotProductAttentionVJP::use_fallback(primals[0], s)) { assert(outputs.size() == 1); return Custom::vjp(primals, cotangents, argnums, outputs); } + if (outputs.size() != 2) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + auto fallback = [sdpa = fallback_, s](const std::vector& inputs) { std::vector primals(inputs.begin(), std::prev(inputs.end())); auto [_, vjps] = mlx::core::vjp(sdpa, primals, {inputs.back()}); @@ -1158,6 +1163,54 @@ std::vector gated_delta_update( auto result = fallback({q, k, v, g, beta, h0, mask}); return result; } +std::vector GatedDeltaUpdate::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + const int Hk = primals[0].shape(2); + const int Dk = primals[0].shape(3); + const int Hv = primals[2].shape(2); + const int Dv = primals[2].shape(3); + + if (GatedDeltaUpdateVJP::use_fallback(Hk, Dk, Hv, Dv, stream())) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + + if (primals.size() != 6) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 6 primals (q,k,v,g,beta,h0), got " + + std::to_string(primals.size())); + } + if (cotangents.size() != 2) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 2 cotangents (out, state), got " + + std::to_string(cotangents.size())); + } + + std::vector inputs = primals; // q,k,v,g,beta,h0 (indices 0..5) + inputs.push_back(cotangents[0]); // cot_o -> index 6 + inputs.push_back(cotangents[1]); // cot_h -> index 7 + + std::vector shapes; + std::vector dtypes; + for (int i = 0; i < 6; ++i) { + shapes.push_back(primals[i].shape()); + dtypes.push_back(i == 5 ? float32 : primals[i].dtype()); // dh is float32 + } + + auto vjps = array::make_arrays( + std::move(shapes), + std::move(dtypes), + std::make_shared(stream(), fallback_), + std::move(inputs)); + + std::vector returned_vjps; + for (int arg : argnums) { + returned_vjps.push_back(std::move(vjps[arg])); + } + return returned_vjps; +} bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); diff --git a/Source/Cmlx/mlx/mlx/fast_primitives.h b/Source/Cmlx/mlx/mlx/fast_primitives.h index cef13bb0f..8117a7a7b 100644 --- a/Source/Cmlx/mlx/mlx/fast_primitives.h +++ b/Source/Cmlx/mlx/mlx/fast_primitives.h @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #include #include @@ -289,6 +289,7 @@ class ScaledDotProductAttention : public Custom { bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, @@ -418,6 +419,12 @@ class GatedDeltaUpdate : public Custom { void eval_gpu(const std::vector& inputs, std::vector& outputs) override; + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + DEFINE_NAME(GatedDeltaUpdate); DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { @@ -425,6 +432,39 @@ class GatedDeltaUpdate : public Custom { } private: + bool is_training_; +}; + +class GatedDeltaUpdateVJP : public Custom { + public: + GatedDeltaUpdateVJP( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + DEFINE_NAME(GatedDeltaUpdateVJP); + DEFINE_INPUT_OUTPUT_SHAPE() + auto state() const { + return std::make_tuple(nullptr); /* TODO */ + } + + private: + bool has_cache_; }; class Quantize : public Custom { diff --git a/Source/Cmlx/mlx/mlx/fft.cpp b/Source/Cmlx/mlx/mlx/fft.cpp index 55596472a..e1c5756e1 100644 --- a/Source/Cmlx/mlx/mlx/fft.cpp +++ b/Source/Cmlx/mlx/mlx/fft.cpp @@ -1,4 +1,5 @@ // Copyright © 2023 Apple Inc. +#include #include #include #include @@ -114,7 +115,12 @@ array fft_impl( auto out = array( out_shape, out_type, - std::make_shared(to_stream(s), valid_axes, inverse, real), + std::make_shared( + to_stream(s), + valid_axes, + inverse, + real, + real && inverse && (n.back() % 2 != 0)), {astype(in, in_type, s)}); auto scale = fft_scale_factor(n, norm, inverse); if (scale != 1.0) { @@ -335,4 +341,230 @@ array rfftfreq(int n, double d /* = 1.0 */, StreamOrDevice s /* = {} */) { return multiply(freqs, scale, s); } +namespace { + +// Center a window of length <= n_fft inside an n_fft-length frame. +array prepare_window( + const std::optional& window, + int win_length, + int n_fft, + const std::string& op, + StreamOrDevice s) { + array win = + window.has_value() ? window.value() : ones({win_length}, float32, s); + if (win.ndim() != 1) { + throw std::invalid_argument( + "[" + op + "] window must be a one-dimensional array."); + } + int wl = win.shape(0); + if (wl > n_fft) { + throw std::invalid_argument( + "[" + op + "] window length must not exceed n_fft."); + } + if (wl < n_fft) { + int pad_left = (n_fft - wl) / 2; + int pad_right = n_fft - wl - pad_left; + win = + pad(win, + {0}, + {pad_left}, + {pad_right}, + array(0, win.dtype()), + "constant", + s); + } + return win; +} + +} // namespace + +array stft( + const array& x_in, + int n_fft /* = 2048 */, + const std::optional& hop_length_ /* = std::nullopt */, + const std::optional& win_length_ /* = std::nullopt */, + const std::optional& window /* = std::nullopt */, + bool center /* = true */, + const std::string& pad_mode /* = "reflect" */, + FFTNorm norm /* = FFTNorm::Backward */, + const std::optional& onesided_ /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + if (n_fft <= 0) { + throw std::invalid_argument("[stft] n_fft must be positive."); + } + if (x_in.ndim() < 1) { + throw std::invalid_argument( + "[stft] Input must have at least one dimension."); + } + int hop_length = hop_length_.value_or(n_fft / 4); + int win_length = win_length_.value_or(n_fft); + if (hop_length <= 0) { + throw std::invalid_argument("[stft] hop_length must be positive."); + } + bool onesided = onesided_.value_or(x_in.dtype() != complex64); + + array win = prepare_window(window, win_length, n_fft, "stft", s); + + // Collapse leading dims into a single batch axis. + Shape lead = x_in.shape(); + lead.pop_back(); + int t = x_in.shape(-1); + array x = reshape(x_in, {-1, t}, s); + + if (center) { + int p = n_fft / 2; + x = pad(x, {1}, {p}, {p}, array(0, x.dtype()), pad_mode, s); + } + int tp = x.shape(-1); + if (tp < n_fft) { + throw std::invalid_argument( + "[stft] Input is too short for the given n_fft."); + } + int n_frames = 1 + (tp - n_fft) / hop_length; + int b = x.shape(0); + + // Overlapping frames via as_strided (element strides: tp, hop, 1). + array frames = as_strided( + x, + {b, n_frames, n_fft}, + Strides{static_cast(tp), static_cast(hop_length), 1}, + 0, + s); + frames = multiply(frames, win, s); + + array spec = onesided ? rfftn(frames, Shape{n_fft}, {-1}, norm, s) + : fftn(frames, Shape{n_fft}, {-1}, norm, s); + spec = swapaxes(spec, -1, -2, s); + + Shape out_shape = lead; + out_shape.push_back(spec.shape(-2)); + out_shape.push_back(n_frames); + return reshape(spec, out_shape, s); +} + +array istft( + const array& stft_matrix, + const std::optional& n_fft_ /* = std::nullopt */, + const std::optional& hop_length_ /* = std::nullopt */, + const std::optional& win_length_ /* = std::nullopt */, + const std::optional& window /* = std::nullopt */, + bool center /* = true */, + FFTNorm norm /* = FFTNorm::Backward */, + const std::optional& onesided_ /* = std::nullopt */, + const std::optional& length_ /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + if (stft_matrix.ndim() < 2) { + throw std::invalid_argument( + "[istft] Input must have at least two dimensions (freq, frames)."); + } + bool onesided = onesided_.value_or(true); + int n_freq = stft_matrix.shape(-2); + int n_fft = n_fft_.value_or(onesided ? (n_freq - 1) * 2 : n_freq); + if (n_fft <= 0) { + throw std::invalid_argument("[istft] n_fft must be positive."); + } + int hop_length = hop_length_.value_or(n_fft / 4); + int win_length = win_length_.value_or(n_fft); + if (hop_length <= 0) { + throw std::invalid_argument("[istft] hop_length must be positive."); + } + + array win = prepare_window(window, win_length, n_fft, "istft", s); + + Shape lead = stft_matrix.shape(); + int n_frames = lead.back(); + lead.pop_back(); + lead.pop_back(); + array z = reshape(stft_matrix, {-1, n_freq, n_frames}, s); + int b = z.shape(0); + z = swapaxes(z, -1, -2, s); + + array frames = onesided ? irfftn(z, Shape{n_fft}, {-1}, norm, s) + : real(ifftn(z, Shape{n_fft}, {-1}, norm, s), s); + frames = multiply(frames, win, s); + + // Overlap-add: sub-block k of frame i lands on output block i + k, so + // summing seg = ceil(n_fft / hop) shifted copies is independent of n_frames. + int seg = (n_fft + hop_length - 1) / hop_length; + int wp = seg * hop_length; + if (wp > n_fft) { + frames = + pad(frames, + {2}, + {0}, + {wp - n_fft}, + array(0, frames.dtype()), + "constant", + s); + } + array c = reshape(frames, {b, n_frames, seg, hop_length}, s); + + array win_sq = square(win, s); + if (wp > n_fft) { + win_sq = + pad(win_sq, + {0}, + {0}, + {wp - n_fft}, + array(0, win_sq.dtype()), + "constant", + s); + } + array cw = broadcast_to( + reshape(win_sq, {1, 1, seg, hop_length}, s), + {1, n_frames, seg, hop_length}, + s); + + int out_blocks = n_frames + seg - 1; + int out_len = out_blocks * hop_length; + array signal = zeros({b, out_len}, frames.dtype(), s); + array envelope = zeros({1, out_len}, frames.dtype(), s); + for (int k = 0; k < seg; ++k) { + array ck = reshape( + slice(c, {0, 0, k, 0}, {b, n_frames, k + 1, hop_length}, s), + {b, n_frames, hop_length}, + s); + ck = pad(ck, {1}, {k}, {seg - 1 - k}, array(0, ck.dtype()), "constant", s); + signal = add(signal, reshape(ck, {b, out_len}, s), s); + + array wk = reshape( + slice(cw, {0, 0, k, 0}, {1, n_frames, k + 1, hop_length}, s), + {1, n_frames, hop_length}, + s); + wk = pad(wk, {1}, {k}, {seg - 1 - k}, array(0, wk.dtype()), "constant", s); + envelope = add(envelope, reshape(wk, {1, out_len}, s), s); + } + signal = + divide(signal, maximum(envelope, array(1e-8f, envelope.dtype()), s), s); + + // Drop the centering pad from the start. The end is bounded by `length` when + // given, otherwise by removing the centering pad from the end as well. The + // full reconstruction spans n_fft + (n_frames - 1) * hop before centering. + int sig_len = n_fft + (n_frames - 1) * hop_length; + int start = center ? n_fft / 2 : 0; + if (length_.has_value()) { + int length = length_.value(); + int stop = std::min(sig_len, start + length); + signal = slice(signal, {0, start}, {b, stop}, s); + int cur = signal.shape(-1); + if (cur < length) { + signal = + pad(signal, + {1}, + {0}, + {length - cur}, + array(0, signal.dtype()), + "constant", + s); + } + } else { + int stop = center ? sig_len - n_fft / 2 : sig_len; + signal = slice(signal, {0, start}, {b, stop}, s); + } + + Shape out_shape = lead; + out_shape.push_back(signal.shape(-1)); + return reshape(signal, out_shape, s); +} + } // namespace mlx::core::fft diff --git a/Source/Cmlx/mlx/mlx/fft.h b/Source/Cmlx/mlx/mlx/fft.h index f69b8aa2c..f067ff49c 100644 --- a/Source/Cmlx/mlx/mlx/fft.h +++ b/Source/Cmlx/mlx/mlx/fft.h @@ -3,6 +3,7 @@ #pragma once #include +#include #include #include "array.h" @@ -212,6 +213,7 @@ inline array irfft2( StreamOrDevice s = {}) { return irfftn(a, axes, norm, s); } + /** Compute the discrete Fourier Transform sample frequencies. */ MLX_API array fftfreq(int n, double d = 1.0, StreamOrDevice s = {}); @@ -234,4 +236,30 @@ MLX_API array ifftshift(const array& a, StreamOrDevice s = {}); MLX_API array ifftshift(const array& a, const std::vector& axes, StreamOrDevice s = {}); +/** Compute the Short-Time Fourier Transform of the last axis of `x`. */ +MLX_API array stft( + const array& x, + int n_fft = 2048, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + const std::string& pad_mode = "reflect", + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + StreamOrDevice s = {}); + +/** Compute the inverse Short-Time Fourier Transform. */ +MLX_API array istft( + const array& stft_matrix, + const std::optional& n_fft = std::nullopt, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + const std::optional& length = std::nullopt, + StreamOrDevice s = {}); + } // namespace mlx::core::fft diff --git a/Source/Cmlx/mlx/mlx/io/gguf.cpp b/Source/Cmlx/mlx/mlx/io/gguf.cpp index dddcbfc77..bc5405624 100644 --- a/Source/Cmlx/mlx/mlx/io/gguf.cpp +++ b/Source/Cmlx/mlx/mlx/io/gguf.cpp @@ -76,7 +76,7 @@ Shape get_shape(const gguf_tensor& tensor) { return shape; } -std::tuple extract_tensor_data(gguf_tensor* tensor) { +std::tuple extract_tensor_data(gguf_tensor* tensor) { if (tensor == nullptr) { throw std::invalid_argument( "[extract_tensor_data] Input tensor pointer is null."); @@ -87,12 +87,12 @@ std::tuple extract_tensor_data(gguf_tensor* tensor) { if (tensor->weights_data == nullptr) { throw std::runtime_error("[load_gguf] NULL tensor data pointer"); } - allocator::Buffer buffer = allocator::malloc(tensor->bsize); + auto out = allocator::malloc(tensor->bsize); memcpy( - buffer.raw_ptr(), + out.buffer().raw_ptr(), tensor->weights_data, tensor->num_weights * equivalent_dtype.value().size()); - return {buffer, equivalent_dtype.value()}; + return {std::move(out), equivalent_dtype.value()}; } // Otherwise, we convert to float16. // TODO: Add other dequantization options. @@ -101,10 +101,10 @@ std::tuple extract_tensor_data(gguf_tensor* tensor) { throw std::runtime_error("[load_gguf] gguf_tensor_to_f16 failed"); } const size_t new_size = tensor->num_weights * sizeof(int16_t); - allocator::Buffer buffer = allocator::malloc(new_size); - memcpy(buffer.raw_ptr(), data, new_size); + auto out = allocator::malloc(new_size); + memcpy(out.buffer().raw_ptr(), data, new_size); free(data); - return {buffer, float16}; + return {std::move(out), float16}; } void set_mx_value_from_gguf( @@ -354,8 +354,8 @@ std::unordered_map load_arrays(gguf_ctx* ctx) { gguf_load_quantized(array_map, tensor); } else { std::string name(tensor.name, tensor.namelen); - const auto& [data, dtype] = extract_tensor_data(&tensor); - array loaded_array = array(data, get_shape(tensor), dtype); + auto [data, dtype] = extract_tensor_data(&tensor); + array loaded_array = array(std::move(data), get_shape(tensor), dtype); check_insert(array_map.insert({name, loaded_array})); } } diff --git a/Source/Cmlx/mlx/mlx/io/load.cpp b/Source/Cmlx/mlx/mlx/io/load.cpp index cd34ddc9c..5941d2edc 100644 --- a/Source/Cmlx/mlx/mlx/io/load.cpp +++ b/Source/Cmlx/mlx/mlx/io/load.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2026 Apple Inc. #include +#include #include #include #include @@ -305,10 +306,23 @@ array load(std::shared_ptr in_stream, StreamOrDevice s) { size_t ed = header.find_last_of(')'); std::string shape_str = header.substr(st, ed - st); + size_t max_size = std::numeric_limits::max() / size_of(dtype); + size_t total_size = 1; while (!shape_str.empty()) { // Read current number and get position of comma size_t pos; int dim = std::stoi(shape_str, &pos); + + // Reject negative dims and a total byte count that overflows size_t + if (dim < 0) { + throw std::runtime_error( + "[load] Negative dimension in shape in " + in_stream->label()); + } + if (dim != 0 && total_size > max_size / dim) { + throw std::runtime_error( + "[load] Shape is too large in " + in_stream->label()); + } + total_size *= dim; shape.push_back(dim); // Skip the comma and space and read the next number @@ -405,7 +419,15 @@ ThreadPool& ParallelFileReader::thread_pool() { void ParallelFileReader::read(char* data, size_t n) { while (n != 0) { auto m = ::read(fd_, data, std::min(n, static_cast(INT32_MAX))); - if (m <= 0) { + if (m < 0) { + if (errno == EINTR) { + continue; + } + std::ostringstream msg; + msg << "[read] Unable to read " << n << " bytes from file."; + throw std::runtime_error(msg.str()); + } + if (m == 0) { std::ostringstream msg; msg << "[read] Unable to read " << n << " bytes from file."; throw std::runtime_error(msg.str()); @@ -419,11 +441,18 @@ void ParallelFileReader::read(char* data, size_t n, size_t offset) { auto readfn = [fd = fd_](size_t offset, size_t size, char* buffer) -> bool { while (size != 0) { auto m = pread(fd, buffer, size, offset); - if (m <= 0) { + if (m < 0) { + if (errno == EINTR) { + continue; + } + return false; + } + if (m == 0) { return false; } buffer += m; size -= m; + offset += m; } return true; }; diff --git a/Source/Cmlx/mlx/mlx/ops.cpp b/Source/Cmlx/mlx/mlx/ops.cpp index 095ba9dcf..866349e03 100644 --- a/Source/Cmlx/mlx/mlx/ops.cpp +++ b/Source/Cmlx/mlx/mlx/ops.cpp @@ -96,7 +96,7 @@ array indices_or_default( << " Got indices with invalid dtype. Indices must be integral."; throw std::invalid_argument(msg.str()); } - return astype(indices.value(), uint32); + return astype(indices.value(), uint32, s); } Shape shape(x.shape().begin(), x.shape().end() - 2); @@ -726,7 +726,7 @@ array flip(const array& a, StreamOrDevice s /* = {} */) { namespace { inline auto -normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { +normalize_slice(const Shape& shape, Shape& start, Shape& stop, Shape& strides) { // - Start indices are normalized // - End indices are unchanged as -1 means something different // pre-normalization (the end of the axis) versus post normalization (the @@ -774,7 +774,13 @@ normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { } // Simplify the stride if it's unused if (out_shape[i] == 1) { - strides[i] = 1; + if (strides[i] < 0 && start[i] > 0) { + strides[i] = -1; + stop[i] = start[i] - 1; + } else { + strides[i] = 1; + stop[i] = start[i] + 1; + } } } @@ -2651,6 +2657,31 @@ array min( return min(a, std::vector{axis}, keepdims, s); } +std::tuple +minmax(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) { + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + return {min(a, axes, keepdims, s), max(a, axes, keepdims, s)}; +} + +std::tuple minmax( + const array& a, + const std::vector& axes, + bool keepdims /* = false */, + StreamOrDevice s /* = {}*/) { + return {min(a, axes, keepdims, s), max(a, axes, keepdims, s)}; +} + +std::tuple minmax( + const array& a, + int axis, + bool keepdims /* = false */, + StreamOrDevice s /* = {} */) { + return { + min(a, std::vector{axis}, keepdims, s), + max(a, std::vector{axis}, keepdims, s)}; +} + array argmin(const array& a, bool keepdims, StreamOrDevice s /* = {} */) { auto result = argmin(flatten(a, s), 0, true, s); if (keepdims) { @@ -2955,6 +2986,148 @@ array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) { return slice(a_partitioned, slice_starts, slice_ends, s); } +std::vector unique( + const array& a, + int size, + bool return_index /* = false */, + bool return_inverse /* = false */, + bool return_counts /* = false */, + const std::optional& fill_value /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + // Validate args + if (size < 0) { + std::ostringstream msg; + msg << "[unique] Received negative size " << size << "."; + throw std::invalid_argument(msg.str()); + } + if (fill_value && fill_value->size() != 1) { + std::ostringstream msg; + msg << "[unique] Fill value must have one element, but got shape " + << fill_value->shape() << "."; + throw std::invalid_argument(msg.str()); + } + const auto flat = flatten(a, s); + const int n = flat.size(); + + // Handle the edge case of an empty array. + if (n == 0) { + // Without an element there is no default fill value. + if (size > 0 && !fill_value) { + throw std::invalid_argument( + "[unique] A fill value is required for an empty input with a" + " non-zero size."); + } + std::vector out; + out.push_back( + size == 0 ? flat + : full( + {size}, + astype(reshape(*fill_value, {}, s), flat.dtype(), s), + flat.dtype(), + s)); + if (return_index) { + out.push_back(zeros({size}, uint32, s)); + } + if (return_inverse) { + out.push_back(zeros(a.shape(), uint32, s)); + } + if (return_counts) { + out.push_back(zeros({size}, int32, s)); + } + return out; + } + + // Sort the array. The values stay differentiable, the permutation does not. + std::optional order; + if (return_index || return_inverse) { + order = stop_gradient(argsort(flat, 0, s), s); + } + const auto sorted = order ? take(flat, *order, 0, s) : sort(flat, 0, s); + + // Edge detection on the sorted array, true where a new unique + // value starts in the sorted array. + const auto boundary = concatenate( + {array({true}), + not_equal( + slice(sorted, {1}, {n}, s), slice(sorted, {0}, {n - 1}, s), s)}, + 0, + s); + + // Cumsum on boundary gives each sorted element the index of the unique + // value it belongs to. It indexes the output, so it is not differentiable. + const auto group = stop_gradient( + subtract( + cumsum(astype(boundary, uint32, s), 0, false, true, s), + array(1, uint32), + s), + s); + + // Use smallest element of the sorted array (index 0) as padding by default. + const auto fill = fill_value + ? astype(reshape(*fill_value, {}, s), flat.dtype(), s) + : slice(sorted, {0}, {1}, s); + + const int buffer_size = std::max(n, size); + // Scatter positions rather than values in the sorted array (64 bit values + // would fail to scatter on the GPU). + const auto slots = stop_gradient( + slice( + scatter_max( + zeros({buffer_size}, uint32, s), + group, + expand_dims( + where( + boundary, + add(arange(n, uint32, s), array(1, uint32), s), + array(0, uint32), + s), + 1, + s), + 0, + s), + {0}, + {size}, + s), + s); + const auto used = greater(slots, array(0, uint32), s); + const auto positions = + subtract(maximum(slots, array(1, uint32), s), array(1, uint32), s); + + // Build the output arrays. + std::vector out; + out.push_back(where(used, take(sorted, positions, 0, s), fill, s)); + if (return_index) { + // The sort is stable, so the first sorted position of a group holds the + // first occurrence in the input. Unused slots point at position zero, + // which pads with the index of the smallest unique value. + out.push_back(take(*order, positions, 0, s)); + } + if (return_inverse) { + // Clamp so that the indices stay inside a truncated output. Without a + // truncation every group index is already smaller than size. + const auto clamped = + minimum(group, array(std::max(size - 1, 0), uint32), s); + const auto inverse = scatter( + zeros({n}, uint32, s), *order, expand_dims(clamped, 1, s), 0, s); + out.push_back(reshape(inverse, a.shape(), s)); + } + // Padding entries count zero, so counts sum to the input size unless the + // output was truncated. + if (return_counts) { + out.push_back(slice( + scatter_add( + zeros({buffer_size}, int32, s), + group, + ones({n, 1}, int32, s), + 0, + s), + {0}, + {size}, + s)); + } + return out; +} + array logsumexp(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) { std::vector axes(a.ndim()); std::iota(axes.begin(), axes.end(), 0); @@ -3159,27 +3332,24 @@ array floor_divide( const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); - if (issubdtype(dtype, inexact)) { - return floor(divide(a, b, s), s); - } - auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); auto shape = inputs[0].shape(); - auto quotient = array( - shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); - // Fork: unsigned integers never need the sign correction below (truncation - // and floor agree), and it adds six ops to every call. The MoE expert sort - // calls this on uint32 indices in every layer; with the correction a - // streamed-experts decode step hit a Metal GPU timeout. + // Fork: unsigned integers never need the sign correction (truncation and + // floor agree). The MoE expert sort calls this on uint32 indices in every + // layer; a streamed-experts decode step hit a Metal GPU timeout on the + // correction path (see UPSTREAM_MERGE_PLAN.md, floor_divide). if (issubdtype(dtype, unsignedinteger)) { - return quotient; + return array( + shape, + dtype, + std::make_shared(to_stream(s)), + std::move(inputs)); } - auto rem = remainder(a, b, s); - auto zero = array(0, dtype); - auto step = logical_and( - not_equal(rem, zero, s), - not_equal(less(a, zero, s), less(b, zero, s), s)); - return subtract(quotient, astype(step, dtype, s), s); + return array( + shape, + dtype, + std::make_shared(to_stream(s)), + std::move(inputs)); } array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { @@ -4299,12 +4469,13 @@ array logcumsumexp( bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { axis = normalize_axis_index(axis, a.ndim(), "[logcumsumexp] "); + auto out_type = at_least_float(a.dtype()); return array( a.shape(), - a.dtype(), + out_type, std::make_shared( to_stream(s), Scan::ReduceType::LogAddExp, axis, reverse, inclusive), - {a}); + {astype(a, out_type, s)}); } array logcumsumexp( @@ -5035,10 +5206,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()); } @@ -5059,14 +5230,22 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { w_max = astype(w_max, float32, s); w_min = astype(w_min, float32, s); - array mask = greater(abs(w_min, s), abs(w_max, s), s); - array scales = - maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); - scales = where(mask, scales, negative(scales, s), s); - array edge = where(mask, w_min, w_max, s); - array q0 = round(divide(edge, scales, s), s); - scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); - array biases = where(equal(q0, zero, s), zero, edge, s); + array scales(0, float32); + array biases(0, float32); + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scales = maximum(subtract(w_max, w_min, s), eps, s); + biases = w_min; + } else { + array mask = greater(abs(w_min, s), abs(w_max, s), s); + scales = maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); + scales = where(mask, scales, negative(scales, s), s); + array edge = where(mask, w_min, w_max, s); + array q0 = round(divide(edge, scales, s), s); + scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); + biases = where(equal(q0, zero, s), zero, edge, s); + } packed_w = pack_and_quantize(packed_w, scales, biases, bits, s); diff --git a/Source/Cmlx/mlx/mlx/ops.h b/Source/Cmlx/mlx/mlx/ops.h index ab676ca6c..017fcbd70 100644 --- a/Source/Cmlx/mlx/mlx/ops.h +++ b/Source/Cmlx/mlx/mlx/ops.h @@ -786,6 +786,23 @@ min(const array& a, MLX_API array min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); +MLX_API std::tuple +minmax(const array& a, bool keepdims, StreamOrDevice s = {}); +inline std::tuple minmax(const array& a, StreamOrDevice s = {}) { + return minmax(a, false, to_stream(s)); +} + +/** The minimum and maximum of the elements of an array along the given axes. */ +MLX_API std::tuple minmax( + const array& a, + const std::vector& axes, + bool keepdims = false, + StreamOrDevice s = {}); + +/** The minimum and maximum of the elements of an array along the given axis. */ +MLX_API std::tuple +minmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); + /** Returns the Hanning window of size M. */ MLX_API array hanning(int M, StreamOrDevice s = {}); @@ -881,6 +898,27 @@ MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); /** Returns topk elements of the array along a given axis. */ MLX_API array topk(const array& a, int k, int axis, StreamOrDevice s = {}); +/** + * Returns the sorted unique elements of the flattened array, and optionally + * the index of the first occurrence of each, the inverse indices and the + * counts. The output has the given size, and + * is truncated if ``size`` is smaller than the number of unique values. + * If ``size`` is larger, the result is padded with ``fill_value``, or the + * first of the sorted unique values if no fill value is provided. An empty + * input has none, so it throws unless ``size`` is zero or a fill value is + * given. + * A truncated output loses the counts of the values it drops, and its inverse + * indices are clamped to the last entry. + */ +MLX_API std::vector unique( + const array& a, + int size, + bool return_index = false, + bool return_inverse = false, + bool return_counts = false, + const std::optional& fill_value = std::nullopt, + StreamOrDevice s = {}); + /** Cumulative logsumexp of an array. */ MLX_API array logcumsumexp( const array& a, diff --git a/Source/Cmlx/mlx/mlx/primitives.cpp b/Source/Cmlx/mlx/mlx/primitives.cpp index f046670ee..d4035247e 100644 --- a/Source/Cmlx/mlx/mlx/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/primitives.cpp @@ -390,7 +390,10 @@ std::vector ArcCos::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcCos::jvp( @@ -418,7 +421,10 @@ std::vector ArcCosh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcCosh::jvp( @@ -428,6 +434,14 @@ std::vector ArcCosh::jvp( assert(primals.size() == 1); assert(argnums.size() == 1); array one = array(1., primals[0].dtype()); + if (issubdtype(primals[0].dtype(), complexfloating)) { + // rsqrt(x^2 - 1) has the wrong sign for complex x with Re(x) < 0. + array d = multiply( + rsqrt(subtract(primals[0], one, stream()), stream()), + rsqrt(add(primals[0], one, stream()), stream()), + stream()); + return {multiply(tangents[0], d, stream())}; + } array t = subtract(square(primals[0], stream()), one, stream()); return {multiply(tangents[0], rsqrt(t, stream()), stream())}; } @@ -445,7 +459,10 @@ std::vector ArcSin::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcSin::jvp( @@ -472,7 +489,10 @@ std::vector ArcSinh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcSinh::jvp( @@ -499,7 +519,10 @@ std::vector ArcTan::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcTan::jvp( @@ -587,7 +610,10 @@ std::vector ArcTanh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcTanh::jvp( @@ -1477,7 +1503,7 @@ std::vector Convolution::vjp( 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"); + padding_hi[i] = safe_cast(in_size - out_size + padding_lo_[i], "conv"); } // Check for negative padding @@ -1489,13 +1515,21 @@ std::vector Convolution::vjp( has_neg_padding |= (pd < 0); } + // Negative padding is cropped from the output below + std::vector conv_padding_lo = padding_lo; + std::vector conv_padding_hi = padding_hi; + for (int i = 0; i < conv_padding_lo.size(); ++i) { + conv_padding_lo[i] = std::max(conv_padding_lo[i], 0); + conv_padding_hi[i] = std::max(conv_padding_hi[i], 0); + } + auto wt_trans = group_transpose(wt, 0, 1, -1); auto grad = conv_general( /* const array& input = */ cotan, /* const array& weight = */ wt_trans, /* std::vector stride = */ input_dilation_, - /* std::vector padding_lo = */ padding_lo, - /* std::vector padding_hi = */ padding_hi, + /* std::vector padding_lo = */ conv_padding_lo, + /* std::vector padding_hi = */ conv_padding_hi, /* std::vector kernel_dilation = */ kernel_dilation_, /* std::vector input_dilation = */ kernel_strides_, /* int groups = */ groups_, @@ -1693,7 +1727,10 @@ std::vector Cos::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return {jvp(primals, cotangents, argnums)}; + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Cos::jvp( @@ -2192,12 +2229,6 @@ std::vector Flatten::output_shapes(const std::vector& inputs) { return {Flatten::output_shape(inputs[0], start_axis_, end_axis_)}; } -bool FFT::is_equivalent(const Primitive& other) const { - const FFT& r_other = static_cast(other); - return axes_ == r_other.axes_ && inverse_ == r_other.inverse_ && - real_ == r_other.real_; -} - std::vector Unflatten::vjp( const std::vector&, const std::vector& cotangents, @@ -2263,13 +2294,13 @@ std::pair, std::vector> FFT::vmap( // Only the last transformed axis changes size in a real transform if (real_) { auto n = out_shape[fft_axes.back()]; - out_shape[fft_axes.back()] = inverse_ ? 2 * (n - 1) : n / 2 + 1; + out_shape[fft_axes.back()] = inverse_ ? 2 * (n - 1) + odd_out_ : n / 2 + 1; } return { {array( out_shape, real_ && inverse_ ? float32 : complex64, - std::make_shared(stream(), fft_axes, inverse_, real_), + std::make_shared(stream(), fft_axes, inverse_, real_, odd_out_), {in})}, {ax}}; } @@ -2380,6 +2411,12 @@ std::vector FFT::jvp( } } +bool FFT::is_equivalent(const Primitive& other) const { + const FFT& r_other = static_cast(other); + return axes_ == r_other.axes_ && inverse_ == r_other.inverse_ && + real_ == r_other.real_ && odd_out_ == r_other.odd_out_; +} + std::vector Floor::vjp( const std::vector& primals, const std::vector& cotangents, @@ -2405,6 +2442,32 @@ std::pair, std::vector> Floor::vmap( return {{floor(inputs[0], stream())}, axes}; } +std::vector FloorDivide::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 FloorDivide::jvp( + const std::vector& primals, + const std::vector&, + const std::vector&) { + return {zeros_like(primals[0], stream())}; +} + +std::pair, std::vector> FloorDivide::vmap( + const std::vector& inputs, + const std::vector& axes) { + auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream()); + return {{floor_divide(a, b, stream())}, {to_ax}}; +} + std::vector Full::vjp( const std::vector& primals, const std::vector& cotangents, @@ -4443,6 +4506,49 @@ std::vector Scan::jvp( if (reduce_type_ == Scan::Sum) { return {cumsum(tangents[0], axis_, reverse_, inclusive_, stream())}; + } else if (reduce_type_ == Scan::LogAddExp) { + if (issubdtype(tangents[0].dtype(), complexfloating)) { + throw std::invalid_argument( + "[logcumsumexp] JVP is not supported for complex inputs."); + } + auto x = primals[0]; + auto t = tangents[0]; + auto y = logcumsumexp(x, axis_, reverse_, inclusive_, stream()); + + auto zero = zeros({1}, t.dtype(), stream()); + auto t_min = array(finfo(t.dtype()).min, t.dtype()); + auto log_abs_t = log(abs(t, stream()), stream()); + auto log_t_positive = + where(greater(t, zero, stream()), log_abs_t, t_min, stream()); + auto log_t_negative = + where(less(t, zero, stream()), log_abs_t, t_min, stream()); + + auto out_pos = exp(subtract( + logcumsumexp( + add(log_t_positive, x, stream()), + axis_, + reverse_, + inclusive_, + stream()), + y, + stream())); + + auto out_neg = exp(subtract( + logcumsumexp( + add(log_t_negative, x, stream()), + axis_, + reverse_, + inclusive_, + stream()), + y, + stream())); + + return {where( + isneginf(y, stream()), + zeros_like(y, stream()), + subtract(out_pos, out_neg, stream()), + stream())}; + } else { throw std::runtime_error( "JVP is not implemented for cumulative prod/min/max"); @@ -4610,14 +4716,19 @@ std::pair, std::vector> Scatter::vmap( // Clone updates along the vmap dimension so they can be applied to each // source tensor in the vmap. auto& updates = inputs.back(); + int idx_ndim = static_cast(inputs[1].ndim()); + int upd_src_ax = idx_ndim + src_ax; if (vmap_axes.back() < 0) { - updates = expand_dims( - updates, {0, static_cast(inputs[1].ndim())}, stream()); + updates = expand_dims(updates, {0, upd_src_ax}, stream()); updates = repeat(updates, vmap_size, 0, stream()); } else { - updates = - expand_dims(updates, static_cast(inputs[1].ndim()), stream()); - updates = moveaxis(updates, vmap_axes.back(), 0, stream()); + int upd_vmap_ax = vmap_axes.back(); + // expand_dims shifts any axis at or after the insertion point + if (upd_vmap_ax >= upd_src_ax) { + upd_vmap_ax++; + } + updates = expand_dims(updates, upd_src_ax, stream()); + updates = moveaxis(updates, upd_vmap_ax, 0, stream()); } } @@ -5601,16 +5712,17 @@ std::vector Sqrt::vjp( assert(primals.size() == 1); assert(cotangents.size() == 1); auto dtype = primals[0].dtype(); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). if (recip_) { auto one_over_x_root_x = divide(outputs[0], primals[0], stream()); return {multiply( multiply(array(-0.5, dtype), cotangents[0], stream()), - one_over_x_root_x, + conjugate(one_over_x_root_x, stream()), stream())}; } else { return {divide( multiply(array(0.5, dtype), cotangents[0], stream()), - outputs[0], + conjugate(outputs[0], stream()), stream())}; } } @@ -5619,10 +5731,21 @@ std::vector Sqrt::jvp( const std::vector& primals, const std::vector& tangents, const std::vector& argnums) { + assert(primals.size() == 1); + assert(tangents.size() == 1); + auto dtype = primals[0].dtype(); if (recip_) { - return vjp(primals, tangents, argnums, {rsqrt(primals[0], stream())}); + auto one_over_x_root_x = + divide(rsqrt(primals[0], stream()), primals[0], stream()); + return {multiply( + multiply(array(-0.5, dtype), tangents[0], stream()), + one_over_x_root_x, + stream())}; } else { - return vjp(primals, tangents, argnums, {sqrt(primals[0], stream())}); + return {divide( + multiply(array(0.5, dtype), tangents[0], stream()), + sqrt(primals[0], stream()), + stream())}; } } @@ -6261,7 +6384,15 @@ std::pair, std::vector> View::vmap( const std::vector& axes) { assert(inputs.size() == 1); assert(axes.size() == 1); - return {{view(inputs[0], dtype_, stream())}, axes}; + auto a = inputs[0]; + int ax = axes[0]; + // Changing the item size resizes the last axis, so move the vmapped axis + if (ax == static_cast(a.ndim()) - 1 && + size_of(a.dtype()) != size_of(dtype_)) { + a = moveaxis(a, ax, 0, stream()); + ax = 0; + } + return {{view(a, dtype_, stream())}, {ax}}; } const char* View::name() const { diff --git a/Source/Cmlx/mlx/mlx/primitives.h b/Source/Cmlx/mlx/mlx/primitives.h index f5ffc7d86..3ec8be22f 100644 --- a/Source/Cmlx/mlx/mlx/primitives.h +++ b/Source/Cmlx/mlx/mlx/primitives.h @@ -1079,8 +1079,13 @@ class FFT : public UnaryPrimitive { Stream stream, const std::vector& axes, bool inverse, - bool real) - : UnaryPrimitive(stream), axes_(axes), inverse_(inverse), real_(real) {} + bool real, + bool odd_out = false) + : UnaryPrimitive(stream), + axes_(axes), + inverse_(inverse), + real_(real), + odd_out_(odd_out) {} void eval_cpu(const std::vector& inputs, array& out) override; void eval_gpu(const std::vector& inputs, array& out) override; @@ -1091,13 +1096,16 @@ class FFT : public UnaryPrimitive { bool is_equivalent(const Primitive& other) const override; auto state() const { - return std::make_tuple(axes_, inverse_, real_); + return std::make_tuple(axes_, inverse_, real_, odd_out_); } private: std::vector axes_; bool inverse_; bool real_; + // The output length of an inverse real transform is not recoverable from + // the input shape. + bool odd_out_; }; class Flatten : public UnaryPrimitive { @@ -1139,6 +1147,20 @@ class Floor : public UnaryPrimitive { DEFINE_INPUT_OUTPUT_SHAPE() }; +class FloorDivide : public UnaryPrimitive { + public: + explicit FloorDivide(Stream stream) : UnaryPrimitive(stream) {} + + 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(FloorDivide) + DEFINE_DEFAULT_IS_EQUIVALENT() + DEFINE_INPUT_OUTPUT_SHAPE() +}; + class Full : public UnaryPrimitive { public: explicit Full(Stream stream) : UnaryPrimitive(stream) {} diff --git a/Source/Cmlx/mlx/mlx/transforms.cpp b/Source/Cmlx/mlx/mlx/transforms.cpp index 77d81e11e..cad704bd4 100644 --- a/Source/Cmlx/mlx/mlx/transforms.cpp +++ b/Source/Cmlx/mlx/mlx/transforms.cpp @@ -332,6 +332,15 @@ array eval_impl(std::vector outputs, bool async) { } catch (...) { } } + // Commit GPU streams before the wait. + for (auto& s : open_streams) { + if (s.device == Device::gpu) { + try { + gpu::finalize(s); + } catch (...) { + } + } + } for (auto& s : open_streams) { try { synchronize(s); diff --git a/Source/Cmlx/mlx/mlx/types/complex.h b/Source/Cmlx/mlx/mlx/types/complex.h index 65370ca89..ee1af4dfd 100644 --- a/Source/Cmlx/mlx/mlx/types/complex.h +++ b/Source/Cmlx/mlx/mlx/types/complex.h @@ -45,6 +45,11 @@ struct complex64_t : public std::complex { operator float() const { return real(); }; + + // Explicit, so int casts still go through operator float + explicit operator bool() const { + return real() != 0 || imag() != 0; + }; }; inline bool operator>=(const complex64_t& a, const complex64_t& b) { diff --git a/Source/Cmlx/mlx/mlx/version.h b/Source/Cmlx/mlx/mlx/version.h index a403f6f8b..cd08a874b 100644 --- a/Source/Cmlx/mlx/mlx/version.h +++ b/Source/Cmlx/mlx/mlx/version.h @@ -6,7 +6,7 @@ #define MLX_VERSION_MAJOR 0 #define MLX_VERSION_MINOR 32 -#define MLX_VERSION_PATCH 3 +#define MLX_VERSION_PATCH 4 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) diff --git a/Source/Cmlx/mlx/pyproject.toml b/Source/Cmlx/mlx/pyproject.toml index 45675ae85..dd13cfab2 100644 --- a/Source/Cmlx/mlx/pyproject.toml +++ b/Source/Cmlx/mlx/pyproject.toml @@ -5,3 +5,13 @@ requires = [ "typing_extensions", ] build-backend = "setuptools.build_meta" + +[tool.ruff.lint] +select = ["E", "F", "I"] + +[tool.ruff.lint.per-file-ignores] +# These modules expose imports as their public API. +"python/mlx/nn/__init__.py" = ["F401", "F403"] +"python/mlx/nn/layers/__init__.py" = ["F401"] +"python/mlx/optimizers/__init__.py" = ["F403"] +"examples/extensions/mlx_sample_extensions/__init__.py" = ["F401"] diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py index 6289da0c4..6294a67da 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py @@ -48,10 +48,13 @@ def from_file(cls, hostfile): "MLX_METAL_FAST_SYNCH=1" ], "hosts": [ - {"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]}, - {"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]}, + {"ssh": "hostname1", "ips": ["123.123.123.1"], + "rdma": [null, "rdma_en2", "rdma_en3"]}, + {"ssh": "hostname2", "ips": ["123.123.123.2"], + "rdma": ["rdma_en2", null, "rdma_en3"]}, ... - {"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]}, + {"ssh": "hostnameN", "ips": ["123.123.123.N"], + "rdma": ["rdma_en2", "rdma_en3", null]}, ] } diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py index a4dabed62..ce595a49d 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py @@ -10,8 +10,6 @@ from subprocess import DEVNULL, run from typing import Optional -import mlx.core as mx - from .common import ( Host, Hostfile, @@ -112,7 +110,8 @@ def can_auto_setup(hosts, sshinfo, auto_setup=False): has_sudo = all(info.has_sudo for info in sshinfo) if not has_sudo and auto_setup: log_warning( - "Automatic setup requested but the following hosts do not have passwordless sudo" + "Automatic setup requested but the following hosts do not have " + "passwordless sudo" ) for h, i in zip(hosts, sshinfo): if not i.has_sudo: @@ -188,11 +187,11 @@ def setup(self, verbose=False, auto_setup=False): def parse_hardware_ports(ports_string): ports = {} port_name = None - for l in ports_string.decode("utf-8").split("\n"): - if l.startswith("Hardware Port:"): - port_name = l.strip()[15:] - elif l.startswith("Device:"): - ports[port_name] = l.strip()[8:] + for line in ports_string.decode("utf-8").split("\n"): + if line.startswith("Hardware Port:"): + port_name = line.strip()[15:] + elif line.startswith("Device:"): + ports[port_name] = line.strip()[8:] port_name = None return ports @@ -343,7 +342,8 @@ def check_valid_mesh(hosts, connectivity, strict=True): if connectivity[i][j] <= 0: if strict: log_error( - f"Incomplete mesh, {hosts[i].ssh_hostname} is not connected to {hosts[j].ssh_hostname}" + f"Incomplete mesh, {hosts[i].ssh_hostname} is not connected to " + f"{hosts[j].ssh_hostname}" ) log_error() log_error("Try passing --dot to visualize the connectivity") @@ -432,7 +432,7 @@ def _check(hostname, i): def prepare_ethernet_hostfile(args, hosts): - log(args.verbose, f"Preparing an ethernet hostfile") + log(args.verbose, "Preparing an ethernet hostfile") add_ips(hosts, args.verbose) hostfile = Hostfile( @@ -466,14 +466,19 @@ def configure_jaccl(args, hosts, ips, sshinfo): log(args.verbose, "Prepare a jaccl hostfile") add_ips(hosts, args.verbose) + num_nodes = len(hosts) + count = min( + len(ips.ips[i, j]) for i in range(num_nodes) for j in range(num_nodes) if i != j + ) jaccl_hosts = [] for i, h in enumerate(hosts): rdma = [] - for j in range(len(hosts)): + for j in range(num_nodes): if i == j: rdma.append(None) else: - rdma.append(f"rdma_{ips.ips[i, j][0][0]}") + devices = [f"rdma_{ips.ips[i, j][c][0]}" for c in range(count)] + rdma.append(devices[0] if count == 1 else devices) jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdma)) hostfile = Hostfile(jaccl_hosts, "jaccl", args.env) @@ -513,7 +518,7 @@ def configure_jaccl_ring(args, hosts, ips, ring, sshinfo): def prepare_tb_hostfile(args, hosts, sshinfo): - log(args.verbose, f"Preparing for communication over thunderbolt") + log(args.verbose, "Preparing for communication over thunderbolt") tb_hosts, uuid_reverse_index = extract_connectivity(hosts, args.verbose) if args.dot: diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py index 0e661e358..73ff5f1b3 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py @@ -1,9 +1,9 @@ # Copyright © 2025 Apple Inc. import argparse -import base64 import json import os +import platform import shlex import shutil import sys @@ -193,27 +193,40 @@ def _thread_fn(rank, *args, **kwargs): to_write = [p.stdin.fileno()] stdin_buffer = b"" - while p.poll() is None: + while True: + exited = p.poll() is not None try: stdin_buffer += stdin_queue.get_nowait() except QueueEmpty: pass - rlist, wlist, _ = select(to_read, to_write, [], 1.0) + rlist, wlist, _ = select(to_read, to_write, [], 0 if exited else 1.0) for fd in rlist: is_stdout = fd == p.stdout.fileno() - msg = os.read(fd, 8192).decode(errors="ignore") - msg = command.preprocess_output(msg, is_stdout) + try: + data = os.read(fd, 8192) + except BlockingIOError: + continue + if not data: + to_read.remove(fd) + continue + msg = command.preprocess_output(data.decode(errors="ignore"), is_stdout) if is_stdout: stdout_queue.put(msg.encode()) else: stderr_queue.put(msg.encode()) for fd in wlist: if len(stdin_buffer) > 0: - n = os.write(fd, stdin_buffer) + try: + n = os.write(fd, stdin_buffer) + except BrokenPipeError: + to_write.remove(fd) + continue stdin_buffer = stdin_buffer[n:] - if stop: - command.terminate() + if exited and not rlist: break + if stop and not exited: + command.terminate() + exit_codes[rank] = command.exit_status if exit_codes[rank][1]: @@ -440,7 +453,7 @@ def get_mpi_libname(): filter(lambda line: "libmpi" in line, otool_output.splitlines()) ) return libmpi_line.strip().split()[0].removeprefix("@rpath/") - except: + except BaseException: return None @@ -592,5 +605,6 @@ def main(): launch_jaccl(parser, hostfile.hosts, args, rest) else: parser.error( - "The backend should be one of {'ring', 'mpi', 'nccl', 'jaccl', 'jaccl-ring'}" + "The backend should be one of " + "{'ring', 'mpi', 'nccl', 'jaccl', 'jaccl-ring'}" ) diff --git a/Source/Cmlx/mlx/python/mlx/extension.py b/Source/Cmlx/mlx/python/mlx/extension.py index d794fbe28..96dc09013 100644 --- a/Source/Cmlx/mlx/python/mlx/extension.py +++ b/Source/Cmlx/mlx/python/mlx/extension.py @@ -1,10 +1,14 @@ # Copyright © 2023 Apple Inc. +import logging import os +import platform import re +import shlex import subprocess import sys from pathlib import Path +from typing import Sequence from setuptools import Extension from setuptools.command.build_ext import build_ext @@ -12,6 +16,572 @@ import mlx _MLX_PATH = str(mlx.__path__[0]) +_MLX_PACKAGE_PATH = str(Path(__file__).resolve().parent) +_HOST_SOURCE_SUFFIXES = {".cc", ".cpp", ".cxx"} +_METAL_SOURCE_SUFFIX = ".metal" +_MIN_MACOS_DEPLOYMENT_TARGET = "14.0" +_LOGGER = logging.getLogger(__name__) + + +def _is_ninja_available() -> bool: + try: + subprocess.check_output(["ninja", "--version"]) + except (OSError, subprocess.SubprocessError): + return False + return True + + +def _cmake_quote(value: os.PathLike[str] | str) -> str: + value = os.fspath(value) + delimiter = "=" + while f"]{delimiter}]" in value: + delimiter += "=" + return f"[{delimiter}[{value}]{delimiter}]" + + +def _cmake_arguments(values: Sequence[os.PathLike[str] | str]) -> list[str]: + return [f" {_cmake_quote(value)}" for value in values] + + +def _extension_paths(command: build_ext, build_py, ext: Extension) -> tuple[Path, Path]: + fullname = command.get_ext_fullname(ext.name) + filename = command.get_ext_filename(fullname) + package = ".".join(fullname.split(".")[:-1]) + package_dir = build_py.get_package_dir(package) + inplace_file = Path(package_dir) / Path(filename).name + regular_file = Path(command.build_lib) / filename + return inplace_file, regular_file + + +def _macos_deployment_target(target: str) -> str: + if not re.fullmatch(r"\d+(?:\.\d+)*", target): + raise ValueError(f"Invalid macOS deployment target: {target!r}.") + if int(target.partition(".")[0]) < 14: + raise ValueError( + "macOS deployment target must be at least " + f"{_MIN_MACOS_DEPLOYMENT_TARGET}, got {target!r}." + ) + return target + + +def _macos_deployment_target_key(target: str) -> tuple[int, ...]: + components = [int(component) for component in target.split(".")] + while len(components) > 1 and components[-1] == 0: + components.pop() + return tuple(components) + + +def _mlx_macos_deployment_target() -> str: + libmlx = Path(_MLX_PACKAGE_PATH) / "lib" / "libmlx.dylib" + try: + output = subprocess.check_output( + ["xcrun", "vtool", "-show-build", str(libmlx)], + stderr=subprocess.STDOUT, + text=True, + ) + except (OSError, subprocess.SubprocessError) as error: + raise RuntimeError( + f"Failed to read the macOS deployment target from {libmlx}." + ) from error + + targets = re.findall(r"^\s*minos\s+(\d+(?:\.\d+)*)\s*$", output, re.MULTILINE) + if not targets: + raise RuntimeError(f"Could not find a macOS deployment target in {libmlx}.") + return max(targets, key=_macos_deployment_target_key) + + +def _cmake_generator( + extra_cmake_args: Sequence[str], use_ninja: bool +) -> tuple[str, bool]: + generator = None + generator_is_configured = False + for index, argument in enumerate(extra_cmake_args): + if argument in ("-G", "--generator"): + generator_is_configured = True + if index + 1 < len(extra_cmake_args): + generator = extra_cmake_args[index + 1] + elif argument.startswith("-G") and argument != "-G": + generator_is_configured = True + generator = argument[2:] + elif argument.startswith("--generator="): + generator_is_configured = True + generator = argument.partition("=")[2] + + if generator_is_configured: + return generator or "explicit", False + if use_ninja: + return "Ninja", True + return os.environ.get("CMAKE_GENERATOR") or "default", False + + +def _cmake_generator_build_directory(generator: str) -> str: + generator = re.sub(r"[^a-z0-9]+", "_", generator.lower()).strip("_") + return f"build-{generator or 'default'}" + + +class MetalLibrary: + """A named Metal library built as part of a :class:`MetalExtension`.""" + + def __init__( + self, + name: str, + sources: Sequence[os.PathLike[str] | str], + *, + extra_compile_args: Sequence[str] = (), + deployment_target: str | None = None, + ) -> None: + if not name or Path(name).name != name: + raise ValueError(f"Invalid Metal library name: {name!r}.") + source_strings = [os.fspath(source) for source in sources] + if not source_strings: + raise ValueError("MetalLibrary requires at least one Metal source file.") + if any( + Path(source).suffix.lower() != _METAL_SOURCE_SUFFIX + for source in source_strings + ): + raise ValueError("MetalLibrary sources must be Metal source files.") + if isinstance(extra_compile_args, (str, bytes)): + raise TypeError("MetalLibrary extra_compile_args must be a sequence.") + + self.name = name + self.sources = source_strings + self.extra_compile_args = list(extra_compile_args) + self.deployment_target = ( + _macos_deployment_target(deployment_target) + if deployment_target is not None + else None + ) + + +class MetalExtension(Extension): + """A setuptools extension built from C++ and Metal sources.""" + + def __init__( + self, + name: str, + sources: Sequence[os.PathLike[str] | str], + *args, + metal_libraries: Sequence[MetalLibrary] | None = None, + **kwargs, + ) -> None: + source_strings = [os.fspath(source) for source in sources] + suffixes = {Path(source).suffix.lower() for source in source_strings} + supported_suffixes = _HOST_SOURCE_SUFFIXES | {_METAL_SOURCE_SUFFIX} + unsupported_sources = [ + source + for source in source_strings + if Path(source).suffix.lower() not in supported_suffixes + ] + if unsupported_sources: + raise ValueError( + "MetalExtension received unsupported source files: " + + ", ".join(unsupported_sources) + ) + if not suffixes.intersection(_HOST_SOURCE_SUFFIXES): + raise ValueError("MetalExtension requires at least one C++ source file.") + if metal_libraries is None and _METAL_SOURCE_SUFFIX not in suffixes: + raise ValueError("MetalExtension requires at least one Metal source file.") + if metal_libraries is not None and _METAL_SOURCE_SUFFIX in suffixes: + raise ValueError( + "Metal sources must be specified in metal_libraries when it is used." + ) + if kwargs.get("py_limited_api", False): + raise ValueError("MetalExtension does not support py_limited_api.") + + compile_args = kwargs.get("extra_compile_args") + if compile_args is None: + compile_args = {"cxx": [], "metal": []} + elif isinstance(compile_args, (list, tuple)): + compile_args = {"cxx": list(compile_args), "metal": []} + elif isinstance(compile_args, dict): + unknown_groups = set(compile_args).difference({"cxx", "metal"}) + if unknown_groups: + raise ValueError( + "extra_compile_args only supports the cxx and metal groups." + ) + if any( + not isinstance(compile_args.get(group, []), (list, tuple)) + for group in ("cxx", "metal") + ): + raise TypeError("cxx and metal compile arguments must be lists.") + compile_args = { + "cxx": list(compile_args.get("cxx", [])), + "metal": list(compile_args.get("metal", [])), + } + else: + raise TypeError("extra_compile_args must be a list or a dictionary.") + + kwargs["extra_compile_args"] = compile_args + kwargs.setdefault("language", "c++") + super().__init__(name, sources, *args, **kwargs) + self.metal_library_name = name.rsplit(".", 1)[-1] + if metal_libraries is None: + self.metal_libraries = [ + MetalLibrary( + self.metal_library_name, + [ + source + for source in source_strings + if Path(source).suffix.lower() == _METAL_SOURCE_SUFFIX + ], + ) + ] + else: + self.metal_libraries = list(metal_libraries) + if not self.metal_libraries: + raise ValueError("metal_libraries must contain at least one library.") + if not all( + isinstance(library, MetalLibrary) for library in self.metal_libraries + ): + raise TypeError("metal_libraries must contain MetalLibrary objects.") + names = [library.name for library in self.metal_libraries] + if len(names) != len(set(names)): + raise ValueError("Metal library names must be unique.") + + +def _metal_extension_cmake( + ext: MetalExtension, output_dir: Path, *, generate_stubs: bool = True +) -> str: + host_sources = [ + str(Path(source).resolve()) + for source in ext.sources + if Path(source).suffix.lower() in _HOST_SOURCE_SUFFIXES + ] + include_dirs = [str(Path(path).resolve()) for path in ext.include_dirs] + depends = [str(Path(path).resolve()) for path in ext.depends] + library_dirs = [str(Path(path).resolve()) for path in ext.library_dirs] + extra_objects = [str(Path(path).resolve()) for path in ext.extra_objects] + languages = ["CXX"] + + lines = [ + "cmake_minimum_required(VERSION 3.25)", + "", + f"project(mlx_extension LANGUAGES {' '.join(languages)})", + "", + "set(CMAKE_CXX_STANDARD 20)", + "set(CMAKE_CXX_STANDARD_REQUIRED ON)", + "set(CMAKE_POSITION_INDEPENDENT_CODE ON)", + "", + "find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED)", + "if(NOT nanobind_ROOT)", + " execute_process(", + ' COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir', + " OUTPUT_STRIP_TRAILING_WHITESPACE", + " OUTPUT_VARIABLE nanobind_ROOT", + " COMMAND_ERROR_IS_FATAL ANY)", + "endif()", + "find_package(nanobind CONFIG REQUIRED)", + "if(NOT MLX_ROOT)", + f" set(MLX_ROOT {_cmake_quote(_MLX_PACKAGE_PATH)})", + "endif()", + "find_package(MLX CONFIG REQUIRED)", + "if(NOT MLX_BUILD_METAL)", + ' message(FATAL_ERROR "MetalExtension requires an MLX build ' + 'with Metal support.")', + "endif()", + "", + "nanobind_add_module(", + " mlx_extension", + " NB_STATIC", + " LTO", + " NOMINSIZE", + " NB_DOMAIN", + " mlx", + *_cmake_arguments(host_sources), + ")", + "set_target_properties(mlx_extension PROPERTIES OUTPUT_NAME " + f"{_cmake_quote(ext.metal_library_name)})", + "set_target_properties(", + " mlx_extension", + " PROPERTIES", + f" LIBRARY_OUTPUT_DIRECTORY {_cmake_quote(output_dir)}", + f" LIBRARY_OUTPUT_DIRECTORY_RELEASE {_cmake_quote(output_dir)}", + f" LIBRARY_OUTPUT_DIRECTORY_DEBUG {_cmake_quote(output_dir)}", + f" LIBRARY_OUTPUT_DIRECTORY_RELWITHDEBINFO {_cmake_quote(output_dir)}", + f" LIBRARY_OUTPUT_DIRECTORY_MINSIZEREL {_cmake_quote(output_dir)})", + "target_link_libraries(mlx_extension PRIVATE mlx)", + ] + + if include_dirs: + lines.extend( + [ + "target_include_directories(mlx_extension PRIVATE", + *_cmake_arguments(include_dirs), + ")", + ] + ) + + definitions = [ + f"MLX_EXTENSION_NAME={ext.metal_library_name}", + f'MLX_METAL_LIBRARY_NAME="{ext.metal_library_name}"', + *( + name if value is None else f"{name}={value}" + for name, value in ext.define_macros + ), + ] + if definitions: + lines.extend( + [ + "target_compile_definitions(mlx_extension PRIVATE", + *_cmake_arguments(definitions), + ")", + ] + ) + + cxx_compile_args = list(ext.extra_compile_args["cxx"]) # type: ignore + cxx_compile_args.extend(f"-U{name}" for name in ext.undef_macros) + if cxx_compile_args: + lines.extend( + [ + "target_compile_options(mlx_extension PRIVATE", + *_cmake_arguments(cxx_compile_args), + ")", + ] + ) + + if library_dirs: + lines.extend( + [ + "target_link_directories(mlx_extension PRIVATE", + *_cmake_arguments(library_dirs), + ")", + ] + ) + + link_libraries = [*ext.libraries, *extra_objects] + if link_libraries: + lines.extend( + [ + "target_link_libraries(mlx_extension PRIVATE", + *_cmake_arguments(link_libraries), + ")", + ] + ) + + if ext.extra_link_args: + lines.extend( + [ + "target_link_options(mlx_extension PRIVATE", + *_cmake_arguments(ext.extra_link_args), + ")", + ] + ) + + mlx_rpath = "@loader_path/" + "../" * ext.name.count(".") + "mlx/lib" + rpaths = ["@loader_path", mlx_rpath, *ext.runtime_library_dirs] + lines.extend( + [ + "set_target_properties(", + " mlx_extension", + " PROPERTIES", + " BUILD_WITH_INSTALL_RPATH TRUE", + f" INSTALL_RPATH {_cmake_quote(';'.join(rpaths))})", + "", + ] + ) + metallib_targets = [] + for index, library in enumerate(ext.metal_libraries): + target = ( + "mlx_extension_metallib" + if len(ext.metal_libraries) == 1 + else f"mlx_extension_metallib_{index}" + ) + metallib_targets.append(target) + lines.extend( + [ + "mlx_build_metallib(", + f" TARGET {target}", + f" TITLE {_cmake_quote(library.name)}", + " SOURCES", + *_cmake_arguments( + list(str(Path(source).resolve()) for source in library.sources) + ), + " INCLUDE_DIRS", + *_cmake_arguments(include_dirs), + " ${MLX_INCLUDE_DIRS}", + " DEPS", + *_cmake_arguments(depends), + ] + ) + compile_options = [ + *ext.extra_compile_args["metal"], # type: ignore + *library.extra_compile_args, + ] + if compile_options: + lines.extend( + [ + " COMPILE_OPTIONS", + *_cmake_arguments(compile_options), + ] + ) + if library.deployment_target is not None: + lines.append( + f" DEPLOYMENT_TARGET {_cmake_quote(library.deployment_target)}" + ) + lines.extend( + [ + f" OUTPUT_DIRECTORY {_cmake_quote(output_dir)})", + "", + ] + ) + lines.extend( + [ + f"add_dependencies(mlx_extension {' '.join(metallib_targets)})", + "", + ] + ) + if generate_stubs: + stub_path = output_dir / f"{ext.metal_library_name}.pyi" + stubgen_code = "import mlx.core; from nanobind.stubgen import main; main()" + lines.extend( + [ + "add_custom_command(", + f" OUTPUT {_cmake_quote(stub_path)}", + ' COMMAND "${Python_EXECUTABLE}"', + " -c", + f" {_cmake_quote(stubgen_code)}", + " -q", + " -m", + f" {_cmake_quote(ext.metal_library_name)}", + " -i", + f" {_cmake_quote(output_dir)}", + " -o", + f" {_cmake_quote(stub_path)}", + " DEPENDS mlx_extension", + f' COMMENT "Generating {ext.metal_library_name}.pyi"', + " VERBATIM)", + "add_custom_target(mlx_extension_stub ALL DEPENDS " + f"{_cmake_quote(stub_path)})", + "", + ] + ) + return "\n".join(lines) + + +class BuildExtension(build_ext): + """Build :class:`MetalExtension` modules with CMake and Metal. + + CMake uses the platform default generator. Pass ``use_ninja=True`` with + :meth:`with_options` to use Ninja instead. If Ninja is requested but + unavailable, the build warns and falls back to the platform default generator. + Nanobind stubs are generated by default. Pass ``generate_stubs=False`` to + disable stub generation. + """ + + @classmethod + def with_options(cls, **options): + """Return a command subclass initialized with the given options.""" + + class cls_with_options(cls): # type: ignore + def __init__(self, *args, **kwargs) -> None: + kwargs.update(options) + super().__init__(*args, **kwargs) + + return cls_with_options + + def __init__(self, *args, **kwargs) -> None: + self.use_ninja = kwargs.pop("use_ninja", False) + self.generate_stubs = kwargs.pop("generate_stubs", True) + super().__init__(*args, **kwargs) + if self.use_ninja and not _is_ninja_available(): + _LOGGER.warning( + "Attempted to use Ninja as the BuildExtension CMake generator, " + "but Ninja was not found. Falling back to CMake's platform " + "default generator." + ) + self.use_ninja = False + + def build_extension(self, ext: Extension) -> None: + if not isinstance(ext, MetalExtension): + super().build_extension(ext) + return + if platform.system() != "Darwin": + raise RuntimeError("MetalExtension is only supported on macOS.") + + ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] + output_dir = ext_fullpath.parent.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + extra_cmake_args = shlex.split(os.environ.get("CMAKE_ARGS", "")) + generator, add_ninja_generator = _cmake_generator( + extra_cmake_args, self.use_ninja + ) + build_root = Path(self.build_temp) / ext.name + source_dir = build_root / "source" + binary_dir = build_root / _cmake_generator_build_directory(generator) + source_dir.mkdir(parents=True, exist_ok=True) + binary_dir.mkdir(parents=True, exist_ok=True) + (source_dir / "CMakeLists.txt").write_text( + _metal_extension_cmake(ext, output_dir, generate_stubs=self.generate_stubs), + encoding="utf-8", + ) + + debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug + cfg = "Debug" if debug else "Release" + cmake_args = [ + "-S", + str(source_dir), + "-B", + str(binary_dir), + f"-DCMAKE_BUILD_TYPE={cfg}", + f"-DPython_EXECUTABLE={sys.executable}", + "-DBUILD_SHARED_LIBS=ON", + f"-DMLX_METAL_DEBUG={'ON' if debug else 'OFF'}", + ] + + deployment_target = os.environ.get("MACOSX_DEPLOYMENT_TARGET") + if deployment_target is None: + deployment_target = _mlx_macos_deployment_target() + else: + deployment_target = _macos_deployment_target(deployment_target) + cmake_args.append(f"-DCMAKE_OSX_DEPLOYMENT_TARGET={deployment_target}") + + deployment_target_override = None + for argument in extra_cmake_args: + match = re.fullmatch( + r"(?:-D)?CMAKE_OSX_DEPLOYMENT_TARGET(?::[^=]+)?=(.*)", argument + ) + if match: + deployment_target_override = match.group(1) + if deployment_target_override is not None: + _macos_deployment_target(deployment_target_override) + if add_ninja_generator: + cmake_args.extend(["-G", "Ninja"]) + cmake_args.extend(extra_cmake_args) + + archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", "")) + if archs: + cmake_args.append(f"-DCMAKE_OSX_ARCHITECTURES={';'.join(archs)}") + + build_args = ["--build", str(binary_dir), "--config", cfg] + if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ: + jobs = self.parallel + if not isinstance(jobs, int) or isinstance(jobs, bool): + jobs = os.cpu_count() + build_args.append(f"-j{jobs}") + + subprocess.run(["cmake", *cmake_args], check=True) + subprocess.run(["cmake", *build_args], check=True) + + def run(self) -> None: + super().run() + if not self.inplace: + return + + build_py = self.get_finalized_command("build_py") + for ext in self.extensions: + if not isinstance(ext, MetalExtension): + continue + inplace_file, regular_file = _extension_paths(self, build_py, ext) + sidecar_names = [ + f"{library.name}.metallib" for library in ext.metal_libraries + ] + if self.generate_stubs: + sidecar_names.append(f"{ext.metal_library_name}.pyi") + for name in sidecar_names: + regular_sidecar = regular_file.parent / name + inplace_sidecar = inplace_file.parent / regular_sidecar.name + self.copy_file(str(regular_sidecar), str(inplace_sidecar)) # A CMakeExtension needs a sourcedir instead of a file list. @@ -40,6 +610,18 @@ def build_extension(self, ext: CMakeExtension) -> None: f"-DPython_EXECUTABLE={sys.executable}", ] build_args = [] + loader_path = None + if sys.platform.startswith("darwin"): + loader_path = "@loader_path" + elif sys.platform.startswith("linux"): + loader_path = "$ORIGIN" + if loader_path is not None: + mlx_rpath = loader_path + "/" + "../" * ext.name.count(".") + "mlx/lib" + cmake_args += [ + "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON", + f"-DCMAKE_INSTALL_RPATH={loader_path};{mlx_rpath}", + ] + # Adding CMake arguments set as environment variable # (needed e.g. to build for ARM OSx on conda-forge) if "CMAKE_ARGS" in os.environ: @@ -79,11 +661,9 @@ def run(self) -> None: if isinstance(ext, CMakeExtension): # Resolve inplace package dir build_py = self.get_finalized_command("build_py") - inplace_file, regular_file = self._get_inplace_equivalent( - build_py, ext - ) + inplace_file, regular_file = _extension_paths(self, build_py, ext) - inplace_dir = str(Path(inplace_file).parent.resolve()) - regular_dir = str(Path(regular_file).parent.resolve()) + inplace_dir = str(inplace_file.parent.resolve()) + regular_dir = str(regular_file.parent.resolve()) self.copy_tree(regular_dir, inplace_dir) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py b/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py index 305b5ec7c..1511c129e 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py @@ -180,7 +180,8 @@ def gelu_approx(x): .. math:: - x = 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) + x = 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * \left(x + 0.044715 * + x^3\right)\right)\right) """ return 0.5 * x * (1 + mx.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x**3))) @@ -557,7 +558,8 @@ class GELU(Module): However, if ``approx`` is set to 'precise' or 'fast' it applies .. math:: - \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) \\ + \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * + \left(x + 0.044715 * x^3\right)\right)\right) \\ \textrm{GELUFast}(x) &= x * \sigma\left(1.702 * x\right) respectively. diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/base.py b/Source/Cmlx/mlx/python/mlx/nn/layers/base.py index e980f8343..820afb769 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/base.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/base.py @@ -115,7 +115,7 @@ def __setattr__(self, key: str, val: Any): self.pop(key, None) def __delattr__(self, name): - if (val := self.get(name, None)) is not None: + if self.get(name, None) is not None: del self[name] else: super().__delattr__(name) @@ -208,7 +208,8 @@ def load_weights( def save_weights(self, file: str): """ - Save the model's weights to a file. The saving method is determined by the file extension: + Save the model's weights to a file. The saving method is determined by the file + extension: - ``.npz`` will use :func:`mx.savez` - ``.safetensors`` will use :func:`mx.save_safetensors` """ @@ -306,8 +307,8 @@ def update(self, parameters: dict, strict: bool = True) -> Module: dict of dicts and lists. Commonly used by the optimizer to change the model to the updated - (optimized) parameters. Also used by the :meth:`mlx.nn.value_and_grad` to set the - tracers in the model in order to compute gradients. + (optimized) parameters. Also used by the :meth:`mlx.nn.value_and_grad` to set + the tracers in the model in order to compute gradients. The passed in parameters dictionary need not be a full dictionary similar to :meth:`parameters`. Only the provided locations will be @@ -331,7 +332,8 @@ def apply(dst, parameters): if isinstance(current_value, mx.array): if strict and not isinstance(new_value, mx.array): raise ValueError( - f"Received invalid type: {type(new_value).__name__}." + f"Received invalid type: " + f"{type(new_value).__name__}." ) dst[k] = new_value else: @@ -478,8 +480,8 @@ def freeze( keys: Optional[Union[str, List[str]]] = None, strict: bool = False, ) -> Module: - """Freeze the Module's parameters or some of them. Freezing a parameter means not - computing gradients for it. + """Freeze the Module's parameters or some of them. Freezing a parameter means + not computing gradients for it. This function is idempotent i.e. freezing a frozen model is a no-op. @@ -490,7 +492,9 @@ def freeze( model = nn.Transformer() model.freeze() - model.apply_to_modules(lambda k, v: v.unfreeze() if k.endswith("attention") else None) + model.apply_to_modules( + lambda k, v: v.unfreeze() if k.endswith("attention") else None + ) Args: recurse (bool, optional): If True then freeze the parameters of the @@ -499,8 +503,8 @@ def freeze( parameters will be frozen otherwise all the parameters of a module. For instance freeze all biases by calling ``module.freeze(keys="bias")``. - strict (bool, optional): If set to ``True`` validate that the passed keys exist. - Default: ``False``. + strict (bool, optional): If set to ``True`` validate that the passed keys + exist. Default: ``False``. Returns: The module instance after freezing the parameters. @@ -511,8 +515,10 @@ def _freeze_impl(_, m): if local_keys is None: local_keys = tree_flatten( m.filter_and_map( - lambda m, k, v: (not isinstance(v, Module)) - and m.valid_parameter_filter(m, k, v) + lambda m, k, v: ( + (not isinstance(v, Module)) + and m.valid_parameter_filter(m, k, v) + ) ) ) local_keys = [k for (k, v) in local_keys] @@ -558,8 +564,8 @@ def unfreeze( parameters will be unfrozen otherwise all the parameters of a module. For instance unfreeze all biases by calling ``module.unfreeze(keys="bias")``. - strict (bool, optional): If set to ``True`` validate that the passed keys exist. - Default: ``False``. + strict (bool, optional): If set to ``True`` validate that the passed keys + exist. Default: ``False``. Returns: The module instance after unfreezing the parameters. @@ -627,7 +633,9 @@ def set_dtype( parameters to the new dtype. """ if predicate is None: - predicate = lambda _: True + + def predicate(_): + return True self.apply(lambda x: x.astype(dtype) if predicate(x.dtype) else x) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py index 40f34ff56..9241e2c14 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py @@ -51,7 +51,7 @@ def __init__( f"divisible by the number of groups ({groups})" ) - scale = math.sqrt(1 / (in_channels * kernel_size)) + scale = math.sqrt(1 / (in_channels // groups * kernel_size)) self.weight = mx.random.uniform( low=-scale, high=scale, @@ -132,7 +132,7 @@ def __init__( lambda x: (x, x) if isinstance(x, int) else x, (kernel_size, stride, padding), ) - scale = math.sqrt(1 / (in_channels * kernel_size[0] * kernel_size[1])) + scale = math.sqrt(1 / (in_channels // groups * kernel_size[0] * kernel_size[1])) self.weight = mx.random.uniform( low=-scale, high=scale, diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py index 1fe8935fc..00f3ae0c3 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py @@ -8,7 +8,8 @@ class ConvTranspose1d(Module): - """Applies a 1-dimensional transposed convolution over the multi-channel input sequence. + """Applies a 1-dimensional transposed convolution over the multi-channel input + sequence. The channels are expected to be last i.e. the input shape should be ``NLC`` where: @@ -83,7 +84,8 @@ def __call__(self, x): class ConvTranspose2d(Module): - """Applies a 2-dimensional transposed convolution over the multi-channel input image. + """Applies a 2-dimensional transposed convolution over the multi-channel input + image. The channels are expected to be last i.e. the input shape should be ``NHWC`` where: @@ -163,7 +165,8 @@ def __call__(self, x): class ConvTranspose3d(Module): - """Applies a 3-dimensional transposed convolution over the multi-channel input image. + """Applies a 3-dimensional transposed convolution over the multi-channel input + image. The channels are expected to be last i.e. the input shape should be ``NDHWC`` where: diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py b/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py index 16979b509..23f521c97 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py @@ -7,7 +7,6 @@ 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_flatten, tree_map_with_path, tree_unflatten diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py b/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py index ef74ba9b6..2c07a5db1 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py @@ -24,7 +24,7 @@ def __init__(self, p: float = 0.5): self._p_1 = 1 - p def _extra_repr(self) -> str: - return f"p={1-self._p_1}" + return f"p={1 - self._p_1}" def __call__(self, x: mx.array) -> mx.array: if self._p_1 == 1 or not self.training: @@ -68,7 +68,7 @@ def __init__(self, p: float = 0.5): self._p_1 = 1 - p def _extra_repr(self) -> str: - return f"p={1-self._p_1}" + return f"p={1 - self._p_1}" def __call__(self, x: mx.array) -> mx.array: if x.ndim not in (3, 4): @@ -118,7 +118,7 @@ def __init__(self, p: float = 0.5): self._p_1 = 1 - p def _extra_repr(self) -> str: - return f"p={1-self._p_1}" + return f"p={1 - self._p_1}" def __call__(self, x: mx.array) -> mx.array: if x.ndim not in (4, 5): diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py b/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py index e05cfb5f0..f3212f934 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py @@ -47,7 +47,9 @@ def to_quantized( mode: str = "affine", quantize_input: bool = False, ): - """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.""" + """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding + layer. + """ if quantize_input: raise ValueError("Quantized input is not supported.") return QuantizedEmbedding.from_embedding(self, group_size, bits, mode) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py b/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py index a1da07c46..ae04c921b 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py @@ -32,9 +32,11 @@ class Linear(Module): y = x W^\top + b - where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``. + 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})`, + The values are initialized from the uniform distribution + :math:`\mathcal{U}(-{k}, {k})`, where :math:`k = \frac{1}{\sqrt{D_i}}` and :math:`D_i` is equal to ``input_dims``. Args: @@ -60,7 +62,10 @@ def __init__(self, input_dims: int, output_dims: int, bias: bool = True) -> None ) def _extra_repr(self) -> str: - return f"input_dims={self.weight.shape[1]}, output_dims={self.weight.shape[0]}, bias={'bias' in self}" + return ( + f"input_dims={self.weight.shape[1]}, " + f"output_dims={self.weight.shape[0]}, bias={'bias' in self}" + ) def __call__(self, x: mx.array) -> mx.array: if "bias" in self: @@ -101,7 +106,8 @@ def to_quantized( if quantize_input: if mode not in ["nvfp4", "mxfp8"]: raise ValueError( - f"Quantized activations are only supported for 'nvfp4' and 'mxfp8' modes, got {mode}." + f"Quantized activations are only supported for 'nvfp4' and 'mxfp8' " + f"modes, got {mode}." ) return QQLinear.from_linear(self, group_size, bits, mode) return QuantizedLinear.from_linear(self, group_size, bits, mode) @@ -117,10 +123,12 @@ class Bilinear(Module): y_i = x_2^\top W_i x_1 + b_i where: - :math:`W` has shape ``[output_dims, input2_dims, input1_dims]``, :math:`b` has shape ``[output_dims]``, + :math:`W` has shape ``[output_dims, input2_dims, input1_dims]``, :math:`b` has + shape ``[output_dims]``, and :math:`i` indexes the output dimension. - The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, + The values are initialized from the uniform distribution + :math:`\mathcal{U}(-{k}, {k})`, where :math:`k = \frac{1}{\sqrt{D_1}}` and :math:`D_1` is ``input1_dims``. Args: diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py b/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py index 97f6942f0..4de267451 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py @@ -13,7 +13,8 @@ class InstanceNorm(Module): .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta, + y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. Both are of size :attr:`dims`, @@ -21,7 +22,8 @@ class InstanceNorm(Module): Args: dims (int): The number of features of the input. - eps (float): A value added to the denominator for numerical stability. Default: ``1e-5``. + eps (float): A value added to the denominator for numerical stability. Default: + ``1e-5``. affine (bool): Default: ``False``. Shape: diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py b/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py index eaf50228e..051c59e25 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py @@ -48,7 +48,8 @@ def _sliding_windows(x, window_shape, window_strides): raise ValueError( f"To extract sliding windows the window shapes and strides must have " f"the same number of spatial dimensions as the signal but the signal " - f"has {len(spatial_dims)} dims and the window shape has {len(window_shape)} " + f"has {len(spatial_dims)} dims and the window shape has " + f"{len(window_shape)} " f"and strides have {len(window_strides)}." ) @@ -89,8 +90,7 @@ def __init__(self, pooling_function, kernel_size, stride, padding, padding_value 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)}." + f"[{class_name}] '{name}' must be positive but got {tuple(values)}." ) if any(p[0] < 0 for p in padding): raise ValueError( diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py b/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py index 52a6733f2..cbf877016 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py @@ -24,7 +24,8 @@ class RoPE(Module): implementation which is slightly less efficient. Default: ``False``. base (float, optional): The base used to compute angular frequency for each dimension in the positional encodings. Default: ``10000``. - scale (float, optional): The scale used to scale the positions. Default: ``1.0``. + scale (float, optional): The scale used to scale the positions. + Default: ``1.0``. """ def __init__( @@ -89,7 +90,8 @@ def __init__( # embedding dimension must be positive and even. if dims <= 0 or dims % 2 != 0: raise ValueError( - f"[SinusoidalPositionalEncoding] dims must be positive and even but got {dims}." + f"[SinusoidalPositionalEncoding] dims must be positive and even but " + f"got {dims}." ) # Avoid division by zero when dims == 2 (one frequency pair). @@ -132,7 +134,8 @@ class ALiBi(Module): ALiBi adds a static, non-learnable bias matrix to attention scores proportional to the distance between query and key tokens. - For more details see `Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation `_. + For more details see `Train Short, Test Long: Attention with Linear Biases Enables + Input Length Extrapolation `_. """ @staticmethod @@ -150,7 +153,8 @@ def create_alibi_matrix( k_sequence_length (int): The key sequence length. num_heads (int): The number of attention heads. offset (int, optional): The position offset. Default: ``0``. - dtype (Dtype, optional): Data type of the output array. Default: ``mx.float32``. + dtype (Dtype, optional): Data type of the output array. Default: + ``mx.float32``. Returns: array: The ALiBi bias matrix. @@ -170,7 +174,8 @@ def create_alibi_slope(num_heads: int, dtype: mx.Dtype = mx.float32) -> mx.array Args: num_heads (int): The number of attention heads. - dtype (Dtype, optional): Data type of the output array. Default: ``mx.float32``. + dtype (Dtype, optional): Data type of the output array. Default: + ``mx.float32``. Returns: array: The slopes array expanded for head broadcasting. diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py b/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py index 7cebb5b2d..556057f63 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py @@ -67,7 +67,9 @@ def quantize( Weight and input quantization for all linear layers: >>> predicate = lambda p, m: isinstance(m, nn.Linear) - >>> nn.quantize(model, mode="nvfp4", quantize_input=True, class_predicate=predicate) + >>> nn.quantize( + ... model, mode="nvfp4", quantize_input=True, class_predicate=predicate, + ... ) """ class_predicate = class_predicate or (lambda _, m: hasattr(m, "to_quantized")) @@ -318,7 +320,8 @@ class QQLinear(Module): the fly during computation so that gradients with respect to the weights can be computed. - To switch between the two cases, use ``layer.eval()`` and ``layer.train()`` respectively. + To switch between the two cases, use ``layer.eval()`` and ``layer.train()`` + respectively. Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer quantizes the input as well and includes weights in gradient computations. diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py b/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py index 31ec916d1..dfcbd1514 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py @@ -173,9 +173,9 @@ def __init__( self.checkpoint = checkpoint def __call__(self, x, mask): - for l in self.layers: - l = checkpoint(l) if self.checkpoint else l - x = l(x, mask) + for layer in self.layers: + layer = checkpoint(layer) if self.checkpoint else layer + x = layer(x, mask) return self.ln(x) @@ -264,9 +264,9 @@ def __init__( self.checkpoint = checkpoint def __call__(self, x, memory, x_mask, memory_mask): - for l in self.layers: - l = checkpoint(l) if self.checkpoint else l - x = l(x, memory, x_mask, memory_mask) + for layer in self.layers: + layer = checkpoint(layer) if self.checkpoint else layer + x = layer(x, memory, x_mask, memory_mask) return self.ln(x) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py index b04f6e9ba..208b15571 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py @@ -435,7 +435,8 @@ def __init__( 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" + "[Upsample] Antialiasing is not supported for nearest neighbor " + "upsampling" ) if isinstance(scale_factor, (list, tuple)): scale_factor = tuple(map(float, scale_factor)) diff --git a/Source/Cmlx/mlx/python/mlx/nn/losses.py b/Source/Cmlx/mlx/python/mlx/nn/losses.py index b72b91816..f9d5775dc 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/losses.py +++ b/Source/Cmlx/mlx/python/mlx/nn/losses.py @@ -178,7 +178,9 @@ def binary_cross_entropy( >>> probs = mx.array([0.1, 0.1, 0.4, 0.4]) >>> targets = mx.array([0, 0, 1, 1]) - >>> loss = nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean") + >>> loss = nn.losses.binary_cross_entropy( + ... probs, targets, with_logits=False, reduction="mean", + ... ) >>> loss array(0.510826, dtype=float32) """ @@ -304,8 +306,8 @@ def gaussian_nll_loss( inputs (array): The predicted expectation of the Gaussian distribution. targets (array): The target values (samples from the Gaussian distribution). vars (array): The predicted variance of the Gaussian distribution. - full (bool, optional): Whether to include the constant term in the loss calculation. - Default: ``False``. + full (bool, optional): Whether to include the constant term in the loss + calculation. Default: ``False``. eps (float, optional): Small positive constant for numerical stability. Default: ``1e-6``. reduction (str, optional): Specifies the reduction to apply to the output: @@ -439,9 +441,9 @@ def triplet_loss( ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: - 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. + 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 @@ -498,8 +500,8 @@ def huber_loss( Args: inputs (array): The predicted values. targets (array): The target values. - delta (float, optional): The threshold at which to change between L1 and L2 loss. - Default: ``1.0``. + delta (float, optional): The threshold at which to change between L1 and L2 + loss. Default: ``1.0``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. @@ -590,22 +592,22 @@ def margin_ranking_loss( reduction: Reduction = "none", ) -> mx.array: r""" - Calculate the margin ranking loss that loss given inputs :math:`x_1`, :math:`x_2` and a label - :math:`y` (containing 1 or -1). + Calculate the margin ranking loss that loss given inputs :math:`x_1`, :math:`x_2` + and a label :math:`y` (containing 1 or -1). The loss is given by: .. math:: \text{loss} = \max (0, -y * (x_1 - x_2) + \text{margin}) - Where :math:`y` represents ``targets``, :math:`x_1` represents ``inputs1`` and :math:`x_2` - represents ``inputs2``. + Where :math:`y` represents ``targets``, :math:`x_1` represents ``inputs1`` and + :math:`x_2` represents ``inputs2``. Args: inputs1 (array): Scores for the first input. inputs2 (array): Scores for the second input. - targets (array): Labels indicating whether samples in ``inputs1`` should be ranked higher - than samples in ``inputs2``. Values should be 1 or -1. + targets (array): Labels indicating whether samples in ``inputs1`` should be + ranked higher than samples in ``inputs2``. Values should be 1 or -1. margin (float, optional): The margin by which the scores should be separated. Default: ``0.0``. reduction (str, optional): Specifies the reduction to apply to the output: diff --git a/Source/Cmlx/mlx/python/mlx/nn/utils.py b/Source/Cmlx/mlx/python/mlx/nn/utils.py index a60468db6..2f057235b 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/utils.py +++ b/Source/Cmlx/mlx/python/mlx/nn/utils.py @@ -128,12 +128,14 @@ def average_gradients( if all_reduce_size <= 0: return tree_map( - lambda x: mx.distributed.all_sum( - x, - group=group, - stream=communication_stream, - ) - / N, + lambda x: ( + mx.distributed.all_sum( + x, + group=group, + stream=communication_stream, + ) + / N + ), gradients, ) diff --git a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py index 3966f961f..e8362d55a 100644 --- a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py +++ b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py @@ -175,7 +175,7 @@ def __init__(self, optimizers, filters: list = []): if len(filters) != len(optimizers) - 1: raise ValueError( - f"Given {len(filters)} filters but {len(optimizers)-1} needed." + f"Given {len(filters)} filters but {len(optimizers) - 1} needed." ) self.optimizers = optimizers @@ -193,7 +193,8 @@ def _split_dictionary(self, gradients: dict): parts[i].append((k, g)) break - return [tree_unflatten(p) for p in parts] + # tree_unflatten([]) returns a list so use a dict for empty parts + return [tree_unflatten(p) if p else {} for p in parts] def init(self, parameters: dict): for o, p in zip(self.optimizers, self._split_dictionary(parameters)): @@ -297,7 +298,8 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): class RMSprop(Optimizer): r"""The RMSprop optimizer [1]. - [1]: Tieleman, T. and Hinton, G. 2012. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning + [1]: Tieleman, T. and Hinton, G. 2012. Lecture 6.5-rmsprop, coursera: Neural + networks for machine learning .. math:: @@ -338,7 +340,9 @@ def init_single(self, parameter: mx.array, state: dict): state["v"] = mx.zeros_like(parameter) def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): - """Performs the RMSprop parameter update and stores :math:`v` in the optimizer state.""" + """Performs the RMSprop parameter update and stores :math:`v` in the + optimizer state. + """ lr = self.learning_rate.astype(gradient.dtype) alpha = self.alpha eps = self.eps @@ -405,7 +409,8 @@ class AdaDelta(Optimizer): Our AdaDelta implementation follows the original paper. In detail, - [1]: Zeiler, M.D., 2012. ADADELTA: an adaptive learning rate method. arXiv preprint arXiv:1212.5701. + [1]: Zeiler, M.D., 2012. ADADELTA: an adaptive learning rate method. arXiv preprint + arXiv:1212.5701. .. math:: @@ -418,8 +423,8 @@ class AdaDelta(Optimizer): learning_rate (float or callable): The learning rate :math:`\lambda`. rho (float, optional): The coefficient :math:`\rho` used for computing a running average of squared gradients. Default: ``0.9`` - eps (float, optional): The term :math:`\epsilon` added to the denominator to improve - numerical stability. Default: ``1e-6`` + eps (float, optional): The term :math:`\epsilon` added to the denominator to + improve numerical stability. Default: ``1e-6`` """ def __init__( @@ -502,8 +507,7 @@ def __init__( for i, beta in enumerate(betas): if not 0.0 <= beta < 1.0: raise ValueError( - f"Adam beta{i + 1} should be in [0, 1), {beta} was provided " - "instead" + f"Adam beta{i + 1} should be in [0, 1), {beta} was provided instead" ) if not 0.0 <= eps: @@ -556,7 +560,8 @@ class AdamW(Adam): m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\ v_{t+1} &= \beta_2 v_t + (1 - \beta_2) g_t^2 \\ - w_{t+1} &= w_t - \alpha (\frac{m_{t+1}}{\sqrt{v_{t+1}} + \epsilon} + \lambda w_t) + w_{t+1} &= w_t - \alpha (\frac{m_{t+1}}{\sqrt{v_{t+1}} + \epsilon} + \lambda + w_t) Args: learning_rate (float or callable): The learning rate :math:`\alpha`. @@ -678,7 +683,8 @@ class Lion(Optimizer): betas (Tuple[float, float], optional): The coefficients :math:`(\beta_1, \beta_2)` used for computing the gradient momentum and update direction. Default: ``(0.9, 0.99)`` - weight_decay (float, optional): The weight decay :math:`\lambda`. Default: ``0.0`` + weight_decay (float, optional): The weight decay :math:`\lambda`. Default: + ``0.0`` """ def __init__( @@ -692,8 +698,7 @@ def __init__( for i, beta in enumerate(betas): if not 0.0 <= beta < 1.0: raise ValueError( - f"Lion beta{i + 1} should be in [0, 1), {beta} was provided " - "instead" + f"Lion beta{i + 1} should be in [0, 1), {beta} was provided instead" ) self._maybe_schedule("learning_rate", learning_rate) @@ -906,9 +911,10 @@ def init_single(self, parameter: mx.array, state: dict): state["v"] = mx.zeros_like(parameter) def _zeropower_via_newtonschulz5(self, X, steps: int): - assert ( - X.ndim == 2 - ), f"Expected a 2D array for Newton-Schulz iteration, got shape {X.shape} instead." + assert X.ndim == 2, ( + f"Expected a 2D array for Newton-Schulz iteration, got shape {X.shape} " + f"instead." + ) a, b, c = (3.4445, -4.7750, 2.0315) transpose_needed = X.shape[-2] > X.shape[-1] @@ -984,8 +990,20 @@ def clip_grad_norm(grads, max_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) + def squared_norm(g): + # Sum float16 squares in float32 so they do not overflow + if g.dtype == mx.float16: + g = g.astype(mx.float32) + return g.square().sum() + + def scale(g): + # Cast after scaling so small normalizers do not underflow in float16 + if mx.issubdtype(g.dtype, mx.floating): + return (g * normalizer).astype(g.dtype) + return g * normalizer + + norm_squared = tree_reduce(lambda acc, g: acc + squared_norm(g), grads, 0.0) total_norm = mx.sqrt(norm_squared) normalizer = mx.minimum(max_norm / (total_norm + 1e-6), 1.0) - clipped_grads = tree_map(lambda g: g * normalizer, grads) + clipped_grads = tree_map(scale, grads) return clipped_grads, total_norm diff --git a/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py b/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py index 99712097c..54569f0cf 100644 --- a/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py +++ b/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py @@ -119,8 +119,7 @@ def join_schedules(schedules: List[Callable], boundaries: List[int]) -> Callable if len(schedules) != len(boundaries) + 1: raise ValueError( - f"Received {len(boundaries)} boundaries but " - f"expected {len(schedules) - 1}." + f"Received {len(boundaries)} boundaries but expected {len(schedules) - 1}." ) def schedule(step): diff --git a/Source/Cmlx/mlx/python/mlx/utils.py b/Source/Cmlx/mlx/python/mlx/utils.py index f8595ef9a..ecc15f6c8 100644 --- a/Source/Cmlx/mlx/python/mlx/utils.py +++ b/Source/Cmlx/mlx/python/mlx/utils.py @@ -208,8 +208,8 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any: # {"hello": {"world": 42}} Args: - tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a Python tree. - For instance as returned by :meth:`tree_flatten`. + tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a + Python tree. For instance as returned by :meth:`tree_flatten`. Returns: A Python tree. @@ -240,12 +240,12 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any: is_list = False if is_list: - l = [] + items = [] 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 + # if i <= len(items), no {} will be appended. + items.extend([{} for _ in range(i - len(items))]) + items.append(tree_unflatten(children[k])) + return items else: return {k: tree_unflatten(v) for k, v in children.items()} diff --git a/Source/Cmlx/mlx/python/src/CMakeLists.txt b/Source/Cmlx/mlx/python/src/CMakeLists.txt index 0798add41..d5cda4dee 100644 --- a/Source/Cmlx/mlx/python/src/CMakeLists.txt +++ b/Source/Cmlx/mlx/python/src/CMakeLists.txt @@ -1,7 +1,6 @@ nanobind_add_module( core NB_STATIC - STABLE_ABI FREE_THREADED LTO NOMINSIZE diff --git a/Source/Cmlx/mlx/python/src/array.cpp b/Source/Cmlx/mlx/python/src/array.cpp index 0f13293a1..e634bf31b 100644 --- a/Source/Cmlx/mlx/python/src/array.cpp +++ b/Source/Cmlx/mlx/python/src/array.cpp @@ -1,9 +1,4 @@ // Copyright © 2023-2024 Apple Inc. -#include -#include -#include -#include - #include #include #include @@ -13,6 +8,11 @@ #include #include +#include +#include +#include +#include + #include "mlx/backend/metal/metal.h" #include "mlx/utils.h" #include "python/src/buffer.h" @@ -129,9 +129,10 @@ void init_array(nb::module_& m) { return nb::isinstance(other) && t == nb::cast(other); }) - .def("__hash__", [](const mx::Dtype& t) { - return static_cast(t.val()); - }); + .def( + "__hash__", + [](const mx::Dtype& t) { return static_cast(t.val()); }) + .freeze(); m.attr("bool_") = nb::cast(mx::bool_); m.attr("uint8") = nb::cast(mx::uint8); @@ -229,13 +230,16 @@ void init_array(nb::module_& m) { &mx::finfo::smallest_normal, R"pbdoc(The smallest positive normal number.)pbdoc") .def_ro("dtype", &mx::finfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc") - .def("__repr__", [](const mx::finfo& f) { - std::ostringstream os; - os << "finfo(" - << "min=" << f.min << ", max=" << f.max << ", dtype=" << f.dtype - << ")"; - return os.str(); - }); + .def( + "__repr__", + [](const mx::finfo& f) { + std::ostringstream os; + os << "finfo(" + << "min=" << f.min << ", max=" << f.max << ", dtype=" << f.dtype + << ")"; + return os.str(); + }) + .freeze(); nb::class_( m, @@ -253,13 +257,16 @@ void init_array(nb::module_& m) { &mx::iinfo::max, R"pbdoc(The largest representable number.)pbdoc") .def_ro("dtype", &mx::iinfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc") - .def("__repr__", [](const mx::iinfo& i) { - std::ostringstream os; - os << "iinfo(" - << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype - << ")"; - return os.str(); - }); + .def( + "__repr__", + [](const mx::iinfo& i) { + std::ostringstream os; + os << "iinfo(" + << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype + << ")"; + return os.str(); + }) + .freeze(); nb::class_( m, @@ -274,7 +281,8 @@ void init_array(nb::module_& m) { .def("multiply", &ArrayAt::multiply, "value"_a) .def("divide", &ArrayAt::divide, "value"_a) .def("maximum", &ArrayAt::maximum, "value"_a) - .def("minimum", &ArrayAt::minimum, "value"_a); + .def("minimum", &ArrayAt::minimum, "value"_a) + .freeze(); nb::class_( m, @@ -283,7 +291,8 @@ void init_array(nb::module_& m) { Any Python object which has an ``__mlx__array__`` method that returns an :obj:`array`. )pbdoc") - .def(nb::init_implicit()); + .def(nb::init_implicit()) + .freeze(); nb::class_( m, @@ -292,7 +301,8 @@ void init_array(nb::module_& m) { A helper object to iterate over the 1st dimension of an array. )pbdoc") .def("__next__", &ArrayPythonIterator::next) - .def("__iter__", [](const ArrayPythonIterator& it) { return it; }); + .def("__iter__", [](const ArrayPythonIterator& it) { return it; }) + .freeze(); // Install buffer protocol functions PyType_Slot array_slots[] = { @@ -506,7 +516,7 @@ void init_array(nb::module_& m) { nd.ndim(), reinterpret_cast(nd.shape_ptr()), owner, - nullptr, + nd.stride_ptr(), nb::dtype()), mx::bfloat16)); } else { @@ -806,6 +816,7 @@ void init_array(nb::module_& m) { }, "other"_a) .def("__neg__", [](const mx::array& a) { return -a; }) + .def("__pos__", [](const mx::array& a) { return mx::copy(a); }) .def("__bool__", [](mx::array& a) { return nb::bool_(to_scalar(a)); }) .def( "__repr__", @@ -1586,5 +1597,6 @@ void init_array(nb::module_& m) { "dtype"_a, nb::kw_only(), "stream"_a = nb::none(), - "See :func:`view`."); + "See :func:`view`.") + .freeze(); } diff --git a/Source/Cmlx/mlx/python/src/buffer.h b/Source/Cmlx/mlx/python/src/buffer.h index 4b194b3d2..553bba80c 100644 --- a/Source/Cmlx/mlx/python/src/buffer.h +++ b/Source/Cmlx/mlx/python/src/buffer.h @@ -88,9 +88,13 @@ extern "C" inline int getbuffer(PyObject* obj, Py_buffer* view, int flags) { std::memset(view, 0, sizeof(Py_buffer)); auto a = nb::cast(nb::handle(obj)); - { + // Exceptions can not propagate through the buffer protocol. + try { nb::gil_scoped_release nogil; a.eval(); + } catch (const std::exception& e) { + PyErr_SetString(PyExc_RuntimeError, e.what()); + return -1; } std::vector shape(a.shape().begin(), a.shape().end()); diff --git a/Source/Cmlx/mlx/python/src/convert.cpp b/Source/Cmlx/mlx/python/src/convert.cpp index ad1bc61a5..ac735e1ff 100644 --- a/Source/Cmlx/mlx/python/src/convert.cpp +++ b/Source/Cmlx/mlx/python/src/convert.cpp @@ -1,13 +1,13 @@ // Copyright © 2024 Apple Inc. +#include +#include + #include #include #include #include -#include -#include - #include "python/src/convert.h" #include "python/src/utils.h" @@ -185,14 +185,17 @@ mx::array cpu_nd_array_to_mlx( // 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. +// dtype width mismatch, a pointer not aligned to the item size, 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) { + // GPU kernels need aligned data + auto ptr = reinterpret_cast(nd_array.data()); if (!mx::metal::is_available() || - nd_array.itemsize() != mx::size_of(dst_dtype)) { + nd_array.itemsize() != mx::size_of(dst_dtype) || + ptr % nd_array.itemsize() != 0) { return std::nullopt; } @@ -206,19 +209,21 @@ std::optional cpu_nd_array_to_mlx_no_copy( return std::nullopt; } + auto byte_offset = nd_array.byte_offset(); 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); - }); + mx::allocator::Data( + buf, + [owner = std::move(nd_array)](mx::allocator::Buffer b) { + mx::allocator::release(b); + }), + storage_size, + std::move(strides), + flags, + byte_offset); out.set_status(mx::array::Status::available); return out; } @@ -238,14 +243,16 @@ mx::array metal_nd_array_to_mlx( } auto [storage_size, strides, flags] = get_strided_layout(nd_array, shape); auto data_handle = nd_array.data_handle(); + auto byte_offset = nd_array.byte_offset(); mx::array out(shape, src_dtype, nullptr, {}); out.set_data( - mx::allocator::Buffer(data_handle), + mx::allocator::Data( + mx::allocator::Buffer(data_handle), + [owner = std::move(nd_array)](mx::allocator::Buffer) {}), storage_size, std::move(strides), flags, - nd_array.byte_offset(), - [owner = std::move(nd_array)](mx::allocator::Buffer) {}); + byte_offset); out.set_status(mx::array::Status::available); if (copy) { @@ -634,6 +641,10 @@ mx::array array_from_list_impl( std::vector vals; fill_vector(pl, vals); return mx::array(vals.begin(), shape, dtype); + } else if (dtype == mx::float64) { + std::vector vals; + fill_vector(pl, vals); + return mx::array(vals.begin(), shape, dtype); } else if (mx::issubdtype(dtype, mx::inexact)) { std::vector vals; fill_vector(pl, vals); @@ -646,7 +657,7 @@ mx::array array_from_list_impl( } case pyfloat: { auto out_type = specified_type.value_or(mx::float32); - if (out_type == mx::float64) { + if (out_type == mx::float64 || !mx::issubdtype(out_type, mx::inexact)) { std::vector vals; fill_vector(pl, vals); return mx::array(vals.begin(), shape, out_type); @@ -741,7 +752,7 @@ mx::array create_array( return mx::array(val, t.value_or(default_type)); } else if (nb::isinstance(v)) { auto out_type = t.value_or(mx::float32); - if (out_type == mx::float64) { + if (out_type == mx::float64 || !mx::issubdtype(out_type, mx::inexact)) { return mx::array(nb::cast(v), out_type); } else { return mx::array(nb::cast(v), out_type); diff --git a/Source/Cmlx/mlx/python/src/device.cpp b/Source/Cmlx/mlx/python/src/device.cpp index 83d32d1cb..da4f919b1 100644 --- a/Source/Cmlx/mlx/python/src/device.cpp +++ b/Source/Cmlx/mlx/python/src/device.cpp @@ -1,14 +1,14 @@ // Copyright © 2023-2025 Apple Inc. -#include -#include - #include #include #include #include #include +#include +#include + #include "mlx/device.h" #include "mlx/utils.h" @@ -43,13 +43,16 @@ void init_device(nb::module_& m) { os << d; return os.str(); }) - .def("__eq__", [](const mx::Device& d, const nb::object& other) { - if (!nb::isinstance(other) && - !nb::isinstance(other)) { - return false; - } - return d == nb::cast(other); - }); + .def( + "__eq__", + [](const mx::Device& d, const nb::object& other) { + if (!nb::isinstance(other) && + !nb::isinstance(other)) { + return false; + } + return d == nb::cast(other); + }) + .freeze(); nb::implicitly_convertible(); diff --git a/Source/Cmlx/mlx/python/src/distributed.cpp b/Source/Cmlx/mlx/python/src/distributed.cpp index 0ad91c6e4..0d1e341f9 100644 --- a/Source/Cmlx/mlx/python/src/distributed.cpp +++ b/Source/Cmlx/mlx/python/src/distributed.cpp @@ -51,7 +51,8 @@ void init_distributed(nb::module_& parent_module) { color (int): A value to group processes into subgroups. key (int, optional): A key to optionally change the rank ordering of the processes. - )pbdoc"); + )pbdoc") + .freeze(); m.def( "is_available", diff --git a/Source/Cmlx/mlx/python/src/export.cpp b/Source/Cmlx/mlx/python/src/export.cpp index 7748f87b1..1becd22bc 100644 --- a/Source/Cmlx/mlx/python/src/export.cpp +++ b/Source/Cmlx/mlx/python/src/export.cpp @@ -298,7 +298,8 @@ void init_export(nb::module_& m) { auto [args_, kwargs_] = validate_and_extract_inputs(args, kwargs, "[export_function]"); exporter(args_, kwargs_); - }); + }) + .freeze(); m.def( "exporter", diff --git a/Source/Cmlx/mlx/python/src/fft.cpp b/Source/Cmlx/mlx/python/src/fft.cpp index 99f8478bb..37a176006 100644 --- a/Source/Cmlx/mlx/python/src/fft.cpp +++ b/Source/Cmlx/mlx/python/src/fft.cpp @@ -634,4 +634,124 @@ void init_fft(nb::module_& parent_module) { Returns: array: The inverse-shifted array with the same shape as the input. )pbdoc"); + m.def( + "stft", + [](const mx::array& x, + int n_fft, + const std::optional& hop_length, + const std::optional& win_length, + const std::optional& window, + bool center, + const std::string& pad_mode, + const std::string& norm, + const std::optional& onesided, + mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "stft"); + return mx::fft::stft( + x, + n_fft, + hop_length, + win_length, + window, + center, + pad_mode, + fft_norm, + onesided, + s); + }, + "x"_a, + "n_fft"_a = 2048, + "hop_length"_a = nb::none(), + "win_length"_a = nb::none(), + "window"_a = nb::none(), + "center"_a = true, + "pad_mode"_a = "reflect", + "norm"_a = "backward", + "onesided"_a = nb::none(), + "stream"_a = nb::none(), + R"pbdoc( + Short-Time Fourier Transform. + + Splits the last axis of ``x`` into overlapping frames of length + ``n_fft``, applies ``window``, and takes the Fourier Transform of each + frame. The output is always complex with shape ``(..., n_freq, n_frames)``. + + Args: + x (array): The input array. + n_fft (int, optional): Size of the Fourier Transform. Default is ``2048``. + hop_length (int, optional): Samples between frames. Default is ``n_fft // 4``. + win_length (int, optional): Window size. Default is ``n_fft``. + window (array, optional): A one-dimensional window, centered and + zero-padded to ``n_fft``. Default is a rectangular window. + center (bool, optional): Pad both ends so frames are centered. Default is ``True``. + pad_mode (str, optional): One of ``"reflect"``, ``"constant"``, or + ``"edge"``, used when ``center`` is ``True``. Default is ``"reflect"``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. + onesided (bool, optional): Return only non-negative frequencies. + Default is ``True`` for real input and ``False`` for complex input. + + Returns: + array: The STFT of ``x``. + )pbdoc"); + m.def( + "istft", + [](const mx::array& stft_matrix, + const std::optional& n_fft, + const std::optional& hop_length, + const std::optional& win_length, + const std::optional& window, + bool center, + const std::string& norm, + const std::optional& onesided, + const std::optional& length, + mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "istft"); + return mx::fft::istft( + stft_matrix, + n_fft, + hop_length, + win_length, + window, + center, + fft_norm, + onesided, + length, + s); + }, + "stft_matrix"_a, + "n_fft"_a = nb::none(), + "hop_length"_a = nb::none(), + "win_length"_a = nb::none(), + "window"_a = nb::none(), + "center"_a = true, + "norm"_a = "backward", + "onesided"_a = nb::none(), + "length"_a = nb::none(), + "stream"_a = nb::none(), + R"pbdoc( + Inverse Short-Time Fourier Transform. + + Inverts :func:`stft` by overlap-adding the windowed inverse transform + of each frame, normalized by the squared-window envelope. + + Args: + stft_matrix (array): The STFT with shape ``(..., n_freq, n_frames)``. + n_fft (int, optional): Size of the Fourier Transform. Default is + inferred from ``n_freq`` and ``onesided``. The onesided inference + is even; pass ``n_fft`` for an odd transform size. + hop_length (int, optional): Samples between frames. Default is ``n_fft // 4``. + win_length (int, optional): Window size. Default is ``n_fft``. + window (array, optional): A one-dimensional window. Default is a + rectangular window. + center (bool, optional): Whether the STFT was centered. Default is ``True``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. + onesided (bool, optional): Whether the STFT has only non-negative + frequencies. Default is ``True``. + length (int, optional): Trim or zero-pad the output to this length. + + Returns: + array: The reconstructed signal. + )pbdoc"); } diff --git a/Source/Cmlx/mlx/python/src/indexing.cpp b/Source/Cmlx/mlx/python/src/indexing.cpp index 5bf467750..679291a71 100644 --- a/Source/Cmlx/mlx/python/src/indexing.cpp +++ b/Source/Cmlx/mlx/python/src/indexing.cpp @@ -1,10 +1,10 @@ // Copyright © 2023-2024 Apple Inc. +#include + #include #include #include -#include - #include "mlx/dtype.h" #include "mlx/ops.h" #include "python/src/convert.h" @@ -75,6 +75,33 @@ void get_slice_params( nb::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size); } +// Resolve negative bounds and clamp into range, mirroring CPython's +// PySlice_AdjustIndices. Needed wherever a slice is expanded into an explicit +// arange, since out-of-range bounds would otherwise be used verbatim. +void adjust_slice_bounds( + mx::ShapeElem& start, + mx::ShapeElem& end, + mx::ShapeElem stride, + int axis_size) { + if (start < 0) { + start += axis_size; + if (start < 0) { + start = (stride < 0) ? -1 : 0; + } + } else if (start >= axis_size) { + start = (stride < 0) ? axis_size - 1 : axis_size; + } + + if (end < 0) { + end += axis_size; + if (end < 0) { + end = (stride < 0) ? -1 : 0; + } + } else if (end >= axis_size) { + end = (stride < 0) ? axis_size - 1 : axis_size; + } +} + mx::array get_int_index(nb::object idx, int axis_size) { int idx_ = safe_to_int32(idx); idx_ = (idx_ < 0) ? idx_ + axis_size : idx_; @@ -155,9 +182,7 @@ mx::array mlx_gather_nd( get_slice_params( start, end, stride, nb::cast(idx), src.shape(i)); - // Handle negative indices - start = (start < 0) ? start + src.shape(i) : start; - end = (end < 0) ? end + src.shape(i) : end; + adjust_slice_bounds(start, end, stride, src.shape(i)); gather_indices.push_back(arange(start, end, stride, mx::uint32)); num_slices++; @@ -686,9 +711,7 @@ mlx_scatter_args_nd( get_slice_params( start, end, stride, nb::cast(pyidx), axis_size); - // Handle negative indices - start = (start < 0) ? start + axis_size : start; - end = (end < 0) ? end + axis_size : end; + adjust_slice_bounds(start, end, stride, axis_size); mx::Shape idx_shape(idx_ndim, 1); @@ -789,7 +812,8 @@ mlx_compute_scatter_args( } else if (nb::isinstance(obj)) { return mlx_scatter_args_nd(src, nb::cast(obj), vals); } else if (nb::isinstance(obj)) { - return {{}, broadcast_to(vals, src.shape()), {}}; + return { + {}, broadcast_to(squeeze_leading_singletons(vals), src.shape()), {}}; } else if (obj.is_none()) { return {{}, broadcast_to(vals, src.shape()), {}}; } else if (nb::isinstance(obj)) { diff --git a/Source/Cmlx/mlx/python/src/linalg.cpp b/Source/Cmlx/mlx/python/src/linalg.cpp index a5e82ae85..4a46e7788 100644 --- a/Source/Cmlx/mlx/python/src/linalg.cpp +++ b/Source/Cmlx/mlx/python/src/linalg.cpp @@ -1,13 +1,13 @@ // Copyright © 2023-2024 Apple Inc. -#include - #include #include #include #include #include +#include + #include "mlx/linalg.h" #include "python/src/small_vector.h" diff --git a/Source/Cmlx/mlx/python/src/memory.cpp b/Source/Cmlx/mlx/python/src/memory.cpp index 529995539..571944953 100644 --- a/Source/Cmlx/mlx/python/src/memory.cpp +++ b/Source/Cmlx/mlx/python/src/memory.cpp @@ -1,8 +1,8 @@ // Copyright © 2025-2026 Apple Inc. -#include "mlx/memory.h" #include +#include "mlx/memory.h" #include "python/src/trees.h" namespace mx = mlx::core; diff --git a/Source/Cmlx/mlx/python/src/metal.cpp b/Source/Cmlx/mlx/python/src/metal.cpp index fb3c51306..8d0d4242e 100644 --- a/Source/Cmlx/mlx/python/src/metal.cpp +++ b/Source/Cmlx/mlx/python/src/metal.cpp @@ -1,6 +1,4 @@ // Copyright © 2023-2024 Apple Inc. -#include - #include #include #include @@ -8,6 +6,8 @@ #include #include +#include + #include "mlx/backend/metal/metal.h" #include "mlx/device.h" #include "mlx/memory.h" diff --git a/Source/Cmlx/mlx/python/src/mlx_func.cpp b/Source/Cmlx/mlx/python/src/mlx_func.cpp index 9955e134f..9f84f5c79 100644 --- a/Source/Cmlx/mlx/python/src/mlx_func.cpp +++ b/Source/Cmlx/mlx/python/src/mlx_func.cpp @@ -2,17 +2,11 @@ #include "python/src/mlx_func.h" -#include - -// A garbage collected function which wraps nb::cpp_function -// See https://github.com/wjakob/nanobind/discussions/919 +// The wrapper owns callbacks because free-threaded nanobind functions are +// immortal struct gc_func { - PyObject_HEAD - // Vector call implementation that forwards calls to nanobind - PyObject* (*vectorcall)(PyObject*, PyObject* const*, size_t, PyObject*); - // The nanobind wrapper func - PyObject* func; + std::unique_ptr func; // The original wrapped func PyObject* orig_func; @@ -22,97 +16,48 @@ struct gc_func { int gc_func_tp_traverse(PyObject* self, visitproc visit, void* arg) { Py_VISIT(Py_TYPE(self)); - gc_func* w = (gc_func*)self; - Py_VISIT(w->func); - for (auto d : w->deps) { - Py_VISIT(d); + if (nb::inst_ready(self)) { + for (auto d : nb::inst_ptr(self)->deps) { + Py_VISIT(d); + } } return 0; -}; +} int gc_func_tp_clear(PyObject* self) { - gc_func* w = (gc_func*)self; - Py_CLEAR(w->func); + auto* w = nb::inst_ptr(self); + w->orig_func = nullptr; + w->deps.clear(); + w->func.reset(); return 0; } -PyObject* gc_func_get_doc(PyObject* self, void*) { - return PyObject_GetAttrString(((gc_func*)self)->func, "__doc__"); -} - -PyObject* gc_func_get_sig(PyObject* self, void*) { - return PyObject_GetAttrString(((gc_func*)self)->func, "__nb_signature__"); -} - -PyObject* gc_func_vectorcall( - PyObject* self, - PyObject* const* args, - size_t nargs, - PyObject* kwnames) { - return PyObject_Vectorcall(((gc_func*)self)->func, args, nargs, kwnames); -} - -void gc_func_dealloc(PyObject* self) { - PyObject_GC_UnTrack(self); - Py_XDECREF(((gc_func*)self)->func); - PyObject_GC_Del(self); -} - -static PyMemberDef gc_func_members[] = { - {"__vectorcalloffset__", - T_PYSSIZET, - (Py_ssize_t)offsetof(gc_func, vectorcall), - READONLY, - nullptr}, - {nullptr, 0, 0, 0, nullptr}}; - -static PyGetSetDef gc_func_getset[] = { - {"__doc__", gc_func_get_doc, nullptr, nullptr, nullptr}, - {"__nb_signature__", gc_func_get_sig, nullptr, nullptr, nullptr}, - {nullptr, nullptr, nullptr, nullptr, nullptr}}; - static PyObject* gc_func_getattro(PyObject* self, PyObject* name_) { - gc_func* w = (gc_func*)self; - return PyObject_GenericGetAttr(w->orig_func, name_); + return PyObject_GenericGetAttr(nb::inst_ptr(self)->orig_func, name_); } // Table of custom type slots we want to install PyType_Slot gc_func_slots[] = { {Py_tp_traverse, (void*)gc_func_tp_traverse}, {Py_tp_clear, (void*)gc_func_tp_clear}, - {Py_tp_getset, (void*)gc_func_getset}, {Py_tp_getattro, (void*)gc_func_getattro}, - {Py_tp_members, (void*)gc_func_members}, - {Py_tp_call, (void*)PyVectorcall_Call}, - {Py_tp_dealloc, (void*)gc_func_dealloc}, {0, 0}}; -static PyType_Spec gc_func_spec = { - /* .name = */ "mlx.gc_func", - /* .basicsize = */ (int)sizeof(gc_func), - /* .itemsize = */ 0, - /* .flags = */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | - Py_TPFLAGS_HAVE_VECTORCALL, - /* .slots = */ gc_func_slots}; - -static PyTypeObject* gc_func_tp = nullptr; - nb::callable mlx_func( - nb::object func, + std::unique_ptr func, const nb::callable& orig_func, std::vector deps) { - gc_func* r = (gc_func*)PyType_GenericAlloc(gc_func_tp, 0); - r->func = func.inc_ref().ptr(); - r->orig_func = orig_func.ptr(); - deps.push_back(r->orig_func); - r->deps = std::move(deps); - r->vectorcall = gc_func_vectorcall; - return nb::steal((PyObject*)r); + deps.push_back(orig_func.ptr()); + return nb::borrow( + nb::cast(gc_func{std::move(func), orig_func.ptr(), std::move(deps)})); } void init_mlx_func(nb::module_& m) { - gc_func_tp = (PyTypeObject*)PyType_FromSpec(&gc_func_spec); - if (!gc_func_tp) { - nb::raise("Could not register MLX function type."); - } + nb::class_(m, "_gc_func", nb::type_slots(gc_func_slots)) + .def( + "__call__", + [](gc_func& self, nb::args args, nb::kwargs kwargs) { + return (*self.func)(args, kwargs); + }) + .freeze(); } diff --git a/Source/Cmlx/mlx/python/src/mlx_func.h b/Source/Cmlx/mlx/python/src/mlx_func.h index 79c8376de..3b0c0a10e 100644 --- a/Source/Cmlx/mlx/python/src/mlx_func.h +++ b/Source/Cmlx/mlx/python/src/mlx_func.h @@ -2,30 +2,36 @@ #pragma once -#include - #include -#include + +#include +#include +#include namespace nb = nanobind; using namespace nb::literals; +struct PyFunction { + virtual ~PyFunction() = default; + virtual nb::object operator()(nb::args& args, nb::kwargs& kwargs) = 0; +}; + nb::callable mlx_func( - nb::object func, + std::unique_ptr func, const nb::callable& orig_func, std::vector deps); template nb::callable mlx_func(F func, const nb::callable& orig_func, Deps&&... deps) { - return mlx_func( - nb::cpp_function(std::move(func)), - orig_func, - std::vector{deps.ptr()...}); -} + struct Callback : PyFunction { + F func; + + explicit Callback(F func) : func(std::move(func)) {} -template -nb::callable -mlx_func(nb::object func, const nb::callable& orig_func, Deps&&... deps) { + nb::object operator()(nb::args& args, nb::kwargs& kwargs) override { + return nb::cast(func(args, kwargs)); + } + }; return mlx_func( - std::move(func), orig_func, std::vector{deps.ptr()...}); + std::make_unique(std::move(func)), orig_func, {deps.ptr()...}); } diff --git a/Source/Cmlx/mlx/python/src/ops.cpp b/Source/Cmlx/mlx/python/src/ops.cpp index fe63bbf71..4d83c6293 100644 --- a/Source/Cmlx/mlx/python/src/ops.cpp +++ b/Source/Cmlx/mlx/python/src/ops.cpp @@ -1,10 +1,5 @@ // Copyright © 2023-2024 Apple Inc. -#include -#include -#include -#include - #include #include #include @@ -13,6 +8,11 @@ #include #include +#include +#include +#include +#include + #include "mlx/einsum.h" #include "mlx/ops.h" #include "mlx/utils.h" @@ -746,6 +746,9 @@ void init_ops(nb::module_& m) { - All but the last two dimensions of each input are broadcast with one another using standard numpy-style broadcasting semantics. + Only inexact types are supported. The promoted type of ``a`` and ``b`` + must be floating point or complex. + Args: a (array): Input array or scalar. b (array): Input array or scalar. @@ -1155,19 +1158,24 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "arctan2", - &mx::arctan2, + [](const ScalarOrArray& a_, + const ScalarOrArray& b_, + mx::StreamOrDevice s) { + auto [a, b] = to_arrays(a_, b_); + return mx::arctan2(a, b, s); + }, nb::arg(), nb::arg(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctan2(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), + "def arctan2(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse tangent of the ratio of two arrays. Args: - a (array): Input array. - b (array): Input array. + a (array or scalar): Input array or scalar. + b (array or scalar): Input array or scalar. Returns: array: The inverse tangent of the ratio of ``a`` and ``b``. @@ -1295,12 +1303,12 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def degrees(a: array, /, *, stream: StreamOrDevice = None) -> array"), + "def degrees(a: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from radians to degrees. Args: - a (array): Input array. + a (array or scalar): Input array. Returns: array: The angles in degrees. @@ -1314,12 +1322,12 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def radians(a: array, /, *, stream: StreamOrDevice = None) -> array"), + "def radians(a: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from degrees to radians. Args: - a (array): Input array. + a (array or scalar): Input array. Returns: array: The angles in radians. @@ -2730,6 +2738,37 @@ void init_ops(nb::module_& m) { Returns: array: The output array with the corresponding axes reduced. )pbdoc"); + m.def( + "minmax", + [](const mx::array& a, + const IntOrVec& axis, + bool keepdims, + mx::StreamOrDevice s) -> std::tuple { + auto axes = get_reduce_axes(axis, a.ndim()); + return mx::minmax(a, axes, keepdims, s); + }, + nb::arg(), + "axis"_a = nb::none(), + "keepdims"_a = false, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def minmax(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> tuple[array, array]"), + R"pbdoc( + Return the minimum and maximum over the given axes. + + Equivalent to calling :func:`min` and :func:`max` separately. + + Args: + a (array): Input array. + axis (int or list(int), optional): Axis or axes to reduce over. + If unspecified, reduce over the entire array. + keepdims (bool, optional): Keep reduced axes as singleton + dimensions. Default: ``False``. + + Returns: + tuple(array, array): The minimum and maximum arrays, in that order. + )pbdoc"); m.def( "logcumsumexp", [](const mx::array& a, @@ -3246,6 +3285,108 @@ void init_ops(nb::module_& m) { Returns: array: The top ``k`` elements from the input. )pbdoc"); + m.def( + "unique", + [](const mx::array& a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const std::optional& fill_value, + mx::StreamOrDevice s) -> nb::object { + std::optional fill_value_ = std::nullopt; + if (fill_value) { + fill_value_ = to_array(fill_value.value(), a.dtype()); + } + auto out = mx::unique( + a, + size, + return_index, + return_inverse, + return_counts, + fill_value_, + s); + if (out.size() == 1) { + return nb::cast(out.at(0)); + } + nb::list result; + for (auto& o : out) { + result.append(o); + } + return nb::tuple(result); + }, + nb::arg(), + "size"_a, + "return_index"_a = false, + "return_inverse"_a = false, + "return_counts"_a = false, + "fill_value"_a = nb::none(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def unique(a: array, /, size: int, return_index: bool = False, return_inverse: bool = False, return_counts: bool = False, fill_value: scalar | array | None = None, *, stream: StreamOrDevice = None) -> array | tuple[array, ...]"), + R"pbdoc( + Returns the sorted unique elements of the flattened array. + + The shape of the output does not depend on the values in ``a``, so + ``a`` is not evaluated. Unlike NumPy, the size is never inferred from + the values and must be provided. + + Entries past the last unique element hold ``fill_value`` and their + count is ``0``, so ``mx.sum(counts > 0)`` gives the number of unique + elements as long as the output is not truncated. + + A truncated output keeps the smallest ``size`` unique elements, and + ``index``, ``inverse`` and ``counts`` then stop describing all of ``a``, + and ``inverse`` and ``counts`` stop agreeing with each other: the counts + of the dropped elements are gone, while their indices in ``inverse`` are + clamped to the last entry. + ``inverse`` is meaningless for a ``size`` of ``0``, since there is no + element left for it to point at. + + This op builds on ``mx.sort``, so a ``bool`` input needs the CPU + stream, which is where ``mx.sort`` supports it. + + Args: + a (array): Input array. + size (int): The size of the output. If the size is smaller than + the number of unique elements of ``a``, the output is truncated. + If it is larger, the output is padded with ``fill_value``. + return_index (bool, optional): If ``True``, also return the index + in the flattened ``a`` of the first occurrence of each unique + element. Default: ``False``. + return_inverse (bool, optional): If ``True``, also return the + indices of the unique array that rebuild ``a``. The indices have + the same shape as ``a``. Default: ``False``. + return_counts (bool, optional): If ``True``, also return the number + of times each unique element occurs in ``a``. Default: ``False``. + fill_value (scalar or array, optional): The value of the entries + past the last unique element. If ``None``, this defaults to the + first of the sorted unique elements. Default: ``None``. + + Returns: + array or tuple(array, ...): The sorted unique elements. If any of + ``return_index``, ``return_inverse`` or ``return_counts`` is + ``True``, a tuple with the requested arrays in the order values, + index, inverse, counts. + + Example: + >>> a = mx.array([2, 1, 2, 3, 1]) + >>> mx.unique(a, 3) + array([1, 2, 3], dtype=int32) + >>> mx.unique(a, 5) + array([1, 2, 3, 1, 1], dtype=int32) + >>> values, index, inverse, counts = mx.unique( + ... a, 4, True, True, True, fill_value=0) + >>> values + array([1, 2, 3, 0], dtype=int32) + >>> index + array([1, 0, 3, 1], dtype=uint32) + >>> inverse + array([1, 0, 1, 2, 0], dtype=uint32) + >>> counts + array([2, 2, 1, 0], dtype=int32) + )pbdoc"); m.def( "broadcast_to", [](const ScalarOrArray& a, const mx::Shape& shape, mx::StreamOrDevice s) { @@ -4727,14 +4868,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. @@ -5825,6 +5966,9 @@ void init_ops(nb::module_& m) { Perform the Einstein summation convention on the operands. + Contractions are implemented with :func:`matmul` and have the same + requirements. + Args: subscripts (str): The Einstein summation convention equation. *operands (array): The input arrays. diff --git a/Source/Cmlx/mlx/python/src/print.cpp b/Source/Cmlx/mlx/python/src/print.cpp index 6c08d88d3..85404bd8a 100644 --- a/Source/Cmlx/mlx/python/src/print.cpp +++ b/Source/Cmlx/mlx/python/src/print.cpp @@ -1,9 +1,9 @@ +#include + #include #include #include -#include - #include "mlx/utils.h" #include "python/src/utils.h" @@ -33,7 +33,8 @@ void init_print(nb::module_& m) { // Expose printing options to Python: allow setting global precision. nb::class_(m, "PrintOptions") .def(nb::init(), "precision"_a = -1) - .def_rw("precision", &mx::PrintOptions::precision); + .def_rw("precision", &mx::PrintOptions::precision) + .freeze(); m.def( "set_printoptions", @@ -64,7 +65,8 @@ void init_print(nb::module_& m) { nb::class_(m, "_PrintOptionsContext") .def(nb::init()) .def("__enter__", &PrintOptionsContext::enter) - .def("__exit__", &PrintOptionsContext::exit); + .def("__exit__", &PrintOptionsContext::exit) + .freeze(); m.def( "printoptions", diff --git a/Source/Cmlx/mlx/python/src/random.cpp b/Source/Cmlx/mlx/python/src/random.cpp index 10b82b892..8710a8d70 100644 --- a/Source/Cmlx/mlx/python/src/random.cpp +++ b/Source/Cmlx/mlx/python/src/random.cpp @@ -74,13 +74,9 @@ void reset_random_state() { class RandomState {}; nb::object random_state_sentinel() { - static nb::object sentinel = []() { - auto sentinel = nb::cast(RandomState{}); - sentinel.inc_ref(); - return sentinel; - }(); + static nb::handle sentinel = nb::cast(RandomState{}).release(); - return sentinel; + return nb::borrow(sentinel); } mx::array random_state_key() { @@ -107,9 +103,10 @@ void init_random(nb::module_& parent_module) { return default_key().state()[0]; }, "index"_a) - .def("__iter__", [](const RandomState&) { - return nb::iter(default_key().state()); - }); + .def( + "__iter__", + [](const RandomState&) { return nb::iter(default_key().state()); }) + .freeze(); m.def("__getattr__", [&](nb::handle key) -> nb::object { // Create random.state lazily to avoid initializing device during import. diff --git a/Source/Cmlx/mlx/python/src/stream.cpp b/Source/Cmlx/mlx/python/src/stream.cpp index 2ddb99179..5706fc878 100644 --- a/Source/Cmlx/mlx/python/src/stream.cpp +++ b/Source/Cmlx/mlx/python/src/stream.cpp @@ -1,13 +1,13 @@ // Copyright © 2023-2024 Apple Inc. -#include -#include - #include #include #include #include +#include +#include + #include "mlx/stream.h" #include "mlx/utils.h" #include "python/src/random.h" @@ -73,10 +73,13 @@ void init_stream(nb::module_& m) { os << s; return os.str(); }) - .def("__eq__", [](const mx::Stream& s, const nb::object& other) { - return nb::isinstance(other) && - s == nb::cast(other); - }); + .def( + "__eq__", + [](const mx::Stream& s, const nb::object& other) { + return nb::isinstance(other) && + s == nb::cast(other); + }) + .freeze(); nb::class_( m, @@ -97,7 +100,8 @@ void init_stream(nb::module_& m) { [](const mx::ThreadLocalStream& s, const nb::object& other) { return nb::isinstance(other) && s == nb::cast(other); - }); + }) + .freeze(); nb::implicitly_convertible(); @@ -179,7 +183,8 @@ void init_stream(nb::module_& m) { const std::optional& traceback) { scm.exit(); }, "exc_type"_a = nb::none(), "exc_value"_a = nb::none(), - "traceback"_a = nb::none()); + "traceback"_a = nb::none()) + .freeze(); m.def( "stream", [](mx::StreamOrDevice s) { return PyStreamContext(s); }, diff --git a/Source/Cmlx/mlx/python/src/transforms.cpp b/Source/Cmlx/mlx/python/src/transforms.cpp index 6c2299b7f..223e232e2 100644 --- a/Source/Cmlx/mlx/python/src/transforms.cpp +++ b/Source/Cmlx/mlx/python/src/transforms.cpp @@ -1,10 +1,5 @@ // Copyright © 2023-2024 Apple Inc. -#include -#include -#include -#include - #include #include #include @@ -13,6 +8,11 @@ #include #include +#include +#include +#include +#include + #include "mlx/array.h" #include "mlx/compile.h" #include "mlx/compile_impl.h" @@ -310,7 +310,11 @@ auto py_vmap( const nb::callable& fun, const nb::object& in_axes, const nb::object& out_axes) { - return [fun, in_axes, out_axes](const nb::args& args) { + return [fun, in_axes, out_axes]( + const nb::args& args, const nb::kwargs& kwargs) { + if (kwargs.size()) { + throw nb::type_error("[vmap] Keyword arguments are not supported."); + } auto axes_to_flat_tree = [](const nb::object& tree, const nb::object& axes, bool output_axes) { @@ -1170,7 +1174,8 @@ void init_transforms(nb::module_& m) { return a pytree with the vectorization axes of each output. If some outputs are no longer vectorized, then their vectorization axis should be ``None``. - )pbdoc"); + )pbdoc") + .freeze(); m.def( "eval", @@ -1471,7 +1476,7 @@ void init_transforms(nb::module_& m) { const nb::object& outputs, bool shapeless) { return mlx_func( - nb::cpp_function(PyCompiledFun{fun, inputs, outputs, shapeless}), + PyCompiledFun{fun, inputs, outputs, shapeless}, fun, inputs, outputs); diff --git a/Source/Cmlx/mlx/python/src/trees.cpp b/Source/Cmlx/mlx/python/src/trees.cpp index cf05d4e25..38ba3a8a9 100644 --- a/Source/Cmlx/mlx/python/src/trees.cpp +++ b/Source/Cmlx/mlx/python/src/trees.cpp @@ -273,14 +273,9 @@ nb::object tree_unflatten( } nb::object structure_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; - }(); + static nb::handle sentinel = PyObject_New(PyObject, &PyBaseObject_Type); - return sentinel; + return nb::borrow(sentinel); } std::pair, nb::object> tree_flatten_with_structure( diff --git a/Source/Cmlx/mlx/python/src/utils.h b/Source/Cmlx/mlx/python/src/utils.h index d2fcbd037..f4d47b3f3 100644 --- a/Source/Cmlx/mlx/python/src/utils.h +++ b/Source/Cmlx/mlx/python/src/utils.h @@ -1,15 +1,15 @@ // Copyright © 2023-2024 Apple Inc. #pragma once -#include -#include -#include -#include - #include #include #include #include +#include +#include +#include +#include + #include "mlx/array.h" #include "python/src/convert.h" diff --git a/Source/Cmlx/mlx/python/tests/mlx_tests.py b/Source/Cmlx/mlx/python/tests/mlx_tests.py index ba5595389..10a7e5c35 100644 --- a/Source/Cmlx/mlx/python/tests/mlx_tests.py +++ b/Source/Cmlx/mlx/python/tests/mlx_tests.py @@ -58,7 +58,8 @@ def setUp(self): def tearDown(self): mx.set_default_device(self.default) - # Note if a tuple is passed into args, it will be considered a shape request and convert to a mx.random.normal with the shape matching the tuple + # Note if a tuple is passed into args, it will be considered a shape request and + # convert to a mx.random.normal with the shape matching the tuple def assertCmpNumpy( self, args: List[Union[Tuple[int], Any]], diff --git a/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py b/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py index db9bbb2e6..be3f38e2e 100644 --- a/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py +++ b/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py @@ -143,7 +143,6 @@ def loss_fn(model, x, y): return ((logits - y) ** 2).mean() for dtype, atol, rtol in dtypes: - mx.random.seed(0xF0F0F0F0) kx, ky = mx.random.split(mx.random.key(rank)) diff --git a/Source/Cmlx/mlx/python/tests/ring_test_distributed.py b/Source/Cmlx/mlx/python/tests/ring_test_distributed.py index 31db2ed5f..aba897ea1 100644 --- a/Source/Cmlx/mlx/python/tests/ring_test_distributed.py +++ b/Source/Cmlx/mlx/python/tests/ring_test_distributed.py @@ -21,7 +21,17 @@ def test_groups(self): self.assertEqual(world.rank(), world2.rank()) with self.assertRaises(RuntimeError): - sub = world.split(world.rank() % 2) + world.split(world.rank() % 2) + + def test_strict_init_after_fallback(self): + with mlx_tests.scoped_env( + JACCL_IBV_DEVICES=None, + MLX_IBV_DEVICES=None, + ): + self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) + + with self.assertRaises(RuntimeError): + mx.distributed.init(backend="jaccl", strict=True) def test_all_reduce_extra(self): world = mx.distributed.init() diff --git a/Source/Cmlx/mlx/python/tests/run.py b/Source/Cmlx/mlx/python/tests/run.py index df96d4dd3..b5c1dbd70 100644 --- a/Source/Cmlx/mlx/python/tests/run.py +++ b/Source/Cmlx/mlx/python/tests/run.py @@ -9,7 +9,7 @@ __unittest = True -import mlx_tests +import mlx_tests # noqa: E402 if __name__ == "__main__": # Run all tests by default. diff --git a/Source/Cmlx/mlx/python/tests/test_array.py b/Source/Cmlx/mlx/python/tests/test_array.py index e9d781e52..425a59268 100644 --- a/Source/Cmlx/mlx/python/tests/test_array.py +++ b/Source/Cmlx/mlx/python/tests/test_array.py @@ -1,7 +1,6 @@ # Copyright © 2023-2024 Apple Inc. import operator -import os import pickle import platform import sys @@ -599,6 +598,18 @@ def test_double_keeps_precision(self): out = mx.array([x], dtype=mx.float64).item() self.assertEqual(out, x) + # Integer and bool targets + v = 2**24 + 1 + self.assertEqual(mx.array(float(v), dtype=mx.int32).item(), v) + self.assertEqual(mx.array([float(v)], dtype=mx.int64).tolist(), [v]) + # Small non-zero values are still True + self.assertTrue(mx.array(1e-50, dtype=mx.bool_).item()) + self.assertEqual(mx.array([1e-50], dtype=mx.bool_).tolist(), [True]) + + # Python ints to float64 + out = mx.array([v], dtype=mx.float64).tolist() + self.assertEqual(out, [float(v)]) + def test_construction_from_lists_wide_ints(self): # A python int that does not fit in int32 widens to int64, the same # rule the scalar path already uses. It used to raise std::bad_cast. @@ -649,7 +660,8 @@ def test_construction_from_lists_of_mlx_arrays(self): expected = mx.stack([x, y], axis=0) self.assertEqualArray(z, expected) - # check heterogeneous construction with mlx arrays and python primitive types + # check heterogeneous construction with mlx arrays and python primitive + # types x, y = mx.array([True], x_t), mx.array([False], y_t) z = mx.array([[x, [2.0]], [[3.0], y]]) expected = mx.array([[[x.item()], [2.0]], [[3.0], [y.item()]]], z.dtype) @@ -1062,11 +1074,14 @@ def test_array_comparison(self): self.assertEqual((a > 1).tolist(), [False, False, True]) self.assertEqual((a >= 1).tolist(), [False, True, True]) - def test_array_neg(self): + def test_array_unary_ops(self): a = mx.array([-1.0, 4.0, 0.0]) self.assertEqual((-a).tolist(), [1.0, -4.0, 0.0]) + self.assertEqual((+a).tolist(), [-1.0, 4.0, 0.0]) + assert +a is not a + def test_array_type_cast(self): a = mx.array([0.1, 2.3, -1.3]) b = [0, 2, -1] @@ -1114,6 +1129,10 @@ def test_array_pickle(self): state = pickle.dumps(x) y = pickle.loads(state) self.assertEqualArray(y, x) + # F-contiguous + x = x.T + y = pickle.loads(pickle.dumps(x)) + self.assertEqualArray(y, x) def test_array_copy(self): dtypes = [ @@ -1255,7 +1274,7 @@ def check_slices(arr_np, *idx_np): idx_mlx = [ mx.array(idx) if isinstance(idx, np.ndarray) else idx for idx in idx_np ] - slice_mlx = arr_mlx[tuple(idx_mlx)] + _slice_mlx = arr_mlx[tuple(idx_mlx)] self.assertTrue( np.array_equal(arr_np[tuple(idx_np)], arr_mlx[tuple(idx_mlx)]) ) @@ -1366,7 +1385,7 @@ def test_setitem(self): a[0:2] = 3 self.assertEqual(a.tolist(), [3, 3, 1]) - # Assigning through a bare Ellipsis, like a[:] and a[None] + # Assigning through a bare Ellipsis, like a[:] and a[(...,)] e = mx.zeros((2, 3), mx.int32) e[...] = 5 self.assertEqual(e.tolist(), [[5, 5, 5], [5, 5, 5]]) @@ -1378,16 +1397,44 @@ def test_setitem(self): e[...] = mx.zeros((2, 3), mx.int32) self.assertEqual(e.tolist(), [[0, 0, 0], [0, 0, 0]]) + # Leading singleton dimensions of the update are squeezed + e[...] = mx.array([[[1, 2, 3], [4, 5, 6]]]) + self.assertEqual(e.tolist(), [[1, 2, 3], [4, 5, 6]]) + + # The target keeps its leading dimension of size 1 + h = mx.zeros((1, 3), mx.int32) + h[...] = mx.array([[[1, 2, 3]]]) + self.assertEqual(h.tolist(), [[1, 2, 3]]) + + # a[...] assigns the same values as a[(...,)] + f = mx.zeros((3,), mx.int32) + g = mx.zeros((3,), mx.int32) + f[...] = mx.array([[1, 2, 3]]) + g[(...,)] = mx.array([[1, 2, 3]]) + self.assertEqual(f.tolist(), [1, 2, 3]) + self.assertEqual(f.tolist(), g.tolist()) + # Scalar array e = mx.array(0) e[...] = 7 self.assertEqual(e.item(), 7) + e[...] = mx.array([[9]]) + self.assertEqual(e.tolist(), 9) + # Shapes that cannot broadcast are still rejected e = mx.zeros((2, 3), mx.int32) with self.assertRaises(ValueError): e[...] = mx.array([1, 2]) + # The squeeze stops at the first dimension that is not 1 + with self.assertRaises(ValueError): + e[...] = mx.zeros((2, 1, 3), mx.int32) + + # A squeezed update that still does not broadcast is rejected + with self.assertRaises(ValueError): + e[...] = mx.zeros((1, 4), mx.int32) + a[0:3] = 4 self.assertEqual(a.tolist(), [4, 4, 4]) @@ -1700,6 +1747,13 @@ 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)) + # Indices and updates with negative strides + idx = mx.array([0, 1, 2, 3]) + upd = mx.array([1.0, 2.0, 3.0, 4.0]) + expected = [4.0, 3.0, 2.0, 1.0] + self.assertEqual(mx.zeros(4).at[idx[::-1]].add(upd).tolist(), expected) + self.assertEqual(mx.zeros(4).at[idx].add(upd[::-1]).tolist(), expected) + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") def test_array_at_complex_add_gpu(self): n = 4096 @@ -1903,6 +1957,33 @@ def test_slice_negative_step(self): b_mx = a_mx[::-1, ::-3, ::-2] self.assertTrue(np.array_equal(b_np, b_mx)) + def test_slice_bounds_with_array_index(self): + # Out-of-range slice bounds must be clamped the same way NumPy clamps + # them when the slice is combined with an array index. + a_np = np.arange(20, dtype=np.int32).reshape(4, 5) + a_mx = mx.array(a_np) + idx_np = np.array([0, 1]) + idx_mx = mx.array([0, 1], dtype=mx.uint32) + + bounds = [None, -100, -6, -4, -1, 0, 1, 3, 4, 6, 100] + steps = [None, 1, 2, 3, -1, -2, -3] + + for start in bounds: + for stop in bounds: + for step in steps: + s = slice(start, stop, step) + with self.subTest(start=start, stop=stop, step=step): + self.assertTrue( + np.array_equal(a_np[s, idx_np], a_mx[s, idx_mx]) + ) + + # Same clamping applies when assigning through the slice + u_np = a_np.copy() + u_mx = mx.array(a_np) + u_np[s, idx_np] = 0 + u_mx[s, idx_mx] = 0 + self.assertTrue(np.array_equal(u_np, u_mx)) + def test_api(self): x = mx.array(np.random.rand(10, 10, 10)) ops = [ @@ -2022,7 +2103,8 @@ def test_buffer_protocol(self): self.assertEqual(mv_mx.shape, mv_np.shape, f"{mlx_dtype}{np_dtype}") # correct buffer format for 8 byte (unsigned) 'long long' is Q/q, see # https://docs.python.org/3.10/library/struct.html#format-characters - # numpy returns L/l, as 'long' is equivalent to 'long long' on 64bit machines, so q and l are equivalent + # numpy returns L/l, as 'long' is equivalent to 'long long' on 64bit + # machines, so q and l are equivalent # see https://github.com/pybind/pybind11/issues/1908 if np_dtype == np.uint64: self.assertEqual(mv_mx.format, "Q", f"{mlx_dtype}{np_dtype}") @@ -2083,9 +2165,15 @@ def test_buffer_protocol_ref_counting(self): mv = memoryview(a) a = None self.assertIsNotNone(wr()) - mv = None + del mv self.assertIsNone(wr()) + def test_buffer_protocol_eval_error(self): + # Errors from evaluating the array are raised instead of aborting + a = mx.linalg.inv(mx.zeros((2, 2)), stream=mx.cpu) + with self.assertRaises(RuntimeError): + memoryview(a) + def test_array_view_ref_counting(self): a = mx.arange(3) wr = weakref.ref(a) @@ -2093,7 +2181,7 @@ def test_array_view_ref_counting(self): a_np = np.array(a, copy=False) a = None self.assertIsNotNone(wr()) - a_np = None + del a_np self.assertIsNone(wr()) def test_create_from_buffer(self): @@ -2469,7 +2557,7 @@ def test_torch_mps_asarray_copy_true_copies_dlpack_input(self): self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") - def test_torch_mps_dlpack_zero_copy_shares_updates(self): + def test_torch_mps_dlpack_zero_copy_one_way_updates(self): assert torch is not None x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) torch.mps.synchronize() @@ -2481,7 +2569,8 @@ def test_torch_mps_dlpack_zero_copy_shares_updates(self): y += 10 mx.eval(y) - self.assertEqual(x.cpu().numpy().tolist(), y.tolist()) + self.assertEqual(y.tolist(), [[10.0] * 4] * 3) + self.assertEqual(x.cpu().numpy().tolist(), [[0.0] * 4] * 3) @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") def test_torch_mps_dlpack_matching_dtype_argument_shares_updates(self): @@ -2646,7 +2735,7 @@ def test_mlx_dlpack_export_torch_update_writes_mlx_buffer(self): self.assertEqual(x.tolist(), [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 6.0, 7.0]) @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") - def test_from_dlpack_torch_mps_copy_none_shares_updates(self): + def test_from_dlpack_torch_mps_copy_none_one_way_updates(self): assert torch is not None x = torch.arange(3, device="mps", dtype=torch.float32) torch.mps.synchronize() @@ -2658,7 +2747,8 @@ def test_from_dlpack_torch_mps_copy_none_shares_updates(self): y += 10 mx.eval(y) - self.assertEqual(x.cpu().numpy().tolist(), [10.0, 10.0, 10.0]) + self.assertEqual(y.tolist(), [10.0, 10.0, 10.0]) + self.assertEqual(x.cpu().numpy().tolist(), [0.0, 0.0, 0.0]) @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") def test_from_dlpack_torch_mps_copy_false_shares_updates(self): @@ -2767,6 +2857,13 @@ def test_setitem_with_boolean_mask(self): a[mask] = 5.0 self.assertTrue(mx.array_equal(a, expected)) + # Scalar target and scalar mask + for selected in (False, True): + a = mx.array(False) + a[mx.array(selected)] = True + self.assertEqual(a.shape, ()) + self.assertEqual(a.item(), selected) + mask_np = np.zeros((1, 10, 10), dtype=bool) with self.assertRaises(ValueError): mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0 @@ -2891,7 +2988,7 @@ def test_format(self): self.assertEqual(f"{b:.1f}", "0.4") with self.assertRaises(TypeError): - s = f"{a:.2f}" + f"{a:.2f}" a = mx.array([1, 2, 3]) self.assertEqual(f"{a}", "array([1, 2, 3], dtype=int32)") @@ -2947,6 +3044,41 @@ def test_large_indices(self): with self.assertRaises(ValueError): x[2**32] + @unittest.skipUnless(sys.version_info >= (3, 15), "requires Python 3.15") + def test_frozen_types(self): + types = ( + mx.array, + mx.Dtype, + mx.finfo, + mx.iinfo, + mx.ArrayAt, + mx.ArrayLike, + mx.ArrayIterator, + mx.Device, + mx.Stream, + mx.ThreadLocalStream, + mx.StreamContext, + mx.PrintOptions, + mx.custom_function, + mx.FunctionExporter, + mx.distributed.Group, + ) + for bound_type in types: + with self.subTest(type=bound_type): + with self.assertRaises(TypeError): + bound_type._test_attribute = None + with self.assertRaises(TypeError): + bound_type.__doc__ = "modified" + with self.assertRaises(TypeError): + del bound_type.__doc__ + + def test_array_subclass(self): + class Array(mx.array): + def total(self): + return self.sum().item() + + self.assertEqual(Array([1, 2, 3]).total(), 6) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_autograd.py b/Source/Cmlx/mlx/python/tests/test_autograd.py index 13cc3c42a..92441ad51 100644 --- a/Source/Cmlx/mlx/python/tests/test_autograd.py +++ b/Source/Cmlx/mlx/python/tests/test_autograd.py @@ -4,6 +4,7 @@ import itertools import math import unittest +import weakref import mlx.core as mx import mlx_tests @@ -19,18 +20,24 @@ class TestAutograd(mlx_tests.MLXTestCase): def test_jvp(self): - fun = lambda x: 2 * x + def fun(x): + return 2 * x + out, dout = mx.jvp(fun, [mx.array(1.0)], [mx.array(2.0)]) self.assertEqual(out[0].item(), 2.0) self.assertEqual(dout[0].item(), 4.0) - fun = lambda x, y: x * y + def fun(x, y): + return x * y + _, out = mx.jvp( fun, [mx.array(4.0), mx.array(2.0)], [mx.array(3.0), mx.array(2.0)] ) self.assertEqual(out[0].item(), 4.0 * 2.0 + 2.0 * 3.0) - fun = lambda x, y, z: (x * y, y * z) + def fun(x, y, z): + return (x * y, y * z) + _, out = mx.jvp( fun, [mx.array(2.0), mx.array(4.0), mx.array(6.0)], @@ -40,6 +47,55 @@ def test_jvp(self): self.assertEqual(out[0].item(), 4.0 * 1.0 + 2.0 * 3.0) self.assertEqual(out[1].item(), 4.0 * 1.0 + 6.0 * 3.0) + def test_logcumsumexp_jvp(self): + x = mx.array([0.0, 0.0, 0.0]) + t = mx.array([1.0, 1.0, 1.0]) + + for rev in (False, True): + + def f(z): + return mx.logcumsumexp(z, axis=0, reverse=rev, inclusive=True) + + _, (dout,) = mx.jvp(f, [x], [t]) + self.assertTrue(mx.allclose(dout, mx.array([1.0, 1.0, 1.0]))) + + def f(z): + return mx.logcumsumexp(z, axis=0, reverse=False, inclusive=False) + + _, (dout,) = mx.jvp(f, [x], [t]) + self.assertTrue(mx.allclose(dout, mx.array([0.0, 1.0, 1.0]))) + + def f(z): + return mx.logcumsumexp(z, axis=0, reverse=True, inclusive=False) + + _, (dout,) = mx.jvp(f, [x], [t]) + self.assertTrue(mx.allclose(dout, mx.array([1.0, 1.0, 0.0]))) + + x = mx.array([1000.0, 1000.0, 1000.0]) + + def f(z): + return mx.logcumsumexp(z, axis=0, reverse=False, inclusive=True) + + _, (dout,) = mx.jvp(f, [x], [t]) + self.assertTrue(mx.allclose(dout, mx.array([1.0, 1.0, 1.0]))) + + x = mx.array([1.0, 2.0, 3.0]) + t = mx.array([1.0, -2.0, 0.5]) + eps = 1e-3 + + def f(z): + return mx.logcumsumexp(z, axis=0) + + _, (dout,) = mx.jvp(f, [x], [t]) + + expected = (f(x + eps * t) - f(x - eps * t)) / (2 * eps) + self.assertTrue(mx.allclose(dout, expected, atol=1e-3)) + + x = mx.array([1.0, 2.0], dtype=mx.complex64) + t = mx.array([1 + 1j, 2 - 1j]) + with self.assertRaises(ValueError): + mx.jvp(f, [x], [t]) + def test_jvp_comparison_tangent_dtype(self): # Comparison op JVP tangents should preserve the input tangent's # dtype (e.g. float32), not return bool. Using bool tangents causes @@ -135,17 +191,23 @@ def fun(x): self.assertTrue(mx.array_equal(dout, mx.array([1.0, 0.0, 1.0]))) def test_vjp(self): - fun = lambda x: 2 * x + def fun(x): + return 2 * x + out, dout = mx.vjp(fun, [mx.array(1.0)], [mx.array(2.0)]) self.assertEqual(out[0].item(), 2.0) self.assertEqual(dout[0].item(), 4.0) - fun = lambda x, y: x * y + def fun(x, y): + return x * y + _, dout = mx.vjp(fun, [mx.array(4.0), mx.array(2.0)], [mx.array(3.0)]) self.assertEqual(dout[0].item(), 6.0) self.assertEqual(dout[1].item(), 12.0) - fun = lambda x, y, z: (x * y, y * z) + def fun(x, y, z): + return (x * y, y * z) + _, out = mx.vjp( fun, [mx.array(2.0), mx.array(4.0), mx.array(6.0)], @@ -158,7 +220,9 @@ def test_vjp(self): def test_jvp_with_partly_traced_inputs(self): # power: each traced input must use its own tangent (issue #3634) - fun = lambda a, b: a**b + def fun(a, b): + return a**b + primals = [mx.array(2.0), mx.array(3.0)] dyda = 12.0 # d/da a^b = b * a^(b - 1) dydb = math.log(2.0) * 8.0 # d/db a^b = ln(a) * a^b @@ -168,10 +232,16 @@ def test_jvp_with_partly_traced_inputs(self): self.assertAlmostEqual(j.item(), dydb, places=4) _, (j,) = mx.jvp(fun, primals, [mx.array(1.0), mx.array(1.0)]) self.assertAlmostEqual(j.item(), dyda + dydb, places=4) - fun = lambda a: a ** mx.array(3.0) + + def fun(a): + return a ** mx.array(3.0) + _, (j,) = mx.jvp(fun, [mx.array(2.0)], [mx.array(1.0)]) self.assertAlmostEqual(j.item(), dyda, places=4) - fun = lambda b: mx.array(2.0) ** b + + def fun(b): + return mx.array(2.0) ** b + _, (j,) = mx.jvp(fun, [mx.array(3.0)], [mx.array(1.0)]) self.assertAlmostEqual(j.item(), dydb, places=4) @@ -190,18 +260,28 @@ def fun(x): src = mx.zeros(4) upd = mx.ones(2) start = mx.array([1]) - fun = lambda u: mx.slice_update(src, u, start, axes=[0]) + + def fun(u): + return mx.slice_update(src, u, start, axes=[0]) + _, (j,) = mx.jvp(fun, [upd], [mx.ones(2)]) self.assertEqual(j.tolist(), [0.0, 1.0, 1.0, 0.0]) - fun = lambda s: mx.slice_update(s, upd, start, axes=[0]) + + def fun(s): + return mx.slice_update(s, upd, start, axes=[0]) + _, (j,) = mx.jvp(fun, [src], [mx.ones(4)]) self.assertEqual(j.tolist(), [1.0, 0.0, 0.0, 1.0]) - fun = lambda s, u: mx.slice_update(s, u, start, axes=[0]) + + def fun(s, u): + return mx.slice_update(s, u, start, axes=[0]) + _, (j,) = mx.jvp(fun, [src, upd], [mx.ones(4), mx.full(2, 2.0)]) self.assertEqual(j.tolist(), [1.0, 2.0, 2.0, 1.0]) def test_grad(self): - fun = lambda x: x * x + def fun(x): + return x * x value, dfdx = mx.value_and_grad(fun)(mx.array(0.5)) self.assertEqual(value.item(), 0.25) @@ -215,7 +295,9 @@ def test_grad(self): df3dx3 = mx.grad(mx.grad(mx.grad(fun)))(mx.array(0.5)) self.assertEqual(df3dx3.item(), 0.0) - fun = lambda x, y: x * y + def fun(x, y): + return x * y + x = mx.array(2.0) y = mx.array(3.0) dfdx = mx.grad(fun, argnums=0)(x, y) @@ -224,7 +306,9 @@ def test_grad(self): self.assertEqual(dfdx.item(), 2.0) # Pass non array args to functions works - fun = lambda x, y: x + def fun(x, y): + return x + value, dfdx = mx.value_and_grad(fun)(mx.array(2.0), "hello") self.assertEqual(value.item(), 2.0) self.assertEqual(dfdx.item(), 1.0) @@ -233,12 +317,16 @@ def test_grad(self): self.assertEqual(dfdx.item(), 1.0) # Raises when function does not return array - fun = lambda x: "hello" + def fun(x): + return "hello" + with self.assertRaises(ValueError): mx.grad(fun)(mx.array(2.0)) # Raises for invalid argument number or argument type - fun = lambda x: x + def fun(x): + return x + with self.assertRaises(ValueError): mx.grad(fun, argnums=2)(mx.array(2.0)) with self.assertRaises(ValueError): @@ -247,43 +335,57 @@ def test_grad(self): mx.grad(fun)("hello") # Raises when output is not a scalar array - fun = lambda x: mx.sum(x, keepdims=True) + def fun(x): + return mx.sum(x, keepdims=True) + with self.assertRaises(ValueError): mx.grad(fun)(mx.ones((2, 2))) def test_grad_trees(self): - fun = lambda x, y: x * y + def fun(x, y): + return x * y + value, dfdx = mx.value_and_grad(fun, (0, 1))(mx.array(0.5), mx.array(2.0)) self.assertEqual(value.item(), 1.0) self.assertTrue(isinstance(dfdx, tuple)) self.assertEqual(dfdx[0].item(), 2.0) self.assertEqual(dfdx[1].item(), 0.5) - fun = lambda x, y: x * y + def fun(x, y): + return x * y + value, dfdx = mx.value_and_grad(fun, 1)(mx.array(0.5), mx.array(2.0)) self.assertEqual(value.item(), 1.0) self.assertEqual(dfdx.item(), 0.5) - fun = lambda p: p["x"] * p["y"] + def fun(p): + return p["x"] * p["y"] + value, dfdx = mx.value_and_grad(fun)({"x": mx.array(0.5), "y": mx.array(2.0)}) self.assertEqual(value.item(), 1.0) self.assertEqual(dfdx["x"].item(), 2.0) self.assertEqual(dfdx["y"].item(), 0.5) - fun = lambda p: p["x"] * p["y"] + def fun(p): + return p["x"] * p["y"] + with self.assertRaises(ValueError): mx.value_and_grad(fun)({"x": 0.5, "y": mx.array(2.0)}) with self.assertRaises(ValueError): mx.value_and_grad(fun, (0, 1))({"x": mx.array(0.5), "y": mx.array(2.0)}) - fun = lambda p, b: mx.square(p[0]["foo"][2]) * b + def fun(p, b): + return mx.square(p[0]["foo"][2]) * b + value, dfdx = mx.value_and_grad(fun)( [{"foo": [[], [], mx.array(2.0)]}], mx.array(0.5) ) self.assertEqual(value.item(), 2.0) self.assertEqual(dfdx[0]["foo"][2].item(), 2.0) - fun = lambda x: x + def fun(x): + return x + with self.assertRaises(TypeError): mx.value_and_grad(fun, (None, None)) with self.assertRaises(ValueError): @@ -293,9 +395,13 @@ def test_grad_trees(self): def test_auxiliary_values(self): def fun(x, y): - l = (x * y).sum() - extra = {"loss": l, "foo": y.square() + x.square(), "bar": [1, 2, 3, y, x]} - return l, extra + loss = (x * y).sum() + extra = { + "loss": loss, + "foo": y.square() + x.square(), + "bar": [1, 2, 3, y, x], + } + return loss, extra fun_value_grad = mx.value_and_grad(fun) fun_grad = mx.grad(fun) @@ -312,7 +418,9 @@ def fun(x, y): _ = fun_grad(mx.ones((2, 2)), mx.ones((2, 2))) def test_grad_kwargs(self): - fun = lambda x, y: x * y + def fun(x, y): + return x * y + a, b = mx.array(0.5), mx.array(2.0) dfdx = mx.grad(fun) self.assertEqual(dfdx(a, b).item(), 2.0) @@ -343,7 +451,9 @@ def test_grad_kwargs(self): self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["y"].item(), 0.5) - fun = lambda x, y, z: x * y * z + def fun(x, y, z): + return x * y * z + dfdxyz = mx.grad(fun, argnums=[0, 1], argnames=["z"]) c = mx.array(4.0) grads = dfdxyz(a, b, z=c) @@ -354,7 +464,9 @@ def test_grad_kwargs(self): self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["z"].item(), 1.0) - fun = lambda x, y: x * y + def fun(x, y): + return x * y + dfdy = mx.grad(fun, argnames=["y"]) grads = dfdy(a, y=b) self.assertTrue(isinstance(grads, tuple)) @@ -364,9 +476,15 @@ def test_grad_kwargs(self): def test_captured(self): a = mx.array(5.0) - f = lambda x: a + x - g = lambda x: a + a - h = lambda x: x + x + + def f(x): + return a + x + + def g(x): + return a + a + + def h(x): + return x + x dfdx = mx.grad(f) self.assertEqual(dfdx(a).item(), 1.0) @@ -404,7 +522,9 @@ def h(w, x): self.assertTrue(mx.allclose(vjps[0], 24.0 * mx.ones(shape_in))) self.assertTrue(mx.allclose(vjps[1], mx.zeros(shape_in))) - g = lambda x: h(w_in, x) + def g(x): + return h(w_in, x) + vals, vjps = mx.vjp(g, [x_in], [cotan]) mx.eval(vjps) @@ -731,7 +851,7 @@ def func(x): _, vjps = mx.vjp(func, (arr,), (cotan,)) self.assertEqual(vjps[0].item(), 8.0) - def test_power_grad(self): + def test_power_grad(self): # noqa: F811 def fun(x, y): res = x - y return res**x @@ -1105,7 +1225,7 @@ def outer_f(a, b): out1, dout1 = mx.jvp(outer_f, inputs, tans) @my_double.jvp - def random_grads(primals, tangents): + def random_grads(primals, tangents): # noqa: F811 return { "out": 2 * primals["x"] * tangents["y"] + 2 * primals["y"] * tangents["x"] @@ -1261,6 +1381,51 @@ def fun(a, b): expected[4:-5:-2] = tan_b self.assertTrue(mx.allclose(grad, expected)) + def test_slice_grads_single_element(self): + # A strided slice selecting exactly one element must place each + # cotangent at the position the forward pass read it from, not over + # the whole half-open span (regression: the stride was collapsed to 1 + # without narrowing stop, so the vjp re-derived a wider region). + def grad_of(fun, x): + return mx.grad(fun)(x) + + # positive stride, first element of a length-2 axis + x = mx.zeros((2, 3)) + Wa = mx.ones((1, 3)) + Wb = mx.ones((1, 3)) * 2 + + def fun(z): + return (z[0::2] * Wa).sum() + (z[1::2] * Wb).sum() + + g = grad_of(fun, x) + self.assertTrue(mx.allclose(g, mx.array([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]))) + + # positive stride, single element in a longer span + x = mx.zeros((8, 2)) + W = mx.ones((1, 2)) + g = grad_of(lambda z: (z[2:5:3] * W).sum(), x) + expected = mx.zeros((8, 2)) + expected[2] = W + self.assertTrue(mx.allclose(g, expected)) + + # negative stride, single element + x = mx.zeros((2, 2)) + W = mx.ones((1, 2)) + g = grad_of(lambda z: (z[::-2] * W).sum(), x) + expected = mx.zeros((2, 2)) + expected[1] = W + self.assertTrue(mx.allclose(g, expected)) + + # negative stride with start == 0 (cannot be expressed as a + # single-element negative-stride span; the unit-stride fallback must + # still land the cotangent at index 0) + x = mx.zeros((3, 2)) + W = mx.ones((1, 2)) + g = grad_of(lambda z: (z[0::-1] * W).sum(), x) + expected = mx.zeros((3, 2)) + expected[0] = W + self.assertTrue(mx.allclose(g, expected)) + def test_leaks(self): for transform in [ mx.grad, @@ -1324,7 +1489,9 @@ def test_conv_second_order_grad_with_padding(self): w = mx.random.normal((1, 2, 2, 1)) def wgrad_sum(x_): - fn = lambda w_: mx.conv2d(x_, w_, padding=(1, 2)).sum() + def fn(w_): + return mx.conv2d(x_, w_, padding=(1, 2)).sum() + return mx.grad(fn)(w).sum() dx = mx.grad(wgrad_sum)(x) @@ -1485,8 +1652,22 @@ def test_complex_unary_vjps(self): mx.tan: lambda x: 1 / mx.cos(x) ** 2, mx.tanh: lambda x: 1 - mx.tanh(x) ** 2, mx.log1p: lambda x: 1 / (1 + x), + mx.cos: lambda x: -mx.sin(x), + mx.sqrt: lambda x: 0.5 / mx.sqrt(x), + mx.rsqrt: lambda x: -0.5 * mx.rsqrt(x) / x, + mx.arcsin: lambda x: 1 / mx.sqrt(1 - mx.square(x)), + mx.arccos: lambda x: -1 / mx.sqrt(1 - mx.square(x)), + mx.arctan: lambda x: 1 / (1 + mx.square(x)), + mx.arcsinh: lambda x: 1 / mx.sqrt(mx.square(x) + 1), + # 1 / sqrt(x**2 - 1) has the wrong sign when Re(x) < 0 + mx.arccosh: lambda x: 1 / (mx.sqrt(x - 1) * mx.sqrt(x + 1)), + mx.arctanh: lambda x: 1 / (1 - mx.square(x)), } for fn, deriv in ops.items(): + _, (jvp,) = mx.jvp(fn, [z], [cotangent]) + expected = cotangent * deriv(z) + self.assertTrue(mx.allclose(jvp, expected, atol=1e-5), msg=str(fn)) + _, (vjp,) = mx.vjp(fn, [z], [cotangent]) expected = cotangent * mx.conj(deriv(z)) self.assertTrue(mx.allclose(vjp, expected, atol=1e-5), msg=str(fn)) @@ -1545,7 +1726,10 @@ def numerical_hvp(f, x, v, eps=1e-3): 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) + + def f(a): + return mx.sum(fn(a) ** 2) + self.assertTrue( mx.allclose(hvp(f, x, v), numerical_hvp(f, x, v), atol=1e-3) ) @@ -1558,11 +1742,39 @@ def numerical_hvp(f, x, v, eps=1e-3): lambda a: mx.partition(a, 1, axis=0), lambda a: mx.cummax(a, axis=0), ): - f = lambda a: mx.sum(fn(a) ** 2) + + def f(a): + return mx.sum(fn(a) ** 2) + self.assertTrue( mx.allclose(hvp(f, y, w), numerical_hvp(f, y, w), atol=1e-3) ) + def test_transform_releases_captures(self): + for transform in ( + mx.compile, + mx.grad, + mx.value_and_grad, + mx.vmap, + mx.checkpoint, + ): + for cycle in (False, True): + with self.subTest(transform=transform, cycle=cycle): + captured = mx.array(2.0) + ref = weakref.ref(captured) + + def fun(x, captured=captured): + return (x * captured).sum() + + fn = transform(fun) + if cycle: + fun.wrapped = fn + result = fn(mx.ones((2, 3))) + mx.eval(result) + del result, fn, fun, captured + gc.collect() + self.assertIsNone(ref()) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_bf16.py b/Source/Cmlx/mlx/python/tests/test_bf16.py index 1f90e1e51..8ef2b97fc 100644 --- a/Source/Cmlx/mlx/python/tests/test_bf16.py +++ b/Source/Cmlx/mlx/python/tests/test_bf16.py @@ -1,8 +1,6 @@ # Copyright © 2023 Apple Inc. -import math import unittest -from itertools import permutations import mlx.core as mx import mlx_tests @@ -12,7 +10,7 @@ import torch has_torch = True -except ImportError as e: +except ImportError: has_torch = False try: @@ -71,8 +69,11 @@ def np_fn(*args): ref_op = np_fn mlx_op = mlx_fn - ref_transform = lambda x: simple_transform(np_transform(x)) - mlx_transform = lambda x: simple_transform(mx.array(x).astype(mx.bfloat16)) + def ref_transform(x): + return simple_transform(np_transform(x)) + + def mlx_transform(x): + return simple_transform(mx.array(x).astype(mx.bfloat16)) self.__test_ops( ref_op, @@ -92,9 +93,10 @@ def torch_fn(*args): return out_bf16.to(torch.float32).numpy() ref_op = torch_fn - ref_transform = lambda x: simple_transform( - torch.from_numpy(x).to(torch.bfloat16) - ) + + def ref_transform(x): + return simple_transform(torch.from_numpy(x).to(torch.bfloat16)) + self.__test_ops( ref_op, mlx_op, diff --git a/Source/Cmlx/mlx/python/tests/test_blas.py b/Source/Cmlx/mlx/python/tests/test_blas.py index 7dbbaea04..eaed51dae 100644 --- a/Source/Cmlx/mlx/python/tests/test_blas.py +++ b/Source/Cmlx/mlx/python/tests/test_blas.py @@ -1,8 +1,6 @@ # Copyright © 2023-2024 Apple Inc. import math -import unittest -from itertools import permutations import mlx.core as mx import mlx_tests @@ -449,14 +447,26 @@ def test_matrix_vector_attn(self): shape_values = (B, ksl, Dk) # Prepare numpy arrays - q_np = np.random.uniform(-scale, scale, size=shape_queries).astype(np_dtype) - k_np = np.random.uniform(-scale, scale, size=shape_keys).astype(np_dtype) - v_np = np.random.uniform(-scale, scale, size=shape_values).astype(np_dtype) + q_np = np.random.uniform(-scale, scale, size=shape_queries).astype( + np_dtype + ) + k_np = np.random.uniform(-scale, scale, size=shape_keys).astype( + np_dtype + ) + v_np = np.random.uniform(-scale, scale, size=shape_values).astype( + np_dtype + ) # Rearrange to move heads up - q_np_reshape = q_np.reshape(B, qsl, n_kv_heads, factor, -1).transpose(0, 2, 3, 1, 4) - k_np_reshape = k_np.reshape(B, ksl, n_kv_heads, 1, -1).transpose(0, 2, 3, 4, 1) - v_np_reshape = v_np.reshape(B, ksl, n_kv_heads, 1, -1).transpose(0, 2, 3, 1, 4) + q_np_reshape = q_np.reshape( + B, qsl, n_kv_heads, factor, -1 + ).transpose(0, 2, 3, 1, 4) + k_np_reshape = k_np.reshape(B, ksl, n_kv_heads, 1, -1).transpose( + 0, 2, 3, 4, 1 + ) + v_np_reshape = v_np.reshape(B, ksl, n_kv_heads, 1, -1).transpose( + 0, 2, 3, 1, 4 + ) # Do attn style matmul s_np = q_np_reshape @ k_np_reshape @@ -469,9 +479,15 @@ def test_matrix_vector_attn(self): v_mx = mx.array(v_np) # Rearrange to move heads up - q_mx_reshape = q_mx.reshape(B, qsl, n_kv_heads, factor, -1).transpose(0, 2, 3, 1, 4) - k_mx_reshape = k_mx.reshape(B, ksl, n_kv_heads, 1, -1).transpose(0, 2, 3, 4, 1) - v_mx_reshape = v_mx.reshape(B, ksl, n_kv_heads, 1, -1).transpose(0, 2, 3, 1, 4) + q_mx_reshape = q_mx.reshape( + B, qsl, n_kv_heads, factor, -1 + ).transpose(0, 2, 3, 1, 4) + k_mx_reshape = k_mx.reshape(B, ksl, n_kv_heads, 1, -1).transpose( + 0, 2, 3, 4, 1 + ) + v_mx_reshape = v_mx.reshape(B, ksl, n_kv_heads, 1, -1).transpose( + 0, 2, 3, 1, 4 + ) # Do attn style matmul s_mx = q_mx_reshape @ k_mx_reshape @@ -598,8 +614,11 @@ def run_test(dtype, shape_a, shape_b, f_np_b, f_mx_b): 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) + def nt_np(b): + return b.swapaxes(-1, -2) + + def nt_mx(b): + return mx.swapaxes(b, -1, -2) for dtype in (mx.float32, mx.float16, mx.bfloat16): for M in (1, 2, 3, 5, 8, 11, 16): @@ -907,6 +926,32 @@ def run_test(dtype, B, M, K, N, c_shape, alpha, beta): 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_thin_matmul(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # few output columns or few rows with a transposed b takes the thin NAX + # kernel on Metal. + np.random.seed(0) + shapes = [(2048, 64, 2048), (2049, 48, 300), (4100, 16, 1), (1025, 24, 63)] + shapes += [(20, 300, 1), (33, 1000, 999), (64, 2560, 128)] + for M, N, K in shapes: + scale = K**-0.5 + a_np = np.random.normal(0.0, scale, (M, K)).astype(np.float32) + b_np = np.random.normal(0.0, scale, (N, K)).astype(np.float32) + c_np = np.random.normal(0.0, scale, (N,)).astype(np.float32) + for dtype in (mx.float16, mx.bfloat16): + a, b, c = (mx.array(x).astype(dtype) for x in (a_np, b_np, c_np)) + a_r, b_r, c_r = (np.array(x.astype(mx.float32)) for x in (a, b, c)) + ref = a_r @ b_r.T + for layout, b_mx in (("nt", b.T), ("nn", mx.contiguous(b.T))): + with self.subTest(shape=(M, N, K), dtype=str(dtype), layout=layout): + out = (a @ b_mx).astype(mx.float32) + self.assertTrue(np.allclose(out, ref, 1e-2, 0.05 * scale)) + out = mx.addmm(c, a, b_mx, 0.5, 2.0).astype(mx.float32) + ref_c = 0.5 * ref + 2.0 * c_r + self.assertTrue(np.allclose(out, ref_c, 1e-2, 0.125 * scale)) + def test_addmm_grad(self): def make_ref_addmm(alpha, beta): return lambda c, a, b: alpha * (a @ b) + beta * c @@ -1222,6 +1267,15 @@ def test_shape( self.assertTrue(np.allclose(c_mx, c_np, atol=1e-5)) self.assertTrue(np.allclose(e_mx, e_np, atol=1e-5)) + # Strided output mask with a vector lhs + a = mx.random.normal((1, 64)) + b = mx.random.normal((64, 128)) + mask = mx.array([[True, False, False, True, True, False, False, True]]) + out_mask = mask[:, ::2] + out = mx.block_masked_mm(a, b, 32, out_mask) + expected = ref_block_masked_mm(a, b, 32, out_mask) + self.assertTrue(mx.allclose(out, expected, atol=1e-5)) + def test_gather_matmul(self): def np_gather_mm(a, b, lhs_indices=None, rhs_indices=None): a = a.reshape((-1, a.shape[-2], a.shape[-1])) @@ -1593,6 +1647,14 @@ def segmented_mm_ref(a, b, s): c2 = mx.segmented_mm(a.T, b.T, segments) self.assertTrue(mx.allclose(c1, c2, atol=1e-4)) + # Negative leading dimension + a = mx.random.normal((8, 64)) + b = mx.random.normal((64, 16))[::-1] + segments = mx.array([[0, 20], [20, 64]], dtype=mx.uint32) + c1 = segmented_mm_ref(a, b, segments) + c2 = mx.segmented_mm(a, b, segments) + self.assertTrue(mx.allclose(c1, c2, atol=1e-4)) + with self.assertRaises(ValueError): a = mx.ones((2, 10, 10)) s = mx.array([[0, 5], [5, 10]]).astype(mx.uint32) diff --git a/Source/Cmlx/mlx/python/tests/test_compile.py b/Source/Cmlx/mlx/python/tests/test_compile.py index 185019342..7436e58cd 100644 --- a/Source/Cmlx/mlx/python/tests/test_compile.py +++ b/Source/Cmlx/mlx/python/tests/test_compile.py @@ -91,7 +91,10 @@ def test_compile_nonfinite_constants(self): def test_compile_float_constant_precision(self): x = mx.ones((4,), dtype=mx.float32) for constant in (1 / 3, 128**-0.5, 0.7071067811865476): - fun = lambda x, constant=constant: (x * x) * constant + + def fun(x, constant=constant): + return (x * x) * constant + self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) def test_compile_tuple_output_in_thread(self): @@ -301,7 +304,7 @@ def count_prims(outputs): buf = io.StringIO() mx.export_to_dot(buf, outputs) buf.seek(0) - return len([l for l in buf.read().split() if "label" in l]) + return len([line for line in buf.read().split() if "label" in line]) x = mx.array(1.0) cfun = mx.compile(fun) @@ -762,6 +765,24 @@ def fn(x): self.assertEqual(cfn(x4).shape, fn(x4).shape) self.assertEqual(cfn(x8).shape, fn(x8).shape) + def test_compile_scalar_constant_output(self): + for shapeless in [False, True]: + fun = mx.compile(lambda x: (x + 1.0, mx.array(1.0)), shapeless=shapeless) + out, c = fun(mx.array(3.0)) + self.assertEqual(out.item(), 4.0) + self.assertEqual(c.item(), 1.0) + + s = mx.array(7.0) + out, c = mx.compile(lambda x: (x + 7.0, s))(mx.array([1.0, 2.0])) + self.assertEqual(out.tolist(), [8.0, 9.0]) + self.assertEqual(c.item(), 7.0) + + # The gradient is the seed constant of the backward pass + fun = mx.compile(mx.value_and_grad(lambda x: x + 1.0)) + value, grad = fun(mx.array(2.0)) + self.assertEqual(value.item(), 3.0) + self.assertEqual(grad.item(), 1.0) + def test_compile_with_constant(self): # Test float @partial(mx.compile) @@ -902,10 +923,10 @@ def fun(x, y): return x + y.value with self.assertRaises(ValueError): - out = fun(mx.array(0.0), MyClass()) + fun(mx.array(0.0), MyClass()) with self.assertRaises(ValueError): - out = fun(mx.array(0.0), y=MyClass()) + fun(mx.array(0.0), y=MyClass()) def test_compile_create_list(self): @mx.compile @@ -933,7 +954,7 @@ def step(w): self.assertTrue(mx.allclose(expected[0], out[0])) self.assertTrue(mx.allclose(expected[1], out[1])) - def fun(w1, w2, x): + def fun(w1, w2, x): # noqa: F811 x = x @ w1 y = x @ w2 x = x + y * y @@ -1544,7 +1565,7 @@ def transform_tuple(t): def transform_vector(t): return Vector([t[0] + 10, t[1] * 10]) - x = State(mx.array(1), mx.array(2)) + _x = State(mx.array(1), mx.array(2)) compiled_transform = mx.compile(transform) compiled_transform_tuple = mx.compile(transform_tuple) @@ -1667,7 +1688,9 @@ def p(x): def test_compile_abs_unsigned(self): # abs has to compile for the wider unsigned types too - fun = lambda x: mx.abs(x) + 1 + def fun(x): + return mx.abs(x) + 1 + for dtype in [mx.uint8, mx.uint16, mx.uint32, mx.uint64]: x = mx.array([1, 2, 3], dtype) self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) diff --git a/Source/Cmlx/mlx/python/tests/test_constants.py b/Source/Cmlx/mlx/python/tests/test_constants.py index cfd971fbe..dfd856002 100644 --- a/Source/Cmlx/mlx/python/tests/test_constants.py +++ b/Source/Cmlx/mlx/python/tests/test_constants.py @@ -1,6 +1,5 @@ # Copyright © 2023 Apple Inc. -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_conv.py b/Source/Cmlx/mlx/python/tests/test_conv.py index 2a317e5b0..684074889 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv.py +++ b/Source/Cmlx/mlx/python/tests/test_conv.py @@ -3,7 +3,6 @@ import math import os import unittest -from itertools import permutations import mlx.core as mx import mlx_tests @@ -60,7 +59,7 @@ def test_torch_conv_1D(self): def run_conv1D( N, C, - O, + OC, iH, kH, stride, @@ -74,7 +73,7 @@ def run_conv1D( dtype=dtype, N=N, C=C, - O=O, + O=OC, iH=iH, kH=kH, stride=stride, @@ -85,7 +84,7 @@ def run_conv1D( np_dtype = getattr(np, dtype) np.random.seed(0) in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, int(C / groups))).astype( + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, int(C / groups))).astype( np_dtype ) @@ -116,7 +115,7 @@ def run_conv1D( self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -127,17 +126,17 @@ def run_conv1D( (3, 3, 1, 0), (31, 5, 5, 2), ): - run_conv1D(N, C, O, iH, kH, stride, padding, dtype=dtype) + run_conv1D(N, C, OC, iH, kH, stride, padding, dtype=dtype) # Groups tests - N, C, O = (4, 32, 64) + N, C, OC = (4, 32, 64) for iH, kH, stride, padding in ( (1, 1, 1, 0), (3, 3, 1, 0), (31, 5, 5, 2), ): for group in (1, 2, 4, 8, 16, 32): - run_conv1D(N, C, O, iH, kH, stride, padding, groups=group, dtype=dtype) + run_conv1D(N, C, OC, iH, kH, stride, padding, groups=group, dtype=dtype) # Strided inputs tests for tpose_in, tpose_wt in ( @@ -169,7 +168,7 @@ def test_torch_conv_1D_grad(self): def run_conv1D_grad( N, C, - O, + OC, iH, kH, stride, @@ -183,7 +182,7 @@ def run_conv1D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, iH=iH, kH=kH, stride=stride, @@ -196,8 +195,8 @@ def run_conv1D_grad( oH = 1 + ((iH + 2 * padding - dilation * (kH - 1) - 1) // stride) in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) - ct_np = np.random.normal(0, 1.0 / C, (N, oH, O)).astype(np_dtype) + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, C)).astype(np_dtype) + ct_np = np.random.normal(0, 1.0 / C, (N, oH, OC)).astype(np_dtype) in_mx, wt_mx, ct_mx = map(mx.array, (in_np, wt_np, ct_np)) in_pt, wt_pt, ct_pt = map( @@ -257,7 +256,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -268,14 +267,14 @@ def f(a, b): (3, 3, 1, 0), (31, 5, 5, 2), ): - run_conv1D_grad(N, C, O, iH, kH, stride, padding, dtype=dtype) + run_conv1D_grad(N, C, OC, iH, kH, stride, padding, dtype=dtype) @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_2D(self): def run_conv2D( N, C, - O, + OC, idim, kdim, stride, @@ -288,7 +287,7 @@ def run_conv2D( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -301,7 +300,7 @@ def run_conv2D( kH, kW = kdim scale = 1.0 / math.sqrt(kH * kW * C) in_np = np.random.normal(0.0, scale, (N, iH, iW, C)) - wt_np = np.random.normal(0.0, 1.0, (O, kH, kW, int(C / groups))) + wt_np = np.random.normal(0.0, 1.0, (OC, kH, kW, int(C / groups))) mx_dtype = getattr(mx, dtype) torch_dtype = getattr(torch, dtype) @@ -341,13 +340,13 @@ def run_conv2D( self.assertEqual(out_pt.shape, out_mx.shape) if dtype == "bfloat16": - atol, rtol = 1e-1, 1e-3 + atol = 1e-1 else: - atol, rtol = 1e-5, 1e-6 + atol = 1e-5 self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32", "bfloat16"): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -358,10 +357,10 @@ def run_conv2D( ((3, 3), (3, 1), (1, 1), (0, 0)), ((31, 31), (5, 5), (5, 5), (2, 2)), ): - run_conv2D(N, C, O, idim, kdim, stride, padding, dtype=dtype) + run_conv2D(N, C, OC, idim, kdim, stride, padding, dtype=dtype) # Groups tests - N, C, O = (4, 32, 64) + N, C, OC = (4, 32, 64) for idim, kdim, stride, padding in ( ((1, 1), (1, 1), (1, 1), (0, 0)), ((3, 3), (3, 1), (1, 1), (0, 0)), @@ -369,7 +368,7 @@ def run_conv2D( ): for group in (1, 2, 4, 8, 16, 32): run_conv2D( - N, C, O, idim, kdim, stride, padding, groups=group, dtype=dtype + N, C, OC, idim, kdim, stride, padding, groups=group, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -377,7 +376,7 @@ def test_torch_conv_2D_grad(self): def run_conv2D_grad( N, C, - O, + OC, idim, kdim, stride, @@ -391,7 +390,7 @@ def run_conv2D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -413,10 +412,10 @@ def run_conv2D_grad( ) in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) - wt_np = np.random.normal(0.0, scale, (O, kH, kW, C // groups)).astype( + wt_np = np.random.normal(0.0, scale, (OC, kH, kW, C // groups)).astype( np_dtype ) - ct_np = np.random.normal(0.0, scale, (N, oH, oW, O)).astype(np_dtype) + ct_np = np.random.normal(0.0, scale, (N, oH, oW, OC)).astype(np_dtype) in_mx, wt_mx, ct_mx = map(mx.array, (in_np, wt_np, ct_np)) in_pt, wt_pt, ct_pt = map( @@ -471,7 +470,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): for idim, kdim, stride, padding, dilation in ( ((1, 1), (1, 1), (1, 1), (0, 0), (1, 1)), ((3, 3), (3, 1), (1, 1), (0, 0), (1, 1)), @@ -481,7 +480,7 @@ def f(a, b): ((32, 32), (3, 3), (2, 2), (1, 1), (3, 2)), ): run_conv2D_grad( - N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dilation, dtype=dtype ) for groups in (2, 4, 32): @@ -497,12 +496,27 @@ def f(a, b): dtype=dtype, ) + # Depthwise with C % 16 != 0 and grouped with C_per_group = 8 + for N, C, OC, groups in ((2, 24, 24, 24), (2, 40, 40, 40), (2, 64, 128, 8)): + for stride in ((1, 1), (2, 2)): + run_conv2D_grad( + N, + C, + OC, + (16, 16), + (3, 3), + stride, + (1, 1), + groups=groups, + dtype=dtype, + ) + @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_3D(self): def run_conv3D( N, C, - O, + OC, idim, kdim, stride, @@ -516,7 +530,7 @@ def run_conv3D( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -532,7 +546,7 @@ def run_conv3D( in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( np_dtype ) - wt_np = np.random.normal(0.0, 1.0, (O, kD, kH, kW, C)).astype(np_dtype) + wt_np = np.random.normal(0.0, 1.0, (OC, kD, kH, kW, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt, wt_pt = map( @@ -562,7 +576,7 @@ def run_conv3D( self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -573,12 +587,12 @@ def run_conv3D( ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), ((31, 31, 31), (5, 5, 5), (5, 5, 5), (2, 2, 2)), ): - run_conv3D(N, C, O, idim, kdim, stride, padding, dtype=dtype) + run_conv3D(N, C, OC, idim, kdim, stride, padding, dtype=dtype) - N, C, O = (2, 4, 4) + N, C, OC = (2, 4, 4) idim, kdim, stride, padding = (6, 6, 6), (3, 1, 1), (1, 1, 1), (0, 0, 0) run_conv3D( - N, C, O, idim, kdim, stride, padding, dilation=(2, 2, 2), dtype=dtype + N, C, OC, idim, kdim, stride, padding, dilation=(2, 2, 2), dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -586,7 +600,7 @@ def test_torch_conv_3D_grad(self): def run_conv3D_grad( N, C, - O, + OC, idim, kdim, stride, @@ -600,7 +614,7 @@ def run_conv3D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -627,10 +641,10 @@ def run_conv3D_grad( in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( np_dtype ) - wt_np = np.random.normal(0.0, scale, (O, kD, kH, kW, C)).astype( + wt_np = np.random.normal(0.0, scale, (OC, kD, kH, kW, C)).astype( np_dtype ) - ct_np = np.random.normal(0.0, scale, (N, oD, oH, oW, O)).astype( + ct_np = np.random.normal(0.0, scale, (N, oD, oH, oW, OC)).astype( np_dtype ) @@ -687,7 +701,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 16, 32), (4, 8, 16)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 16, 32), (4, 8, 16)): for idim, kdim, stride, padding, dilation in ( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), @@ -697,7 +711,7 @@ def f(a, b): ((16, 16, 16), (3, 3, 3), (2, 2, 2), (1, 1, 1), (3, 2, 2)), ): run_conv3D_grad( - N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dilation, dtype=dtype ) def __conv_general_test( @@ -753,7 +767,9 @@ def conv_general_pt( ): C = inp.size()[1] ndim = inp.ndim - 2 - map_ints = lambda x: [x] * ndim if isinstance(x, int) else x + + def map_ints(x): + return [x] * ndim if isinstance(x, int) else x stride, padding, kernel_dilation, input_dilation = map( map_ints, (stride, padding, kernel_dilation, input_dilation) @@ -975,7 +991,6 @@ def fn(x, w): def fn_gt(x, w): num_groups = x.shape[-1] // w.shape[-1] - group_size = w.shape[-1] ws = w.reshape(num_groups, -1, *w.shape[1:]).split(num_groups) xs = x.reshape(*x.shape[:-1], num_groups, -1).split(num_groups, axis=-2) return mx.concatenate( @@ -1025,7 +1040,6 @@ def fn(x, w): def fn_gt(x, w): num_groups = x.shape[-1] // w.shape[-1] - group_size = w.shape[-1] ws = w.reshape(num_groups, -1, *w.shape[1:]).split(num_groups) xs = x.reshape(*x.shape[:-1], num_groups, -1).split(num_groups, axis=-2) return mx.concatenate( @@ -1079,6 +1093,9 @@ def test_torch_conv_depthwise(self): ( 1, 32, 32, 32, 7, 7, 32, (1, 1), (3, 3), 32), ( 3, 32, 32, 32, 5, 5, 32, (1, 2), (0, 0), 32), ( 1, 32, 32, 32, 7, 7, 32, (2, 1), (1, 3), 32), + ( 2, 16, 16, 24, 3, 3, 24, (1, 1), (1, 1), 24), + ( 1, 14, 14, 40, 5, 5, 40, (2, 2), (2, 2), 40), + ( 2, 9, 9, 8, 3, 3, 8, (1, 1), (1, 1), 8), ) # fmt: on @@ -1086,14 +1103,14 @@ def test_torch_conv_depthwise(self): if mx.default_device() == mx.gpu: dtypes += [np.float16] - for N, H, W, C, kH, kW, O, strides, padding, groups in shapes: + for N, H, W, C, kH, kW, OC, strides, padding, groups in shapes: for dtype in dtypes: for flip in [False, True]: Cw = C // groups self.__conv_general_test( (N, H, W, C), - (O, kH, kW, Cw), + (OC, kH, kW, Cw), strides, padding, kernel_dilation=1, @@ -1168,6 +1185,35 @@ def loss_fn(kernel, inputs, strides, groups): ) self.assertEqual(grads.shape, k_shape) + def test_padding_input_grad(self): + def conv_loss(x, w, stride, padding): + out = mx.conv_general(x, w, stride=stride, padding=padding) + return mx.sum(out * out) + + def padded_loss(x, w, stride, padding): + lo, hi = padding + pad_width = [(0, 0)] + list(zip(lo, hi)) + [(0, 0)] + out = mx.conv_general(mx.pad(x, pad_width), w, stride=stride) + return mx.sum(out * out) + + for in_shape, w_shape, stride, padding in [ + ((1, 4, 1), (1, 3, 1), (1,), ([0], [1])), + ((1, 4, 1), (1, 3, 1), (2,), ([0], [1])), + ((1, 4, 1), (1, 3, 1), (1,), ([1], [0])), + ((2, 6, 6, 3), (4, 3, 3, 3), (2, 2), ([0, 1], [1, 2])), + ((2, 5, 5, 5, 2), (3, 2, 2, 2, 2), (2, 1, 2), ([0, 0, 1], [1, 1, 0])), + ((1, 4, 1), (1, 2, 1), (1,), ([3], [3])), + ((1, 4, 1), (1, 2, 1), (1,), ([0], [3])), + ((1, 4, 1), (1, 2, 1), (2,), ([3], [0])), + ((2, 5, 6, 3), (4, 2, 3, 3), (1, 2), ([4, 0], [1, 5])), + ]: + x = mx.random.normal(in_shape) + w = mx.random.normal(w_shape) + expected = mx.grad(padded_loss)(x, w, stride, padding) + grad = mx.grad(conv_loss)(x, w, stride, padding) + self.assertEqual(grad.shape, x.shape) + self.assertTrue(mx.allclose(grad, expected, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "flaky in CI") def test_conv_1d_with_2d(self): x = mx.random.uniform(shape=(2, 10, 16)) @@ -1215,6 +1261,177 @@ def test_conv2d_unaligned_channels(self): y_hat = mx.conv_transpose2d(x, w) self.assertTrue(mx.allclose(y, y_hat)) + def __assert_gpu_matches_cpu(self, f, *args, **kwargs): + out_gpu = f(*args, **kwargs, stream=mx.gpu) + out_cpu = f(*args, **kwargs, stream=mx.cpu) + self.assertEqual(out_gpu.shape, out_cpu.shape) + atol = 1e-4 if args[0].dtype == mx.float32 else 5e-3 + self.assertTrue( + mx.allclose( + out_gpu.astype(mx.float32), out_cpu.astype(mx.float32), atol=atol + ) + ) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv_grouped_c_per_group_8(self): + # fmt: off + shapes_2d = ( + # N, H, W, groups, O_per_group, k, stride, padding, dilation + ( 2, 14, 14, 2, 8, 3, 1, 1, 1), + ( 2, 13, 9, 16, 16, 3, 2, 1, 1), + ( 1, 7, 7, 32, 8, 3, 1, 1, 1), + ( 2, 14, 14, 32, 16, 3, 2, 1, 1), + ( 2, 14, 14, 16, 8, 5, 1, 2, 1), + ( 2, 14, 14, 16, 8, 3, 1, 2, 2), + ( 2, 14, 14, 8, 32, 3, 1, 1, 1), + ( 2, 14, 14, 8, 64, 3, 1, 1, 1), + ( 1, 33, 31, 4, 8, 3, 2, 0, 1), + ( 1, 9, 9, 1, 8, 3, 1, 1, 1), + ) + shapes_1d = ( + # N, L, groups, O_per_group, k, stride, padding + ( 2, 64, 2, 8, 3, 1, 1), + ( 2, 63, 16, 16, 3, 2, 1), + ( 1, 100, 32, 8, 5, 1, 2), + ( 2, 64, 32, 16, 3, 2, 0), + ) + # fmt: on + + for dtype in (mx.float32, mx.float16): + for N, H, W, groups, opg, k, stride, padding, dilation in shapes_2d: + C, OC = 8 * groups, opg * groups + for flip in (False, True): + x = mx.random.normal((N, H, W, C)).astype(dtype) + w = (0.2 * mx.random.normal((OC, k, k, 8))).astype(dtype) + with self.subTest( + dtype=dtype, + in_shape=x.shape, + wt_shape=w.shape, + groups=groups, + stride=stride, + padding=padding, + dilation=dilation, + flip=flip, + ): + self.__assert_gpu_matches_cpu( + mx.conv_general, + x, + w, + stride=stride, + padding=padding, + kernel_dilation=dilation, + groups=groups, + flip=flip, + ) + + for N, L, groups, opg, k, stride, padding in shapes_1d: + C, OC = 8 * groups, opg * groups + x = mx.random.normal((N, L, C)).astype(dtype) + w = (0.2 * mx.random.normal((OC, k, 8))).astype(dtype) + with self.subTest(dtype=dtype, in_shape=x.shape, wt_shape=w.shape): + self.__assert_gpu_matches_cpu( + mx.conv1d, + x, + w, + stride=stride, + padding=padding, + groups=groups, + ) + + # Non-contiguous input and weight go through a contiguous copy + x = mx.random.normal((2, 128, 14, 14)).transpose(0, 2, 3, 1) + w = 0.2 * mx.random.normal((3, 3, 8, 256)).transpose(3, 0, 1, 2) + self.__assert_gpu_matches_cpu(mx.conv2d, x, w, padding=1, groups=16) + + # Gradients with respect to the input and the weight + def vjp(f, x, w, stream): + out, grads = mx.vjp( + lambda a, b: f(a, b, stream=stream), + [x, w], + [mx.ones(f(x, w, stream=mx.cpu).shape)], + ) + return grads + + x = mx.random.normal((2, 14, 14, 128)) + w = 0.2 * mx.random.normal((256, 3, 3, 8)) + + def f(a, b, stream): + return mx.conv2d(a, b, stride=2, padding=1, groups=16, stream=stream) + + for g_gpu, g_cpu in zip(vjp(f, x, w, mx.gpu), vjp(f, x, w, mx.cpu)): + self.assertTrue(mx.allclose(g_gpu, g_cpu, atol=1e-4)) + + x = mx.random.normal((2, 64, 128)) + w = 0.2 * mx.random.normal((128, 3, 8)) + + def f(a, b, stream): + return mx.conv1d(a, b, padding=1, groups=16, stream=stream) + + for g_gpu, g_cpu in zip(vjp(f, x, w, mx.gpu), vjp(f, x, w, mx.cpu)): + self.assertTrue(mx.allclose(g_gpu, g_cpu, atol=1e-4)) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv_depthwise_c_multiple_of_8(self): + # fmt: off + shapes = ( + # N, H, W, C, k, stride, padding + ( 2, 16, 16, 24, 3, 1, 1), + ( 2, 17, 15, 24, 3, 2, 1), + ( 1, 14, 14, 40, 5, 1, 2), + ( 1, 14, 14, 40, 5, 2, 2), + ( 2, 9, 9, 8, 3, 1, 1), + ( 1, 12, 12, 56, 7, 1, 3), + ) + # fmt: on + + for dtype in (mx.float32, mx.float16): + for N, H, W, C, k, stride, padding in shapes: + for flip in (False, True): + x = mx.random.normal((N, H, W, C)).astype(dtype) + w = (0.2 * mx.random.normal((C, k, k, 1))).astype(dtype) + with self.subTest( + dtype=dtype, + in_shape=x.shape, + wt_shape=w.shape, + stride=stride, + padding=padding, + flip=flip, + ): + self.__assert_gpu_matches_cpu( + mx.conv_general, + x, + w, + stride=stride, + padding=padding, + groups=C, + flip=flip, + ) + + # conv1d depthwise with padding also uses the 2D depthwise kernel + for C, k, stride in ((24, 3, 1), (40, 5, 2)): + x = mx.random.normal((2, 64, C)).astype(dtype) + w = (0.2 * mx.random.normal((C, k, 1))).astype(dtype) + with self.subTest(dtype=dtype, in_shape=x.shape, wt_shape=w.shape): + self.__assert_gpu_matches_cpu( + mx.conv1d, x, w, stride=stride, padding=k // 2, groups=C + ) + + # Gradients with respect to the input and the weight + for C, k, stride in ((24, 3, 1), (40, 5, 2)): + x = mx.random.normal((2, 16, 16, C)) + w = 0.2 * mx.random.normal((C, k, k, 1)) + + def f(a, b, stream): + return mx.conv2d( + a, b, stride=stride, padding=k // 2, groups=C, stream=stream + ) + + cotan = mx.ones(f(x, w, mx.cpu).shape) + _, grads_gpu = mx.vjp(lambda a, b: f(a, b, mx.gpu), [x, w], [cotan]) + _, grads_cpu = mx.vjp(lambda a, b: f(a, b, mx.cpu), [x, w], [cotan]) + for g_gpu, g_cpu in zip(grads_gpu, grads_cpu): + self.assertTrue(mx.allclose(g_gpu, g_cpu, atol=1e-4)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_conv2d_winograd_batch_tiling(self): # Use envs to test tiling without allocating large buffers. @@ -1268,13 +1485,13 @@ def run(x, w, env={}): # Budget for ~2 batch elements so the selector itself must # tile; mirrors the winograd_batch_step arithmetic. n, iH, iW, C = in_shape - O = wt_shape[0] + OC = wt_shape[0] pH = 6 * ((iH + 2 - 2 + 5) // 6) + 2 pW = 6 * ((iW + 2 - 2 + 5) // 6) + 2 per_n = ( - pH * pW * C * 4 + 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + O) * 4 + pH * pW * C * 4 + 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + OC) * 4 ) - used = (n * iH * iW * (C + O) + 64 * C * O) * 4 + used = (n * iH * iW * (C + OC) + 64 * C * OC) * 4 with self.subTest(in_shape=in_shape, budget="tiled"): budget = str(int((used + 5 * per_n // 2) / 0.75)) tiled = run(x, w, {ws_key: budget}) @@ -1293,6 +1510,25 @@ def run(x, w, env={}): self.assertFalse(np.array_equal(untiled, capped)) self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv2d_winograd_disable(self): + # N * iH * iW >= 4096 selects Winograd only from a batch of 3. + np.random.seed(0) + x = mx.array(np.random.normal(size=(1, 40, 40, 128)).astype(np.float32)) + w = mx.array( + (np.random.normal(size=(128, 3, 3, 128)) * 0.05).astype(np.float32) + ) + xb = mx.repeat(x, 3, axis=0) + + single = np.array(mx.conv2d(x, w, padding=1)) + winograd = np.array(mx.conv2d(xb, w, padding=1)) + self.assertFalse(np.array_equal(winograd[0], single[0])) + + with mlx_tests.scoped_env(MLX_CONV_WINOGRAD="0"): + gemm = np.array(mx.conv2d(xb, w, padding=1)) + for i in range(3): + self.assertTrue(np.array_equal(gemm[i], single[0])) + 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)) @@ -1326,6 +1562,78 @@ def test_conv_3D_small_kd_decomposition(self): f"C{Cin}->{Cout} k{kd}{kh}{kw} flip={flip}", ) + def test_grouped_conv_padding_stays_in_output(self): + # groups takes the direct CPU kernel. Padding taller than the output + # used to run the border loop past the buffer. + mx.random.seed(0) + x = mx.random.normal((1, 1, 8, 2)) + w = mx.random.normal((2, 5, 1, 1)) + y = mx.conv2d(x, w, stride=(1, 1), padding=(2, 0), groups=2, stream=mx.cpu) + parts = [ + mx.conv2d( + x[..., g : g + 1], + w[g : g + 1], + stride=(1, 1), + padding=(2, 0), + stream=mx.cpu, + ) + for g in range(2) + ] + expected = mx.concatenate(parts, axis=-1) + self.assertEqual(tuple(y.shape), (1, 1, 8, 2)) + self.assertTrue(mx.allclose(y, expected, atol=1e-4)) + + # A 1-wide kernel makes weight dilation a no-op, but any dilation + # other than 1 sends 3D conv to the direct kernel. Depth padding is + # the same overhang as the 2D case. + x3 = mx.random.normal((1, 1, 4, 2, 1)) + w3 = mx.random.normal((1, 5, 1, 1, 1)) + y3 = mx.conv3d( + x3, + w3, + padding=(2, 0, 0), + dilation=(1, 2, 1), + stream=mx.cpu, + ) + expected3 = mx.conv3d( + x3, + w3, + padding=(2, 0, 0), + stream=mx.cpu, + ) + self.assertEqual(tuple(y3.shape), tuple(expected3.shape)) + self.assertTrue(mx.allclose(y3, expected3, atol=1e-4)) + + wide = mx.random.normal((1, 1, 500000, 2)) + wide_w = mx.random.normal((2, 5, 1, 1)) + wide_y = mx.conv2d( + wide, + wide_w, + stride=(1, 1), + padding=(2, 0), + groups=2, + stream=mx.cpu, + ) + mx.eval(wide_y) + self.assertEqual(tuple(wide_y.shape), (1, 1, 500000, 2)) + + @unittest.skipIf(not mx.is_available(mx.gpu), "requires GPU") + def test_grouped_depthwise_conv_padding_bounds(self): + x = mx.random.normal((1, 16, 64)) + w = mx.random.normal((64, 3, 1)) + y = mx.conv1d(x, w, padding=1, groups=64) + mx.eval(y) + self.assertEqual(tuple(y.shape), (1, 16, 64)) + + @unittest.skipIf(not mx.is_available(mx.gpu), "requires GPU") + def test_conv2d_large_output_channels(self): + x = mx.random.normal((1, 33, 33, 5)) + w = mx.random.normal((32800, 1, 1, 5)) + yg = mx.conv2d(x, w, stride=(2, 2), stream=mx.gpu) + yc = mx.conv2d(x, w, stride=(2, 2), stream=mx.cpu) + mx.eval(yg, yc) + self.assertTrue(mx.allclose(yg, yc, atol=1e-3)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py index 28593a530..07ce6e74b 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py +++ b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py @@ -1,9 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math -import os import unittest -from itertools import permutations import mlx.core as mx import mlx_tests @@ -14,7 +12,7 @@ import torch.nn.functional as F has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -24,7 +22,7 @@ def test_torch_conv_transpose_1D(self): def run_conv_transpose_1D( N, C, - O, + OC, iH, kH, stride, @@ -39,7 +37,7 @@ def run_conv_transpose_1D( dtype=dtype, N=N, C=C, - O=O, + O=OC, iH=iH, kH=kH, stride=stride, @@ -50,7 +48,7 @@ def run_conv_transpose_1D( np_dtype = getattr(np, dtype) np.random.seed(0) in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, int(C / groups))).astype( + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, int(C / groups))).astype( np_dtype ) @@ -80,7 +78,7 @@ def run_conv_transpose_1D( self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -91,10 +89,12 @@ def run_conv_transpose_1D( (3, 3, 1, 0), (31, 5, 5, 2), ): - run_conv_transpose_1D(N, C, O, iH, kH, stride, padding, dtype=dtype) + run_conv_transpose_1D( + N, C, OC, iH, kH, stride, padding, dtype=dtype + ) # Groups tests - N, C, O = (4, 32, 64) + N, C, OC = (4, 32, 64) for iH, kH, stride, padding in ( (1, 1, 1, 0), (3, 3, 1, 0), @@ -102,7 +102,7 @@ def run_conv_transpose_1D( ): for group in (1,): run_conv_transpose_1D( - N, C, O, iH, kH, stride, padding, groups=group, dtype=dtype + N, C, OC, iH, kH, stride, padding, groups=group, dtype=dtype ) # Strided inputs tests @@ -133,7 +133,7 @@ def test_torch_conv_transpose_1D_grad(self): def run_conv_transpose1D_grad( N, C, - O, + OC, iH, kH, stride, @@ -147,7 +147,7 @@ def run_conv_transpose1D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, iH=iH, kH=kH, stride=stride, @@ -160,7 +160,7 @@ def run_conv_transpose1D_grad( # oH = 1 + ((iH + 2 * padding - dilation * (kH - 1) - 1) // stride) in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)).requires_grad_(True) @@ -211,7 +211,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -223,7 +223,7 @@ def f(a, b): (31, 5, 5, 2), ): run_conv_transpose1D_grad( - N, C, O, iH, kH, stride, padding, dtype=dtype + N, C, OC, iH, kH, stride, padding, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -231,7 +231,7 @@ def test_torch_conv_transpose_2D(self): def run_conv_transpose2D( N, C, - O, + OC, idim, kdim, stride, @@ -245,7 +245,7 @@ def run_conv_transpose2D( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -259,9 +259,9 @@ def run_conv_transpose2D( kH, kW = kdim scale = 1.0 / math.sqrt(kH * kW * C) in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) - wt_np = np.random.normal(0.0, 1.0, (O, kH, kW, int(C / groups))).astype( - np_dtype - ) + wt_np = np.random.normal( + 0.0, 1.0, (OC, kH, kW, int(C / groups)) + ).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu") @@ -289,7 +289,7 @@ def run_conv_transpose2D( self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -301,11 +301,11 @@ def run_conv_transpose2D( ((31, 31), (5, 5), (5, 5), (2, 2)), ): run_conv_transpose2D( - N, C, O, idim, kdim, stride, padding, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dtype=dtype ) # Groups tests - N, C, O = (4, 32, 64) + N, C, OC = (4, 32, 64) for idim, kdim, stride, padding in ( ((1, 1), (1, 1), (1, 1), (0, 0)), ((3, 3), (3, 1), (1, 1), (0, 0)), @@ -313,7 +313,7 @@ def run_conv_transpose2D( ): for group in (1,): run_conv_transpose2D( - N, C, O, idim, kdim, stride, padding, groups=group, dtype=dtype + N, C, OC, idim, kdim, stride, padding, groups=group, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -321,7 +321,7 @@ def test_torch_conv_transpose_2D_grad(self): def run_conv_transpose2D_grad( N, C, - O, + OC, idim, kdim, stride, @@ -335,7 +335,7 @@ def run_conv_transpose2D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -347,10 +347,10 @@ def run_conv_transpose2D_grad( np.random.seed(0) iH, iW = idim kH, kW = kdim - scale = 1.0 / math.sqrt(kH * kW * C * O) + scale = 1.0 / math.sqrt(kH * kW * C * OC) in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) - wt_np = np.random.normal(0.0, scale, (O, kH, kW, C)).astype(np_dtype) + wt_np = np.random.normal(0.0, scale, (OC, kH, kW, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).requires_grad_( @@ -400,7 +400,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): for idim, kdim, stride, padding, dilation in ( ((1, 1), (1, 1), (1, 1), (0, 0), (1, 1)), ((3, 3), (3, 1), (1, 1), (0, 0), (1, 1)), @@ -410,7 +410,7 @@ def f(a, b): ((32, 32), (3, 3), (2, 2), (1, 1), (3, 2)), ): run_conv_transpose2D_grad( - N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dilation, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -418,7 +418,7 @@ def test_torch_conv_transpose_3D(self): def run_conv_transpose3D( N, C, - O, + OC, idim, kdim, stride, @@ -432,7 +432,7 @@ def run_conv_transpose3D( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -444,11 +444,11 @@ def run_conv_transpose3D( np.random.seed(0) iD, iH, iW = idim kD, kH, kW = kdim - scale = 1.0 / math.sqrt(kD * kH * kW * C * O) + scale = 1.0 / math.sqrt(kD * kH * kW * C * OC) in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( np_dtype ) - wt_np = np.random.normal(0.0, 1.0, (O, kD, kH, kW, C)).astype(np_dtype) + wt_np = np.random.normal(0.0, 1.0, (OC, kD, kH, kW, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)) @@ -476,7 +476,7 @@ def run_conv_transpose3D( self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ( + for N, C, OC in ( (1, 1, 1), (1, 6, 1), (1, 1, 6), @@ -490,7 +490,7 @@ def run_conv_transpose3D( ((3, 4, 5), (2, 2, 2), (2, 2, 2), (0, 0, 0)), ): run_conv_transpose3D( - N, C, O, idim, kdim, stride, padding, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -498,7 +498,7 @@ def test_torch_conv_transpose_3D_grad(self): def run_conv_transpose3D_grad( N, C, - O, + OC, idim, kdim, stride, @@ -512,7 +512,7 @@ def run_conv_transpose3D_grad( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -524,12 +524,12 @@ def run_conv_transpose3D_grad( np.random.seed(0) iD, iH, iW = idim kD, kH, kW = kdim - scale = 1.0 / math.sqrt(kD * kH * kW * C * O) + scale = 1.0 / math.sqrt(kD * kH * kW * C * OC) in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( np_dtype ) - wt_np = np.random.normal(0.0, scale, (O, kD, kH, kW, C)).astype( + wt_np = np.random.normal(0.0, scale, (OC, kD, kH, kW, C)).astype( np_dtype ) @@ -586,7 +586,7 @@ def f(a, b): self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (2, 4, 8), (2, 8, 16)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (2, 4, 8), (2, 8, 16)): for idim, kdim, stride, padding, dilation in ( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), @@ -596,19 +596,28 @@ def f(a, b): ((8, 8, 8), (3, 3, 3), (2, 2, 2), (1, 1, 1), (3, 2, 2)), ): run_conv_transpose3D_grad( - N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype + N, C, OC, idim, kdim, stride, padding, dilation, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_tranpose_1d_output_padding(self): def run_conv_transpose_1d_output_padding( - N, C, O, iH, kH, stride, padding, output_padding, dtype="float32", atol=1e-5 + N, + C, + OC, + iH, + kH, + stride, + padding, + output_padding, + dtype="float32", + atol=1e-5, ): with self.subTest( dtype=dtype, N=N, C=C, - O=O, + O=OC, iH=iH, kH=kH, stride=stride, @@ -618,7 +627,7 @@ def run_conv_transpose_1d_output_padding( np_dtype = getattr(np, dtype) np.random.seed(0) in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)) @@ -645,14 +654,14 @@ def run_conv_transpose_1d_output_padding( self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): for iH, kH, stride, padding, output_padding in ( (3, 2, 2, 0, 1), (5, 3, 2, 1, 0), (7, 4, 3, 1, 2), ): run_conv_transpose_1d_output_padding( - N, C, O, iH, kH, stride, padding, output_padding, dtype=dtype + N, C, OC, iH, kH, stride, padding, output_padding, dtype=dtype ) @unittest.skipIf(not has_torch, "requires Torch") @@ -660,7 +669,7 @@ def test_torch_conv_transpose_2d_output_padding(self): def run_conv_transpose_2d_output_padding( N, C, - O, + OC, idim, kdim, stride, @@ -673,7 +682,7 @@ def run_conv_transpose_2d_output_padding( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -685,7 +694,7 @@ def run_conv_transpose_2d_output_padding( iH, iW = idim kH, kW = kdim in_np = np.random.normal(0, 1.0 / C, (N, iH, iW, C)).astype(np_dtype) - wt_np = np.random.normal(0, 1.0 / C, (O, kH, kW, C)).astype(np_dtype) + wt_np = np.random.normal(0, 1.0 / C, (OC, kH, kW, C)).astype(np_dtype) in_mx, wt_mx = map(mx.array, (in_np, wt_np)) in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)) @@ -712,7 +721,7 @@ def run_conv_transpose_2d_output_padding( self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): for idim, kdim, stride, padding, output_padding in ( ((3, 3), (2, 2), (2, 2), (0, 0), (1, 1)), ((5, 5), (3, 3), (2, 2), (1, 1), (0, 0)), @@ -721,7 +730,7 @@ def run_conv_transpose_2d_output_padding( run_conv_transpose_2d_output_padding( N, C, - O, + OC, idim, kdim, stride, @@ -735,7 +744,7 @@ def test_torch_conv_transpose_3d_output_padding(self): def run_conv_transpose_3d_output_padding( N, C, - O, + OC, idim, kdim, stride, @@ -748,7 +757,7 @@ def run_conv_transpose_3d_output_padding( dtype=dtype, N=N, C=C, - O=O, + O=OC, idim=idim, kdim=kdim, stride=stride, @@ -762,7 +771,7 @@ def run_conv_transpose_3d_output_padding( in_np = np.random.normal(0, 1.0 / C, (N, iD, iH, iW, C)).astype( np_dtype ) - wt_np = np.random.normal(0, 1.0 / C, (O, kD, kH, kW, C)).astype( + wt_np = np.random.normal(0, 1.0 / C, (OC, kD, kH, kW, C)).astype( np_dtype ) @@ -790,7 +799,7 @@ def run_conv_transpose_3d_output_padding( self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) for dtype in ("float32",): - for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): + for N, C, OC in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): for idim, kdim, stride, padding, output_padding in ( ((3, 3, 3), (2, 2, 2), (2, 2, 2), (0, 0, 0), (1, 1, 1)), ((5, 5, 5), (3, 3, 3), (2, 2, 2), (1, 1, 1), (0, 0, 0)), @@ -799,7 +808,7 @@ def run_conv_transpose_3d_output_padding( run_conv_transpose_3d_output_padding( N, C, - O, + OC, idim, kdim, stride, diff --git a/Source/Cmlx/mlx/python/tests/test_double.py b/Source/Cmlx/mlx/python/tests/test_double.py index 3186e7e57..7cae2e19e 100644 --- a/Source/Cmlx/mlx/python/tests/test_double.py +++ b/Source/Cmlx/mlx/python/tests/test_double.py @@ -1,8 +1,6 @@ # Copyright © 2024 Apple Inc. import math -import os -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_einsum.py b/Source/Cmlx/mlx/python/tests/test_einsum.py index 08884ae98..d9baa3a77 100644 --- a/Source/Cmlx/mlx/python/tests/test_einsum.py +++ b/Source/Cmlx/mlx/python/tests/test_einsum.py @@ -1,6 +1,5 @@ # Copyright © 2024 Apple Inc. -import unittest import mlx.core as mx import mlx_tests @@ -8,7 +7,6 @@ class TestEinsum(mlx_tests.MLXTestCase): - def test_simple_path(self): a = mx.zeros((5, 5)) path = mx.einsum_path("ii", a) diff --git a/Source/Cmlx/mlx/python/tests/test_eval.py b/Source/Cmlx/mlx/python/tests/test_eval.py index 265a090aa..9801f82c7 100644 --- a/Source/Cmlx/mlx/python/tests/test_eval.py +++ b/Source/Cmlx/mlx/python/tests/test_eval.py @@ -1,7 +1,6 @@ # Copyright © 2023 Apple Inc. import unittest -from functools import partial import mlx.core as mx import mlx_tests @@ -115,10 +114,10 @@ def fn(x): mx.eval(a) return a - out = mx.vjp(fn, (x,), (y,)) - out = mx.vjp(fn, (x,), (y,)) + _out = mx.vjp(fn, (x,), (y,)) + _out = mx.vjp(fn, (x,), (y,)) peak_mem = mx.get_peak_memory() - out = mx.vjp(fn, (x,), (y,)) + _out = mx.vjp(fn, (x,), (y,)) self.assertEqual(peak_mem, mx.get_peak_memory()) def test_async_eval_with_multiple_streams(self): @@ -128,7 +127,7 @@ def test_async_eval_with_multiple_streams(self): b = mx.array([1.0]) d = mx.default_device() - s2 = mx.new_stream(d) + _s2 = mx.new_stream(d) for _ in range(50): for _ in range(20): @@ -138,7 +137,6 @@ def test_async_eval_with_multiple_streams(self): def test_donation_for_noops(self): def fun(x): - s = x.shape for _ in range(10): x = mx.abs(x) x = mx.reshape(x, (-1,)) @@ -236,6 +234,27 @@ def test_async_eval_error_in_synchronize(self): with self.assertRaises(RuntimeError): mx.synchronize(mx.cpu) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_eval_exception_after_cross_stream_wait(self): + # gather_qqmm has no CPU kernel. It fails in eval after the CPU + # stream waits for x from the GPU. + x = mx.multiply( + mx.full((2, 64), 3.0, stream=mx.gpu), + 2.0, + stream=mx.gpu, + ) + wq, scales = mx.quantize( + mx.ones((32, 64), stream=mx.gpu), + mode="nvfp4", + stream=mx.gpu, + )[:2] + y = mx.gather_qqmm(x, wq, scales, mode="nvfp4", stream=mx.cpu) + with self.assertRaises(RuntimeError): + mx.eval(y) + + self.assertTrue(mx.all(x == 6.0).item()) + self.assertEqual((mx.ones((4,), stream=mx.cpu) + 1).sum().item(), 8.0) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_export_import.py b/Source/Cmlx/mlx/python/tests/test_export_import.py index 23f0645b3..b8de6ce48 100644 --- a/Source/Cmlx/mlx/python/tests/test_export_import.py +++ b/Source/Cmlx/mlx/python/tests/test_export_import.py @@ -12,7 +12,6 @@ class TestExportImport(mlx_tests.MLXTestCase): - @classmethod def setUpClass(cls): cls.test_dir_fid = tempfile.TemporaryDirectory() @@ -160,7 +159,7 @@ def fun(x, z=None): return out x = mx.array([1, 2, 3]) - y = mx.array([1, 1, 0]) + _y = mx.array([1, 1, 0]) z = mx.array([2, 2, 2]) mx.export_function(path, fun, (x,), {"z": z}) @@ -332,6 +331,33 @@ def fun(a, v): out = imported_fun(x, y)[0] self.assertTrue(mx.array_equal(expected, out)) + def test_export_unique(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + # the fixed output size is what makes this exportable. The input is + # fixed so that the padded, exact and truncated cases are all covered + # rather than left to chance. + x = mx.array([3, 1, 2, 1, 3, 2, 1, 0, 2, 1]) # four unique values + for size, fill_value in ((10, None), (6, -1), (4, None), (2, None)): + + def fun(a): + return mx.unique( + a, + size, + return_index=True, + return_inverse=True, + return_counts=True, + fill_value=fill_value, + ) + + mx.export_function(path, fun, x) + imported_fun = mx.import_function(path) + expected = fun(x) + out = imported_fun(x) + self.assertEqual(len(out), len(expected)) + for e, o in zip(expected, out): + self.assertTrue(mx.array_equal(e, o)) + def test_export_conv(self): path = os.path.join(self.test_dir, "fn.mlxfn") diff --git a/Source/Cmlx/mlx/python/tests/test_extension.py b/Source/Cmlx/mlx/python/tests/test_extension.py new file mode 100644 index 000000000..6900bbf29 --- /dev/null +++ b/Source/Cmlx/mlx/python/tests/test_extension.py @@ -0,0 +1,393 @@ +# Copyright © 2026 Apple Inc. + +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +import mlx.core as mx +from mlx.extension import ( + BuildExtension, + MetalExtension, + MetalLibrary, + _extension_paths, + _metal_extension_cmake, + _mlx_macos_deployment_target, +) +from setuptools import Distribution + + +@unittest.skipUnless(mx.metal.is_available(), "MetalExtension requires Metal") +class TestMetalExtension(unittest.TestCase): + def test_extension_paths_use_public_build_metadata(self): + extension = MetalExtension("sample._ext", ["bindings.cpp", "kernel.metal"]) + command = Mock() + command.build_lib = "/tmp/build" + command.get_ext_fullname.return_value = extension.name + command.get_ext_filename.return_value = "sample/_ext.so" + build_py = Mock() + build_py.get_package_dir.return_value = "/tmp/source/sample" + + inplace_file, regular_file = _extension_paths(command, build_py, extension) + + self.assertEqual(inplace_file, Path("/tmp/source/sample/_ext.so")) + self.assertEqual(regular_file, Path("/tmp/build/sample/_ext.so")) + command.get_ext_fullname.assert_called_once_with(extension.name) + command.get_ext_filename.assert_called_once_with(extension.name) + build_py.get_package_dir.assert_called_once_with("sample") + + def _cmake_configure_command( + self, command_type, cmake_args="", debug=False, mlx_target="14.0" + ): + extension = MetalExtension("sample._ext", ["bindings.cpp", "kernel.metal"]) + distribution = Distribution({"ext_modules": [extension]}) + with tempfile.TemporaryDirectory() as directory: + command = command_type(distribution) + command.build_lib = str(Path(directory) / "lib") + command.build_temp = str(Path(directory) / "temp") + command.debug = debug + command.parallel = 1 + with patch( + "mlx.extension._mlx_macos_deployment_target", + return_value=mlx_target, + ): + with patch("mlx.extension.subprocess.run") as run: + with patch.dict(os.environ, {"CMAKE_ARGS": cmake_args}): + command.build_extension(extension) + return run.call_args_list[0].args[0] + + def test_reads_mlx_macos_deployment_target(self): + vtool_output = """ +Load command 10 + cmd LC_BUILD_VERSION + platform MACOS + minos 14.0 +Load command 10 + cmd LC_BUILD_VERSION + platform MACOS + minos 15.2 +""" + with ( + patch("mlx.extension._MLX_PACKAGE_PATH", "/tmp/mlx"), + patch( + "mlx.extension.subprocess.check_output", return_value=vtool_output + ) as check_output, + ): + deployment_target = _mlx_macos_deployment_target() + + self.assertEqual(deployment_target, "15.2") + check_output.assert_called_once_with( + [ + "xcrun", + "vtool", + "-show-build", + "/tmp/mlx/lib/libmlx.dylib", + ], + stderr=subprocess.STDOUT, + text=True, + ) + + def test_build_extension_can_enable_ninja(self): + command_type = BuildExtension.with_options(use_ninja=True) + with patch("mlx.extension._is_ninja_available", return_value=True): + command = command_type(Distribution()) + configure_command = self._cmake_configure_command(command_type) + + self.assertTrue(command.use_ninja) + self.assertTrue(command.generate_stubs) + self.assertIn("-DMLX_METAL_DEBUG=OFF", configure_command) + generator_index = configure_command.index("-G") + self.assertEqual( + configure_command[generator_index : generator_index + 2], ["-G", "Ninja"] + ) + binary_dir = Path(configure_command[configure_command.index("-B") + 1]) + self.assertEqual(binary_dir.name, "build-ninja") + + def test_debug_build_enables_metal_debug(self): + command_type = BuildExtension.with_options(use_ninja=False) + configure_command = self._cmake_configure_command(command_type, debug=True) + + self.assertIn("-DCMAKE_BUILD_TYPE=Debug", configure_command) + self.assertIn("-DMLX_METAL_DEBUG=ON", configure_command) + + def test_build_extension_uses_mlx_macos_deployment_target(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {}, clear=True): + configure_command = self._cmake_configure_command( + command_type, mlx_target="15.2" + ) + + self.assertIn("-DCMAKE_OSX_DEPLOYMENT_TARGET=15.2", configure_command) + + def test_environment_overrides_macos_deployment_target(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {"MACOSX_DEPLOYMENT_TARGET": "15.0"}): + configure_command = self._cmake_configure_command( + command_type, mlx_target="16.0" + ) + + self.assertIn("-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0", configure_command) + + def test_environment_rejects_macos_deployment_target_below_minimum(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {"MACOSX_DEPLOYMENT_TARGET": "13.0"}): + with self.assertRaisesRegex(ValueError, "must be at least 14.0"): + self._cmake_configure_command(command_type) + + def test_cmake_args_overrides_macos_deployment_target(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {}, clear=True): + configure_command = self._cmake_configure_command( + command_type, + cmake_args="-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0", + mlx_target="16.0", + ) + + default_index = configure_command.index("-DCMAKE_OSX_DEPLOYMENT_TARGET=16.0") + override_index = configure_command.index("-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0") + self.assertLess(default_index, override_index) + + def test_cmake_args_rejects_macos_deployment_target_below_minimum(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {}, clear=True): + with self.assertRaisesRegex(ValueError, "must be at least 14.0"): + self._cmake_configure_command( + command_type, + cmake_args="-DCMAKE_OSX_DEPLOYMENT_TARGET=13.0", + ) + + def test_build_extension_uses_default_generator(self): + with patch("mlx.extension._is_ninja_available") as is_ninja_available: + command = BuildExtension(Distribution()) + configure_command = self._cmake_configure_command(BuildExtension) + + self.assertFalse(command.use_ninja) + self.assertNotIn("-G", configure_command) + binary_dir = Path(configure_command[configure_command.index("-B") + 1]) + self.assertEqual(binary_dir.name, "build-default") + is_ninja_available.assert_not_called() + + def test_build_extension_falls_back_without_ninja(self): + command_type = BuildExtension.with_options(use_ninja=True) + with patch("mlx.extension._is_ninja_available", return_value=False): + with self.assertLogs("mlx.extension", level="WARNING") as logs: + command = command_type(Distribution()) + configure_command = self._cmake_configure_command(command_type) + + self.assertFalse(command.use_ninja) + self.assertNotIn("-G", configure_command) + binary_dir = Path(configure_command[configure_command.index("-B") + 1]) + self.assertEqual(binary_dir.name, "build-default") + self.assertIn("Falling back", "\n".join(logs.output)) + + def test_explicit_cmake_generator_overrides_ninja(self): + command_type = BuildExtension.with_options(use_ninja=True) + with patch("mlx.extension._is_ninja_available", return_value=True): + configure_command = self._cmake_configure_command( + command_type, cmake_args="-G Xcode" + ) + + self.assertEqual(configure_command.count("-G"), 1) + generator_index = configure_command.index("-G") + self.assertEqual( + configure_command[generator_index : generator_index + 2], ["-G", "Xcode"] + ) + binary_dir = Path(configure_command[configure_command.index("-B") + 1]) + self.assertEqual(binary_dir.name, "build-xcode") + + def test_cmake_generator_environment_uses_distinct_build_directory(self): + command_type = BuildExtension.with_options(use_ninja=False) + with patch.dict(os.environ, {"CMAKE_GENERATOR": "Unix Makefiles"}): + configure_command = self._cmake_configure_command(command_type) + + self.assertNotIn("-G", configure_command) + binary_dir = Path(configure_command[configure_command.index("-B") + 1]) + self.assertEqual(binary_dir.name, "build-unix_makefiles") + + def test_build_extension_can_disable_stub_generation(self): + command_type = BuildExtension.with_options(generate_stubs=False) + command = command_type(Distribution()) + extension = MetalExtension("sample._ext", ["bindings.cpp", "kernel.metal"]) + + self.assertFalse(command.generate_stubs) + with patch( + "mlx.extension._metal_extension_cmake", wraps=_metal_extension_cmake + ) as generate_cmake: + self._cmake_configure_command(command_type) + self.assertFalse(generate_cmake.call_args.kwargs["generate_stubs"]) + + cmake = _metal_extension_cmake( + extension, Path("/tmp/output"), generate_stubs=False + ) + self.assertNotIn("nanobind.stubgen", cmake) + self.assertNotIn("_ext.pyi", cmake) + self.assertIn("mlx_extension_metallib", cmake) + + def test_extension_metadata(self): + extension = MetalExtension( + "sample._ext", + ["bindings.cpp", "silu.metal", "gelu.metal"], + extra_compile_args={"cxx": ["-O2"], "metal": ["-O3"]}, + ) + + self.assertEqual( + extension.sources, + ["bindings.cpp", "silu.metal", "gelu.metal"], + ) + self.assertEqual(extension.language, "c++") + self.assertFalse(extension.py_limited_api) + self.assertEqual(extension.metal_library_name, "_ext") + self.assertEqual(len(extension.metal_libraries), 1) + self.assertEqual(extension.metal_libraries[0].name, "_ext") + self.assertEqual( + extension.metal_libraries[0].sources, + ["silu.metal", "gelu.metal"], + ) + self.assertEqual( + extension.extra_compile_args, + {"cxx": ["-O2"], "metal": ["-O3"]}, + ) + + def test_requires_host_source(self): + with self.assertRaisesRegex(ValueError, "C\\+\\+ source"): + MetalExtension("sample._ext", ["kernel.metal"]) + + def test_compile_argument_list_applies_to_cxx(self): + extension = MetalExtension( + "sample._ext", + ["bindings.cpp", "silu.metal", "gelu.metal"], + extra_compile_args=["-O2"], + ) + + self.assertEqual( + extension.extra_compile_args, + {"cxx": ["-O2"], "metal": []}, + ) + + def test_requires_metal_source(self): + with self.assertRaisesRegex(ValueError, "Metal source"): + MetalExtension("sample._ext", ["bindings.cpp"]) + + def test_explicit_metal_libraries(self): + extension = MetalExtension( + "sample._ext", + ["bindings.cpp"], + metal_libraries=[ + MetalLibrary("first", ["first.metal"]), + MetalLibrary( + "second", + ["second.metal"], + extra_compile_args=["-O3"], + deployment_target="26.2", + ), + ], + extra_compile_args={"metal": ["-Wno-unused"]}, + ) + + self.assertEqual( + [library.name for library in extension.metal_libraries], + ["first", "second"], + ) + cmake = _metal_extension_cmake(extension, Path("/tmp/output")) + self.assertIn("TITLE [=[first]=]", cmake) + self.assertIn("TITLE [=[second]=]", cmake) + self.assertIn("TARGET mlx_extension_metallib_0", cmake) + self.assertIn("TARGET mlx_extension_metallib_1", cmake) + self.assertIn("DEPLOYMENT_TARGET [=[26.2]=]", cmake) + self.assertIn("[=[-Wno-unused]=]\n [=[-O3]=]", cmake) + self.assertIn( + "add_dependencies(mlx_extension mlx_extension_metallib_0 " + "mlx_extension_metallib_1)", + cmake, + ) + + def test_rejects_mixed_metal_source_forms(self): + with self.assertRaisesRegex(ValueError, "specified in metal_libraries"): + MetalExtension( + "sample._ext", + ["bindings.cpp", "kernel.metal"], + metal_libraries=[MetalLibrary("kernels", ["other.metal"])], + ) + + def test_rejects_duplicate_metal_library_names(self): + with self.assertRaisesRegex(ValueError, "must be unique"): + MetalExtension( + "sample._ext", + ["bindings.cpp"], + metal_libraries=[ + MetalLibrary("kernels", ["first.metal"]), + MetalLibrary("kernels", ["second.metal"]), + ], + ) + + def test_rejects_unknown_compile_argument_group(self): + with self.assertRaisesRegex(ValueError, "cxx.*metal"): + MetalExtension( + "sample._ext", + ["bindings.cpp", "kernel.metal"], + extra_compile_args={"nvcc": ["-O3"]}, + ) + + def test_rejects_non_list_compile_arguments(self): + with self.assertRaisesRegex(TypeError, "must be lists"): + MetalExtension( + "sample._ext", + ["bindings.cpp", "kernel.metal"], + extra_compile_args={"metal": "-O3"}, + ) + + def test_generated_cmake_builds_module_and_metallib(self): + extension = MetalExtension( + "sample._ext", + ["bindings.cpp", "silu.metal", "gelu.metal"], + include_dirs=["include"], + extra_compile_args={"cxx": ["-Wall"], "metal": ["-O3"]}, + ) + + cmake = _metal_extension_cmake(extension, Path("/tmp/output")) + + self.assertIn("nanobind_add_module(", cmake) + self.assertIn("cmake_minimum_required(VERSION 3.25)", cmake) + self.assertIn( + "find_package(Python 3.10 COMPONENTS " + "Interpreter Development.Module REQUIRED)", + cmake, + ) + self.assertNotIn("Development.SABIModule", cmake) + self.assertNotIn(" STABLE_ABI", cmake) + self.assertIn("TITLE [=[_ext]=]", cmake) + self.assertIn("COMPILE_OPTIONS\n [=[-O3]=]", cmake) + self.assertIn(str(Path("silu.metal").resolve()), cmake) + self.assertIn(str(Path("gelu.metal").resolve()), cmake) + self.assertIn("[=[MLX_EXTENSION_NAME=_ext]=]", cmake) + self.assertIn('[=[MLX_METAL_LIBRARY_NAME="_ext"]=]', cmake) + self.assertIn("target_compile_options(mlx_extension PRIVATE", cmake) + self.assertIn("[=[-Wall]=]", cmake) + self.assertIn("BUILD_WITH_INSTALL_RPATH TRUE", cmake) + self.assertIn("INSTALL_RPATH [=[@loader_path;@loader_path/../mlx/lib]=]", cmake) + self.assertNotIn("BUILD_RPATH", cmake) + self.assertIn("from nanobind.stubgen import main", cmake) + self.assertIn("import mlx.core", cmake) + self.assertIn("OUTPUT [=[/tmp/output/_ext.pyi]=]", cmake) + self.assertIn("DEPENDS mlx_extension", cmake) + + def test_generated_cmake_uses_nested_package_mlx_rpath(self): + extension = MetalExtension("sample.ops._ext", ["bindings.cpp", "kernel.metal"]) + + cmake = _metal_extension_cmake(extension, Path("/tmp/output")) + + self.assertIn("@loader_path/../../mlx/lib", cmake) + + def test_rejects_python_stable_abi(self): + with self.assertRaisesRegex(ValueError, "does not support py_limited_api"): + MetalExtension( + "sample._ext", + ["bindings.cpp", "kernel.metal"], + py_limited_api=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/Source/Cmlx/mlx/python/tests/test_fast.py b/Source/Cmlx/mlx/python/tests/test_fast.py index 56f6211ef..dfdcc1824 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast.py +++ b/Source/Cmlx/mlx/python/tests/test_fast.py @@ -166,7 +166,7 @@ def test_rope(self): dims, _, base, scale, offset, traditional = defaults x = (mx.random.uniform(shape=(2, T, dims)) * 10).astype(mx.int32) with self.assertRaises(ValueError): - y = mx.fast.rope( + mx.fast.rope( x, dims, traditional=traditional, base=base, scale=scale, offset=offset ) @@ -284,19 +284,22 @@ def test_rope_with_freqs(self): self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5) # Test grad with freqs - f1 = lambda x, y: (rope_orig(x, dims, False, None, 1.0, 0, freqs) * y).sum() - f2 = lambda x, y: ( - mx.fast.rope( - x, - dims, - traditional=False, - base=None, - scale=1.0, - offset=0, - freqs=freqs, - ) - * y - ).sum() + def f1(x, y): + return (rope_orig(x, dims, False, None, 1.0, 0, freqs) * y).sum() + + def f2(x, y): + return ( + mx.fast.rope( + x, + dims, + traditional=False, + base=None, + scale=1.0, + offset=0, + freqs=freqs, + ) + * y + ).sum() x = mx.random.uniform(shape=(2, 4, dims)) y = mx.random.uniform(shape=(2, 4, dims)) @@ -310,20 +313,24 @@ def test_rope_grad(self): for dims in (D, D // 2): for traditional in (True, False): _, base, scale, offset, _ = defaults - f1 = lambda x, y: ( - rope_orig(x, dims, traditional, base, scale, offset) * y - ).sum() - f2 = lambda x, y: ( - mx.fast.rope( - x, - dims, - traditional=traditional, - base=base, - scale=scale, - offset=offset, - ) - * y - ).sum() + + def f1(x, y): + return ( + rope_orig(x, dims, traditional, base, scale, offset) * y + ).sum() + + def f2(x, y): + return ( + mx.fast.rope( + x, + dims, + traditional=traditional, + base=base, + scale=scale, + offset=offset, + ) + * y + ).sum() x = mx.random.uniform(shape=(2, 100, D)) y = mx.random.uniform(shape=(2, 100, D)) @@ -373,7 +380,7 @@ def test_rope_batch(self): rx_fast = mx.fast.rope( x, dims, traditional=traditional, scale=scale, base=base, offset=offset ) - rx_fast_single = mx.fast.rope( + _rx_fast_single = mx.fast.rope( x[0:1], dims, traditional=traditional, scale=scale, base=base, offset=offset ) @@ -484,10 +491,18 @@ def test_rms_norm(self): def test_rms_norm_grad(self): eps = 1e-5 - f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum() - f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum() - f3 = lambda x, y: (rms_norm(x, mx.ones((x.shape[-1],)), eps) * y).sum() - f4 = lambda x, y: (mx.fast.rms_norm(x, None, eps) * y).sum() + + def f1(x, w, y): + return (rms_norm(x, w, eps) * y).sum() + + def f2(x, w, y): + return (mx.fast.rms_norm(x, w, eps) * y).sum() + + def f3(x, y): + return (rms_norm(x, mx.ones((x.shape[-1],)), eps) * y).sum() + + def f4(x, y): + return (mx.fast.rms_norm(x, None, eps) * y).sum() for D in [32, 256]: x = mx.random.uniform(shape=(8, 100, D)) @@ -568,8 +583,11 @@ def ref(logits, targets): ) return mx.logsumexp(logits, axis=-1) - score - f1 = lambda x, y: ref(x, y).mean() - f2 = lambda x, y: mx.fast.cross_entropy(x, y).mean() + def f1(x, y): + return ref(x, y).mean() + + def f2(x, y): + return mx.fast.cross_entropy(x, y).mean() for V in [7, 128, 1000, 4096]: logits = mx.random.normal(shape=(4, 7, V), scale=2.0) @@ -580,8 +598,13 @@ def ref(logits, targets): self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) w = mx.random.uniform(shape=(4, 7)) - f3 = lambda x, y: (ref(x, y) * w).sum() - f4 = lambda x, y: (mx.fast.cross_entropy(x, y) * w).sum() + + def f3(x, y): + return (ref(x, y) * w).sum() + + def f4(x, y): + return (mx.fast.cross_entropy(x, y) * w).sum() + logits = mx.random.normal(shape=(4, 7, 512), scale=2.0) targets = mx.random.randint(0, 512, shape=(4, 7)) g1 = mx.grad(f3, argnums=0)(logits, targets) @@ -687,7 +710,6 @@ def test_layer_norm(self): self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype]) def test_slice_into_layer_norm(self): - dim = 128 eps = 1e-5 x = mx.random.uniform(shape=(8, 100, 128))[:, 99:] rx_fast = mx.fast.layer_norm(x, weight=None, bias=None, eps=eps) @@ -697,8 +719,12 @@ def test_slice_into_layer_norm(self): def test_layer_norm_grad(self): D = 32 eps = 1e-5 - f1 = lambda x, w, b, y: (layer_norm(x, w, b, eps) * y).sum() - f2 = lambda x, w, b, y: (mx.fast.layer_norm(x, w, b, eps) * y).sum() + + def f1(x, w, b, y): + return (layer_norm(x, w, b, eps) * y).sum() + + def f2(x, w, b, y): + return (mx.fast.layer_norm(x, w, b, eps) * y).sum() x = mx.random.uniform(shape=(8, 100, D)) w = mx.random.uniform(shape=(D,)) @@ -749,8 +775,11 @@ def test_layer_norm_grad_no_bias(self): y = mx.random.uniform(shape=(2, 4, D)) mx.eval(x, w, y) - f_ref = lambda x, w, y: (layer_norm(x, w, None, eps) * y).sum() - f_fast = lambda x, w, y: (mx.fast.layer_norm(x, w, None, eps) * y).sum() + def f_ref(x, w, y): + return (layer_norm(x, w, None, eps) * y).sum() + + def f_fast(x, w, y): + return (mx.fast.layer_norm(x, w, None, eps) * y).sum() # First order should match reference gx1, gw1 = mx.grad(f_ref, argnums=(0, 1))(x, w, y) @@ -775,8 +804,13 @@ def inner(x, w, y): def test_layer_norm_grad_no_params(self): eps = 1e-5 - f1 = lambda x: layer_norm(x, None, None, eps).sum() - f2 = lambda x: mx.fast.layer_norm(x, None, None, eps).sum() + + def f1(x): + return layer_norm(x, None, None, eps).sum() + + def f2(x): + return mx.fast.layer_norm(x, None, None, eps).sum() + x = mx.random.normal(shape=(2, 2, 8)) mx.eval(x) @@ -786,8 +820,12 @@ def test_layer_norm_grad_no_params(self): def test_layer_norm_grad_params(self): eps = 1e-5 - f1 = lambda params, x: (layer_norm(x, params[0], params[1], eps)).sum() - f2 = lambda params, x: (mx.fast.layer_norm(x, params[0], params[1], eps)).sum() + + def f1(params, x): + return (layer_norm(x, params[0], params[1], eps)).sum() + + def f2(params, x): + return (mx.fast.layer_norm(x, params[0], params[1], eps)).sum() w = mx.ones((8,)) b = mx.zeros((8,)) @@ -947,7 +985,8 @@ def test_custom_kernel_strides(self): elif mx.cuda.is_available(): source = """ auto elem = cooperative_groups::this_grid().thread_rank(); - auto loc = elem_to_loc(elem, inp_shape.data(), inp_strides.data(), inp_ndim); + auto loc = elem_to_loc( + elem, inp_shape.data(), inp_strides.data(), inp_ndim); T tmp = inp[loc]; out[elem] = exp(tmp) * WARP_SIZE; """ @@ -1242,6 +1281,28 @@ def call_kernel(a: mx.array): out = call_kernel(a).astype(mx.float32) + call_kernel(b) self.assertTrue(mx.allclose(out, mx.full((32,), 8.0))) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_metal_kernel_negative_template_int(self): + # Negative int template values used to produce a '-' in the generated + # kernel name (via make_template_hash), which fails Metal compilation + # (#4579). + kernel = mx.fast.metal_kernel( + name="negtmpl", + input_names=["inp"], + output_names=["out"], + source="out[thread_position_in_grid.x] = T(V);", + ) + (out,) = kernel( + inputs=[mx.zeros((4,), dtype=mx.int32)], + template=[("T", mx.int32), ("V", -1)], + grid=(4, 1, 1), + threadgroup=(4, 1, 1), + output_shapes=[(4,)], + output_dtypes=[mx.int32], + ) + mx.eval(out) + self.assertEqual(out.tolist(), [-1, -1, -1, -1]) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py b/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py index d7a33c67a..f0a5255ff 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py +++ b/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py @@ -9,7 +9,7 @@ import torch has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -32,9 +32,12 @@ def gated_delta_oracle( k: [B, T, H, K] v: [B, T, H, V] beta: [B, T, H] - g: [B, T, H] <--- Difference: with out kernel: this is expected as a log. - scale: float, optional <--- Difference: This is done by the qwen3_5 model. - initial_state: [B, H, K, V], optional <- Difference: last two dimensions are transposed. + g: [B, T, H] <--- Difference: with out kernel: + this is expected as a log. + scale: float, optional <--- Difference: This is done by the + qwen3_5 model. + initial_state: [B, H, K, V], optional <- Difference: last two dimensions + are transposed. output_final_state: bool Returns: @@ -115,7 +118,8 @@ def runner(dims, stream=mx.gpu, reference=True): out_ref = mx.array(out_on) hf_ref = mx.array(hf_on) else: - # use fallback for tests once fallback is validated by setting a mask instead of using the cpu + # use fallback for tests once fallback is validated by setting a mask instead of + # using the cpu mask = mx.ones((B, T)) out_ref, hf_ref = mx.fast.gated_delta_update( @@ -155,6 +159,7 @@ class TestGatedDelta(mlx_tests.MLXTestCase): diff_heads2, large_t_dims, ] + vjp_dims = [base_dims, unaligned_dims, diff_heads, diff_heads2] @unittest.skipIf(not has_torch, "requires Torch") def test_gated_delta_fallback(self): @@ -170,7 +175,6 @@ def test_gated_delta_fallback(self): def test_gated_delta_fallback_masked(self): for dims in self.fallback_dims: - B, Hk, Hv, T, Dk, Dv = dims q = mx.random.normal(shape=(B, T, Hk, Dk)) @@ -225,7 +229,6 @@ def test_gated_delta_dtypes(self): dtypes = [mx.bfloat16, mx.float32] for dtype in dtypes: for dims in [self.base_dims]: - B, Hk, Hv, T, Dk, Dv = dims q = mx.random.normal(shape=(B, T, Hk, Dk), dtype=dtype) @@ -281,6 +284,58 @@ def test_gated_delta_nax(self): mx.allclose(hf_ref, hf, atol=1e-1, rtol=1e-4), msg="State " + msg ) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_gated_delta_vjp(self): + for dims in self.vjp_dims: + B, Hk, Hv, T, Dk, Dv = dims + + mx.random.seed(0) + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.random.normal((B, Hv, Dv, Dk), dtype=mx.float32) + primals = [q, k, v, g, b, h0] + + cotans = [ + mx.random.normal(shape=(B, T, Hv, Dv)) * 1e-2, + mx.random.normal(shape=(B, Hv, Dv, Dk)) * 1e-2, + ] + + def f(q, k, v, g, b, h0): + out, state = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0) + return out, state + + def run(fallback, chunk): + with mlx_tests.scoped_env( + GATED_DELTA_VJP_FALLBACK="1" if fallback else "0", + GATED_DELTA_CHUNK_VJP=str(chunk), + ): + outs, vjps = mx.vjp(f, primals, cotans) + mx.eval(outs, vjps) + return outs, vjps + + o_ref, vjp_ref = run(True, 0) + + for chunk in (0, 16): + rtol = 1e-5 if chunk == 0 else 1e-2 + atol = 1e-5 if chunk == 0 else 1e-2 + with self.subTest(dims=dims, chunk=chunk): + o_out, vjp_out = run(False, chunk) + for i in range(len(o_ref)): + self.assertTrue( + mx.allclose(o_ref[i], o_out[i], rtol=rtol, atol=atol), + msg=f"Output {i}, dims={dims}, chunk={chunk}", + ) + for i in range(len(vjp_ref)): + self.assertTrue( + mx.allclose(vjp_ref[i], vjp_out[i], rtol=rtol, atol=atol), + msg=f"Grad {i}, dims={dims}, chunk={chunk}", + ) + mx.clear_cache() + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True) diff --git a/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py b/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py index 142d34195..b2d7e9121 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py +++ b/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py @@ -5,7 +5,6 @@ import mlx.core as mx import mlx_tests -import numpy as np def mlx_ref_attn(q, k, v, scale=1.0, mask=None, sinks=None): @@ -27,7 +26,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None, sinks=None): scores = q @ mx.swapaxes(k, -1, -2) is_causal = mask == "causal" if mask is not None: - if is_causal: offset = kL - L q_indices = mx.arange(L) + offset @@ -361,6 +359,51 @@ def test_sdpa_full_head_dim_256(self): tol = 5e-3 self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + if not mx.metal.is_available() or mx.default_device() != mx.gpu: + return + + mx.random.seed(0) + cases = ( + (1, 4, 2, 1024, 1057, 1152), + (2, 2, 1, 1025, 1025, None), + (1, 2, 1, 1024, 4103, 4128), + ) + for dtype, (B, qH, kH, qL, kL, cache_len) in product( + (mx.float16, mx.bfloat16), cases + ): + with self.subTest(dtype=dtype, qL=qL, kL=kL): + q = (0.5 * mx.random.normal((B, qH, qL, 512))).astype(dtype) + kv_len = cache_len or kL + k_cache = (0.5 * mx.random.normal((B, kH, kv_len, 512))).astype(dtype) + v_cache = (0.5 * mx.random.normal((B, kH, kv_len, 512))).astype(dtype) + k = k_cache[..., :kL, :] + v = v_cache[..., :kL, :] + scale = 512**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask="causal") + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask="causal" + ) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + q = mx.random.normal((1, 2, 1024, 512), mx.float16) + k = mx.random.normal((1, 1, 1057, 512), mx.float16) + v = mx.random.normal((1, 1, 1057, 512), mx.float16) + scale = 512**-0.5 + mask = mx.random.uniform(shape=(1, 1, 1024, 1057)) > 0.2 + sinks = mx.random.normal((2,), mx.float16) + fallback_cases = ( + (v, mask, None), + (v, "causal", sinks), + (v[..., :128], "causal", None), + ) + for v_arg, mask_arg, sinks_arg in fallback_cases: + ref = mlx_ref_attn(q, k, v_arg, scale=scale, mask=mask_arg, sinks=sinks_arg) + out = mx.fast.scaled_dot_product_attention( + q, k, v_arg, scale=scale, mask=mask_arg, sinks=sinks_arg + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") def test_sdpa_full_head_dim_512_nax(self): if mx.default_device() != mx.gpu: @@ -509,6 +552,7 @@ def test_sdpa_vector(self): Nq = 4 Nkv = 1 scale = 1.0 + tol = 1e-3 if mx.cuda.is_available() else 1e-4 mx.random.seed(0) q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D)) k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D)) @@ -544,7 +588,7 @@ def test_sdpa_vector(self): scale=scale, mask=m, ) - self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) L = 4096 scale = 1.0 @@ -573,7 +617,7 @@ def test_sdpa_vector(self): scale=scale, mask=m, ) - self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) def test_sdpa_vector_gqa_long(self): scale = 1.0 @@ -848,6 +892,40 @@ 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.metal.is_available(), "Metal kernel path only") + def test_sdpa_vector_2pass_threadgroup_limit(self): + for dtype, D, q_heads, kv_heads, q_len in ( + (mx.float16, 192, 32, 8, 8), + (mx.float16, 256, 24, 4, 5), + (mx.float16, 256, 32, 4, 4), + (mx.bfloat16, 256, 24, 4, 5), + (mx.bfloat16, 256, 32, 4, 4), + (mx.float32, 192, 32, 8, 8), + (mx.float32, 256, 24, 4, 5), + (mx.float32, 256, 32, 4, 4), + ): + with self.subTest(dtype=dtype, D=D, q_len=q_len): + self.assertGreater(32 * q_heads // kv_heads * q_len, 896) + q = mx.random.normal(shape=(1, q_heads, q_len, D), dtype=dtype) + k = mx.random.normal(shape=(1, kv_heads, 1029, D), dtype=dtype) + v = mx.random.normal(shape=(1, kv_heads, 1029, D), dtype=dtype) + scale = D**-0.5 + + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask="causal", force_fused=True + ) + ref = mlx_ref_attn( + q.astype(mx.float32), + k.astype(mx.float32), + v.astype(mx.float32), + scale=scale, + mask="causal", + ) + atol = 1e-4 if dtype == mx.float32 else 2e-3 + if dtype == mx.bfloat16: + atol = 2e-2 + self.assertTrue(mx.allclose(out.astype(mx.float32), ref, atol=atol)) + @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 @@ -1013,6 +1091,20 @@ def test_sdpa_noncontiguous_inputs(self): ref = mlx_ref_attn(q, k, v, scale=1.0, mask=mask) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_vector_reversed_mask(self): + mx.random.seed(0) + q = mx.random.normal(shape=(1, 2, 1, 64)) + k = mx.random.normal(shape=(1, 2, 16, 64)) + v = mx.random.normal(shape=(1, 2, 16, 64)) + masks = [ + mx.random.normal(shape=(1, 2, 1, 16))[..., ::-1], + (mx.arange(16) % 3 != 0).reshape(1, 1, 1, 16)[..., ::-1], + ] + for mask in masks: + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=0.125, mask=mask) + ref = mlx_ref_attn(q, k, v, scale=0.125, mask=mask) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_promote_mask(self): mask = mx.array(2.0, mx.bfloat16) D = 64 @@ -1138,20 +1230,25 @@ def test_grad(slow, fast, args): mask_bool = mx.random.uniform(0, 1, (B, N_q, T, T), dtype=mx.float16) < 0.5 for mask in (None, "causal", mask_additive, mask_bool): - sdpa_slow = lambda q, k, v: mlx_ref_attn( - q, k, v, scale=scale, mask=mask - ) - sdpa_fast = lambda q, k, v: mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=mask - ) + + def sdpa_slow(q, k, v): + return mlx_ref_attn(q, k, v, scale=scale, mask=mask) + + def sdpa_fast(q, k, v): + return mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + test_vjp(sdpa_slow, sdpa_fast, [q, k, v]) - loss_slow = lambda q, k, v: mlx_ref_attn( - q, k, v, scale=scale, mask=mask - ).sum() - loss_fast = lambda q, k, v: mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=mask - ).sum() + def loss_slow(q, k, v): + return mlx_ref_attn(q, k, v, scale=scale, mask=mask).sum() + + def loss_fast(q, k, v): + return mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ).sum() + test_grad(loss_slow, loss_fast, [q, k, v]) @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") diff --git a/Source/Cmlx/mlx/python/tests/test_fft.py b/Source/Cmlx/mlx/python/tests/test_fft.py index 60f7610fe..0e410c7c5 100644 --- a/Source/Cmlx/mlx/python/tests/test_fft.py +++ b/Source/Cmlx/mlx/python/tests/test_fft.py @@ -11,7 +11,7 @@ import torch has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -59,6 +59,15 @@ def test_fft(self): x = np.fft.rfft(np.real(a_np)) self.check_mx_np(mx.fft.irfft, np.fft.irfft, x) + # Real transforms of size one + a_np = np.arange(1, 5, dtype=np.float32).reshape(4, 1) + self.check_mx_np(mx.fft.rfft, np.fft.rfft, a_np) + x = np.fft.rfft(a_np) + self.check_mx_np(mx.fft.irfft, np.fft.irfft, x, n=1) + self.check_mx_np(mx.fft.rfft2, np.fft.rfft2, a_np) + x = np.fft.rfft2(a_np) + self.check_mx_np(mx.fft.irfft2, np.fft.irfft2, x, s=(4, 1)) + def test_fftn(self): r = np.random.randn(8, 8, 8).astype(np.float32) i = np.random.randn(8, 8, 8).astype(np.float32) @@ -464,23 +473,165 @@ def make_ffts(self): def test_fft_vmap(self): for fftn, a_np, axes in self.make_ffts(): a = mx.array(a_np) - f = lambda x: fftn(x, axes=axes) + + def f(x): + return fftn(x, axes=axes) + expected = mx.stack([f(a[i]) for i in range(a.shape[0])]) out = mx.vmap(f)(a) self.assertEqual(tuple(out.shape), tuple(expected.shape)) np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + def test_fft_vmap_explicit_size(self): + # An inverse real transform to an odd length is not recoverable from + # the input shape. + shape = (3, 8, 6) + a = mx.array(np.random.rand(*shape) + 1j * np.random.rand(*shape)) + for n in [10, 11, 12, 13]: + for fftn, axes in [ + (mx.fft.irfftn, (-1,)), + (mx.fft.irfftn, (-2, -1)), + (mx.fft.ifftn, (-1,)), + ]: + s_arg = [shape[-2], n] if len(axes) == 2 else [n] + + def f(x): + return fftn(x, s=s_arg, axes=axes) + + expected = mx.stack([f(a[i]) for i in range(a.shape[0])]) + out = mx.vmap(f)(a) + self.assertEqual(tuple(out.shape), tuple(expected.shape)) + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + def test_fft_jvp(self): # The fft is linear so the jvp is the fft of the tangent for fftn, a_np, axes in self.make_ffts(): a = mx.array(a_np) t = mx.array(np.random.rand(*a_np.shape).astype(a_np.dtype)) - f = lambda x: fftn(x, axes=axes) + + def f(x): + return fftn(x, axes=axes) + expected = f(t) out = mx.jvp(f, [a], [t])[1][0] self.assertEqual(tuple(out.shape), tuple(expected.shape)) np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + @unittest.skipIf(not has_torch, "requires PyTorch") + def test_stft(self): + def check(x_np, win_np, n_fft, hop, win_len, center, pad_mode, norm, one): + win_mx = None if win_np is None else mx.array(win_np) + win_pt = None if win_np is None else torch.tensor(win_np) + z_mx = mx.fft.stft( + mx.array(x_np), + n_fft=n_fft, + hop_length=hop, + win_length=win_len, + window=win_mx, + center=center, + pad_mode=pad_mode, + norm=norm, + onesided=one, + ) + z_pt = torch.stft( + torch.tensor(x_np), + n_fft=n_fft, + hop_length=hop, + win_length=win_len, + window=win_pt, + center=center, + pad_mode="replicate" if pad_mode == "edge" else pad_mode, + normalized=(norm == "ortho"), + onesided=one, + return_complex=True, + ) + np.testing.assert_allclose( + np.array(z_mx), z_pt.numpy(), atol=1e-4, rtol=1e-4 + ) + + np.random.seed(0) + x = np.random.randn(16000).astype(np.float32) + han = np.hanning(400).astype(np.float32) + h256 = np.hanning(256).astype(np.float32) + check(x, han, 400, 160, 400, True, "reflect", "backward", True) + check(x, han, 512, 128, 400, False, "constant", "backward", True) + check(x, han, 400, 160, 400, True, "edge", "backward", True) + check(x, han, 400, 160, 400, True, "reflect", "ortho", True) + check(x, h256, 256, 64, 256, True, "reflect", "backward", False) + check(x, None, 256, 64, 256, True, "reflect", "backward", True) + xb = np.random.randn(2, 16000).astype(np.float32) + check(xb, h256, 256, 64, 256, True, "reflect", "backward", True) + + @unittest.skipIf(not has_torch, "requires PyTorch") + def test_istft(self): + def check(x_np, win_np, n_fft, hop, center, norm, one, length): + win_mx = None if win_np is None else mx.array(win_np) + win_pt = None if win_np is None else torch.tensor(win_np) + z = mx.fft.stft( + mx.array(x_np), + n_fft=n_fft, + hop_length=hop, + window=win_mx, + center=center, + norm=norm, + onesided=one, + ) + y_mx = mx.fft.istft( + z, + n_fft=n_fft, + hop_length=hop, + window=win_mx, + center=center, + norm=norm, + onesided=one, + length=length, + ) + y_pt = torch.istft( + torch.tensor(np.array(z)), + n_fft=n_fft, + hop_length=hop, + window=win_pt, + center=center, + normalized=(norm == "ortho"), + onesided=one, + length=length, + ) + np.testing.assert_allclose( + np.array(y_mx), y_pt.numpy(), atol=1e-4, rtol=1e-4 + ) + + np.random.seed(0) + x = np.random.randn(16000).astype(np.float32) + win = np.hanning(512).astype(np.float32) + check(x, win, 512, 128, True, "backward", True, 16000) + # n_fft not divisible by hop, no explicit length (length must still match) + check(x, win, 512, 160, True, "backward", True, None) + # center=False needs an edge-nonzero window for torch's NOLA check + check(x, None, 512, 160, False, "backward", True, None) + check(x, win, 512, 128, True, "ortho", True, 16000) + check(x, None, 512, 256, True, "backward", False, 16000) + # odd n_fft (passed explicitly, since onesided inference is even) + check( + x, + np.hanning(257).astype(np.float32), + 257, + 64, + True, + "backward", + True, + 16000, + ) + xb = np.random.randn(2, 16000).astype(np.float32) + check(xb, win, 512, 128, True, "backward", True, 16000) + + # norm="forward" has no torch.istft flag; check the round-trip directly + win_mx = mx.array(win) + z = mx.fft.stft( + mx.array(x), n_fft=512, hop_length=128, window=win_mx, norm="forward" + ) + y = mx.fft.istft(z, hop_length=128, window=win_mx, norm="forward", length=16000) + np.testing.assert_allclose(np.array(y)[1000:-1000], x[1000:-1000], atol=1e-4) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_graph.py b/Source/Cmlx/mlx/python/tests/test_graph.py index 7c6a11412..147232b66 100644 --- a/Source/Cmlx/mlx/python/tests/test_graph.py +++ b/Source/Cmlx/mlx/python/tests/test_graph.py @@ -1,7 +1,6 @@ # Copyright © 2023 Apple Inc. import io -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_init.py b/Source/Cmlx/mlx/python/tests/test_init.py index bd8d9fa4f..a5ca08f74 100644 --- a/Source/Cmlx/mlx/python/tests/test_init.py +++ b/Source/Cmlx/mlx/python/tests/test_init.py @@ -1,6 +1,5 @@ # Copyright © 2023 Apple Inc. import math -import unittest import mlx.core as mx import mlx.nn.init as init @@ -129,10 +128,11 @@ def test_orthogonal(self): self.assertEqual(result.shape, shape) self.assertEqual(result.dtype, mx.float32) - I = result @ result.T + identity = result @ result.T eye = mx.eye(shape[0], dtype=mx.float32) self.assertTrue( - mx.allclose(I, eye, atol=1e-5), "Orthogonal init failed on a square matrix." + mx.allclose(identity, eye, atol=1e-5), + "Orthogonal init failed on a square matrix.", ) # Test with a rectangular matrix: more rows than cols @@ -141,10 +141,10 @@ def test_orthogonal(self): self.assertEqual(result.shape, shape) self.assertEqual(result.dtype, mx.float32) - I = result.T @ result + identity = result.T @ result eye = mx.eye(shape[1], dtype=mx.float32) self.assertTrue( - mx.allclose(I, eye, atol=1e-5), + mx.allclose(identity, eye, atol=1e-5), "Orthogonal init failed on a rectangular matrix.", ) diff --git a/Source/Cmlx/mlx/python/tests/test_linalg.py b/Source/Cmlx/mlx/python/tests/test_linalg.py index 2da3781b4..12065835d 100644 --- a/Source/Cmlx/mlx/python/tests/test_linalg.py +++ b/Source/Cmlx/mlx/python/tests/test_linalg.py @@ -2,7 +2,6 @@ import itertools import math -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_load.py b/Source/Cmlx/mlx/python/tests/test_load.py index d9ad9751c..d3c17fde7 100644 --- a/Source/Cmlx/mlx/python/tests/test_load.py +++ b/Source/Cmlx/mlx/python/tests/test_load.py @@ -65,7 +65,7 @@ def test_save_and_load(self): load_arr_mlx_npy = np.load(save_file_mlx) self.assertTrue(np.array_equal(load_arr_mlx_npy, save_arr_npy)) - save_file = os.path.join(self.test_dir, f"mlx_path.npy") + save_file = os.path.join(self.test_dir, "mlx_path.npy") save_arr = mx.ones((32,)) mx.save(Path(save_file), save_arr) @@ -99,6 +99,26 @@ def test_load_npy_read_error(self): with self.assertRaises(RuntimeError): mx.eval(out) + def test_load_npy_invalid_shape(self): + def write_npy(path, shape): + header = "{'descr': ' 0).item(), 3 + ) + expect(mx.unique(a, 2, return_counts=True)[1], [2, 2]) + + # the inverse has the shape of the input and indexes into the output, + # so values[inverse] rebuilds the input when nothing was truncated + values, inverse = mx.unique(b, 4, return_inverse=True) + expect(values, [1, 2, 3, 4]) + self.assertEqual(inverse.shape, b.shape) + expect(values[inverse], [[3, 1, 3], [2, 1, 4]]) + + # truncation clamps the inverse to keep it in range, so a dropped value + # comes back as the largest kept one instead of reading past the output + values, inverse = mx.unique(a, 2, return_inverse=True) + self.assertTrue(bool(mx.all(inverse < values.size))) + expect(values[inverse], [2, 1, 2, 2, 1]) + + # a size of zero leaves no index to clamp to, the only case where the + # reconstruction is not a valid array + for x in (a, mx.zeros((2, 3), mx.int32), mx.array(5)): + values, inverse = mx.unique(x, 0, return_inverse=True) + self.assertEqual(values.size, 0) + self.assertTrue(mx.array_equal(inverse, mx.zeros(x.shape, mx.uint32))) + with self.assertRaises(ValueError): + mx.eval(values[inverse]) + + # a 0-d input is flattened like any other + values, inverse = mx.unique(mx.array(5), 1, return_inverse=True) + expect(values, [5]) + self.assertEqual(inverse.shape, ()) + + # an empty input needs a fill value if size > 0 + e = mx.array([], mx.float32) + values, inverse, counts = mx.unique( + e, 0, return_inverse=True, return_counts=True + ) + self.assertEqual(values.dtype, mx.float32) + self.assertEqual((values.shape, inverse.shape, counts.shape), ((0,),) * 3) + with self.assertRaises(ValueError): + mx.unique(e, 3) + # the reconstruction stays valid even with size 0 + for size, want in ((0, []), (3, [7.0, 7.0, 7.0])): + values, inverse = mx.unique(e, size, return_inverse=True, fill_value=7) + expect(values, want) + rebuilt = values[inverse] + mx.eval(rebuilt) + self.assertEqual(rebuilt.shape, (0,)) + + # every element identical, and every element distinct + values, counts = mx.unique(mx.full((7,), 4, mx.int32), 7, return_counts=True) + expect(values, [4] * 7) + expect(counts, [7] + [0] * 6) + expect(mx.unique(mx.arange(6), 6), [0, 1, 2, 3, 4, 5]) + + # check 64 bit types work + for dtype in (mx.int64, mx.uint64, mx.complex64): + with self.subTest(dtype=dtype): + values, inverse, counts = mx.unique( + mx.array([3, 1, 2, 1], dtype), + 4, + return_inverse=True, + return_counts=True, + ) + self.assertEqual(values.dtype, dtype) + self.assertTrue(mx.array_equal(values, mx.array([1, 2, 3, 1], dtype))) + expect(inverse, [2, 0, 1, 0]) + expect(counts, [2, 1, 1, 0]) + # in particular the int64 numpy hands over by default + expect(mx.unique(mx.array(np.arange(5)), 5), [0, 1, 2, 3, 4]) + + # bool only works on the cpu stream, the metal sort has no bool kernel + values, inverse, counts = mx.unique( + mx.array([True, False, True]), + 2, + return_inverse=True, + return_counts=True, + stream=mx.cpu, + ) + expect(values, [False, True]) + expect(inverse, [1, 0, 1]) + expect(counts, [1, 2]) + + # NaN sorts last and never equals itself, so each one is kept, and none + # is ever picked as the default fill even though mx.min would return it + nan = float("nan") + values = mx.unique(mx.array([1.0, nan, 2.0, nan], mx.float32), 4) + expect(values[:2], [1.0, 2.0]) + self.assertTrue(bool(mx.all(mx.isnan(values[2:])))) + values = mx.unique(mx.array([3.0, 1.0, nan, 2.0], mx.float32), 6) + expect(values[:3], [1.0, 2.0, 3.0]) + self.assertTrue(bool(mx.isnan(values[3]))) + expect(values[4:], [1.0, 1.0]) + # with nothing but NaN there is no other element to fall back on, so + # the default fill is NaN as well + values = mx.unique(mx.array([nan, nan], mx.float32), 5) + self.assertTrue(bool(mx.all(mx.isnan(values)))) + + # the gradient credits the first occurrence of each unique value, and + # padding routes to that same element + for size, data, want in ( + (3, [3.0, 1.0, 2.0, 1.0], [1.0, 1.0, 1.0, 0.0]), + (4, [3.0, 1.0, 2.0, 1.0], [1.0, 2.0, 1.0, 0.0]), + (2, [1.0, 1.0, 2.0], [1.0, 0.0, 1.0]), + ): + with self.subTest(size=size, data=data): + grad = mx.grad(lambda x: mx.sum(mx.unique(x, size)))(mx.array(data)) + expect(grad, want) + + # the output shape does not depend on the values, so this composes with + # the function transforms + self.assertTrue( + mx.array_equal(mx.compile(lambda x: mx.unique(x, 3))(a), mx.unique(a, 3)) + ) + + # every optional output and a fill value have to survive compilation too + def all_outputs(x): + return mx.unique( + x, + 5, + return_index=True, + return_inverse=True, + return_counts=True, + fill_value=-1, + ) + + for got, want in zip(mx.compile(all_outputs)(a), all_outputs(a)): + self.assertTrue(mx.array_equal(got, want)) + self.assertEqual(got.dtype, want.dtype) + # the batch rows must have different unique sets, otherwise broadcasting + # one row over the batch would pass + rows = mx.array([[2, 1, 2, 3, 1], [7, 7, 5, 5, 9]]) + values, inverse, counts = mx.vmap( + lambda x: mx.unique(x, 3, return_inverse=True, return_counts=True) + )(rows) + expect(values, [[1, 2, 3], [5, 7, 9]]) + expect(inverse, [[1, 0, 1, 2, 0], [1, 1, 0, 0, 2]]) + expect(counts, [[2, 2, 1], [2, 2, 1]]) + + # compare against numpy on random inputs + rng = np.random.RandomState(0) + for n in (1, 2, 17, 1000, 5000): + with self.subTest(n=n): + a_np = rng.randint(-20, 20, size=n).astype(np.int32) + a_mx = mx.array(a_np) + v_np, i_np, c_np = np.unique( + a_np, return_inverse=True, return_counts=True + ) + values, inverse, counts = mx.unique( + a_mx, len(v_np), return_inverse=True, return_counts=True + ) + self.assertTrue(np.array_equal(np.array(values), v_np)) + self.assertTrue( + np.array_equal(np.array(inverse).reshape(-1), i_np.reshape(-1)) + ) + self.assertTrue(np.array_equal(np.array(counts), c_np)) + self.assertTrue(mx.array_equal(values[inverse], a_mx)) + # a padded output holds the same values up front + padded, pad_counts = mx.unique(a_mx, n, return_counts=True) + self.assertEqual(padded.shape, (n,)) + self.assertTrue(np.array_equal(np.array(padded)[: len(v_np)], v_np)) + self.assertEqual(mx.sum(pad_counts > 0).item(), len(v_np)) + + # non contiguous inputs are flattened over their logical elements, + # broadcast multiplicity included, not over their backing memory + base_np = rng.randint(0, 5, size=(4, 6)).astype(np.int32) + base = mx.array(base_np) + for v_mx, v_np in ( + (base.T, base_np.T), + (base[::2], base_np[::2]), + (base[::-1], base_np[::-1]), + (base[:, ::-1], base_np[:, ::-1]), + (mx.broadcast_to(base[0], (3, 6)), np.broadcast_to(base_np[0], (3, 6))), + ): + want = np.unique(v_np) + values, inverse = mx.unique(v_mx, len(want), return_inverse=True) + self.assertTrue(np.array_equal(np.array(values), want)) + # the inverse follows the shape of the view, not of its backing array + self.assertEqual(inverse.shape, v_mx.shape) + self.assertTrue(np.array_equal(np.array(values[inverse]), v_np)) + def test_large_binary(self): a = mx.ones([1000, 2147484], mx.int8) b = mx.ones([2147484], mx.int8) @@ -3264,6 +3602,14 @@ def test_flatten(self): self.assertEqual(x.flatten(start_axis=1).shape, (2, 3 * 4)) self.assertEqual(x.flatten(end_axis=1).shape, (2 * 3, 4)) + def test_reshape_col_contiguous(self): + # Reshape without a copy + a = mx.arange(12).reshape(3, 4).T.reshape(2, 2, 3) + b = mx.arange(12).reshape(3, 2, 2).T + a_np = np.arange(12).reshape(3, 4).T.reshape(2, 2, 3) + b_np = np.arange(12).reshape(3, 2, 2).T + self.assertTrue(np.array_equal(a + b, a_np + b_np)) + def test_clip(self): a = np.array([1, 4, 3, 8, 5], np.int32) expected = np.clip(a, 2, 6) @@ -3821,6 +4167,11 @@ def test_bitwise_ops(self): self.assertEqual(c.shape, (3, 2, 5)) self.assertTrue(mx.array_equal(c, mx.ones((3, 2, 5), dtype=mx.bool_))) + x = mx.array([1, -1, 1, -1]) + y = mx.array([1, 1, -1, -1]) + out = (x > 0) ^ (y > 0) + self.assertEqual(out.tolist(), [False, True, True, False]) + def test_bitwise_grad(self): a = np.random.randint(0, 10, size=(4, 3)) b = np.random.randint(0, 10, size=(4, 3)) @@ -3887,6 +4238,26 @@ def test_view(self): a_out = out.view(mx.int32) self.assertTrue(mx.array_equal(a_out, a, equal_nan=True)) + # Size one last axis + a = mx.arange(4, dtype=mx.int32).reshape(1, 4).T + expected = np.arange(4, dtype=np.int32).reshape(4, 1).view(np.int16) + self.assertTrue(np.array_equal(a.view(mx.int16), expected)) + + # Byte offsets that are not a multiple of the new item size + a = mx.arange(64, dtype=mx.uint8) + a_np = np.arange(64, dtype=np.uint8) + for start in range(8): + out = a[start : start + 16].view(mx.int32) + expected = a_np[start : start + 16].view(np.int32) + self.assertTrue(np.array_equal(out + 0, expected)) # read in a kernel + + # View of an (n, 1) array + a = mx.arange(4, dtype=mx.int32).reshape(4, 1).view(mx.int16) + b = mx.arange(8, dtype=mx.int16).reshape(2, 4).T + expected = np.arange(4, dtype=np.int32).reshape(4, 1).view(np.int16) + expected = expected + np.arange(8, dtype=np.int16).reshape(2, 4).T + self.assertTrue(np.array_equal(a + b, expected)) + def _hadamard(self, N): # Matches scipy.linalg.hadamard H = np.array([[1]], dtype=np.int64) @@ -4117,6 +4488,23 @@ def test_dynamic_slicing(self): expected[1:, 2:, 3:] = update self.assertTrue(mx.array_equal(expected, out)) + # Strided, broadcast and reversed start indices + x = mx.arange(64).reshape(8, 8) + update = mx.zeros((2, 2), dtype=x.dtype) + starts = [ + (mx.array([1, 5, 2, 5])[::2], (1, 2)), + (mx.broadcast_to(mx.array([2]), (2,)), (2, 2)), + (mx.array([4, 1])[::-1], (1, 4)), + ] + for start, (i, j) in starts: + out = mx.slice(x, start, (0, 1), (2, 2)) + self.assertTrue(mx.array_equal(out, x[i : i + 2, j : j + 2])) + + out = mx.slice_update(x, update, start, (0, 1)) + expected = mx.arange(64).reshape(8, 8) + expected[i : i + 2, j : j + 2] = update + self.assertTrue(mx.array_equal(out, expected)) + def test_broadcast_arrays(self): a = mx.array(1) b = mx.array(1.0) @@ -4305,6 +4693,12 @@ def test_masked_scatter(self): ) ) + # source with a negative stride + a = mx.zeros((4,)) + mask = mx.array([True, False, True, True]) + a[mask] = mx.array([1.0, 2.0, 3.0, 4.0])[::-1] + self.assertTrue(mx.array_equal(a, mx.array([4.0, 0.0, 3.0, 2.0]))) + def test_broadcast_shapes(self): # Basic broadcasting self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3)) @@ -4343,6 +4737,28 @@ def test_broadcast_shapes(self): with self.assertRaises(ValueError): mx.broadcast_shapes() + def test_sort_transposed(self): + # The sorted axis can keep the smallest or largest stride while the + # axes that are not sorted are no longer a row major block, which is + # what the contiguous kernel's row enumeration assumes. + np.random.seed(0) + for shape in [(3, 4, 8), (2, 1, 6), (2, 3, 4, 2), (2, 1, 3, 4)]: + a_np = np.random.uniform(0, 100, size=shape).astype(np.float32) + a_mx = mx.array(a_np) + for perm in permutations(range(len(shape))): + b_np = np.transpose(a_np, perm) + b_mx = mx.transpose(a_mx, perm) + for axis in range(len(shape)): + with self.subTest(shape=shape, perm=perm, axis=axis): + s_np = np.sort(b_np, axis=axis) + self.assertTrue(np.array_equal(s_np, mx.sort(b_mx, axis=axis))) + idx = np.array(mx.argsort(b_mx, axis=axis)) + self.assertTrue( + np.array_equal( + s_np, np.take_along_axis(b_np, idx, axis=axis) + ) + ) + def test_sort_nan(self): for dtype in [mx.float32, mx.float16, mx.bfloat16]: with self.subTest(dtype=dtype): diff --git a/Source/Cmlx/mlx/python/tests/test_optimizers.py b/Source/Cmlx/mlx/python/tests/test_optimizers.py index c7caf5b95..94a1809b0 100644 --- a/Source/Cmlx/mlx/python/tests/test_optimizers.py +++ b/Source/Cmlx/mlx/python/tests/test_optimizers.py @@ -15,10 +15,10 @@ try: import torch - import torch.nn.functional as F + import torch.nn.functional as F # noqa: F401 has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -132,7 +132,7 @@ def test_adagrad(self): "first": [mx.zeros((10,)), mx.zeros((1,))], "second": mx.zeros((1,)), } - grads = tree_map(lambda x: mx.ones_like(x), params) + _grads = tree_map(lambda x: mx.ones_like(x), params) # Explicit init optim = opt.Adagrad(learning_rate=1e-2) @@ -150,7 +150,7 @@ def test_adadelta(self): "first": [mx.zeros((10,)), mx.zeros((1,))], "second": mx.zeros((1,)), } - grads = tree_map(lambda x: mx.ones_like(x), params) + _grads = tree_map(lambda x: mx.ones_like(x), params) # Explicit init optim = opt.AdaDelta(learning_rate=1e-2) @@ -263,7 +263,7 @@ def test_lion(self): "first": [mx.zeros((10,)), mx.zeros((1,))], "second": mx.zeros((1,)), } - grads = tree_map(lambda x: mx.ones_like(x), params) + _grads = tree_map(lambda x: mx.ones_like(x), params) # Explicit init optim = opt.Lion(learning_rate=1e-2) @@ -406,7 +406,7 @@ def step(x): uncompiled_params = model.parameters() # Pure version - def loss(params, x): + def loss(params, x): # noqa: F811 model.update(params) return model(x).sum() @@ -426,7 +426,7 @@ def step(params, opt_state, x): self.assertTrue(mx.allclose(pure_params["bias"], uncompiled_params["bias"])) # Impure version - def loss(model, x): + def loss(model, x): # noqa: F811 return model(x).sum() model.update(orig_params) @@ -588,6 +588,57 @@ def test_clip_grad_norm(self): with self.assertRaises(ValueError): opt.clip_grad_norm(small_grads, -1.0) + def test_clip_grad_norm_float16(self): + grads = { + "first": [mx.full((100,), 30.0, dtype=mx.float16)], + "second": mx.array([400.0], dtype=mx.float16), + } + for clip in (opt.clip_grad_norm, mx.compile(opt.clip_grad_norm)): + for max_norm, first, second in ( + (1.0, 0.06, 0.8), + (1e-5, 6e-7, 8e-6), + (0.0, 0.0, 0.0), + (1000.0, 30.0, 400.0), + ): + with self.subTest(clip=clip, max_norm=max_norm): + clipped, total_norm = clip(grads, max_norm) + self.assertEqual(total_norm.dtype, mx.float32) + self.assertAlmostEqual(total_norm.item(), 500.0, places=3) + for actual, expected in ( + (clipped["first"][0], first), + (clipped["second"], second), + ): + self.assertEqual(actual.dtype, mx.float16) + self.assertTrue( + mx.allclose( + actual, mx.array(expected), rtol=1e-3, atol=1e-7 + ) + ) + self.assertTrue(mx.array_equal(grads["first"][0], mx.full((100,), 30.0))) + self.assertTrue(mx.array_equal(grads["second"], mx.array([400.0]))) + + grads = {"w": mx.full((16,), 30000.0, dtype=mx.float16)} + clipped, total_norm = opt.clip_grad_norm(grads, max_norm=1.0) + self.assertAlmostEqual(total_norm.item(), 120000.0, delta=0.1) + self.assertTrue(mx.array_equal(clipped["w"], mx.full((16,), 0.25))) + + def test_clip_grad_norm_mixed_dtypes(self): + grads = { + "first": mx.array([300.0], dtype=mx.float16), + "second": [ + mx.array([400.0], dtype=mx.bfloat16), + mx.array([0.0], dtype=mx.float32), + ], + } + clipped, total_norm = opt.clip_grad_norm(grads, max_norm=1.0) + self.assertTrue(mx.isfinite(total_norm)) + self.assertAlmostEqual(total_norm.item(), 500.0, delta=1.0) + for (_, g), (_, c), expected in zip( + tree_flatten(grads), tree_flatten(clipped), (0.6, 0.8, 0.0) + ): + self.assertEqual(c.dtype, g.dtype) + self.assertTrue(mx.allclose(c, mx.array(expected), rtol=1e-2, atol=1e-7)) + def test_init_from_state(self): class Model(nn.Module): def __init__(self): @@ -657,6 +708,33 @@ def test_multi_optimizer_with_parameterless_layers(self): self.assertFalse(mx.array_equal(w, model.layers[0].weight)) self.assertFalse(mx.array_equal(b, model.layers[0].bias)) + def test_multi_optimizer_with_empty_group(self): + mx.random.seed(0) + # no biases so all the weights are 2D + model = nn.Sequential( + nn.Linear(4, 4, bias=False), nn.ReLU(), nn.Linear(4, 4, bias=False) + ) + mx.eval(model.parameters()) + + loss_and_grad = nn.value_and_grad(model, lambda m, x: m(x).sum()) + _, grads = loss_and_grad(model, mx.ones((1, 4))) + expected = opt.Adam(learning_rate=0.01).apply_gradients(grads, model) + + # SGD gets no weights, first as the fallback and then as the first optimizer + optimizers = [ + opt.MultiOptimizer( + [opt.Adam(learning_rate=0.01), opt.SGD(learning_rate=0.01)], + [lambda _, w: w.ndim >= 2], + ), + opt.MultiOptimizer( + [opt.SGD(learning_rate=0.01), opt.Adam(learning_rate=0.01)], + [lambda _, w: w.ndim < 2], + ), + ] + for optimizer in optimizers: + new_params = optimizer.apply_gradients(grads, model) + self.assertTrue(tree_equal(mx.array_equal, expected, new_params)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_quantized.py b/Source/Cmlx/mlx/python/tests/test_quantized.py index 7a5547411..3c1bf0bbf 100644 --- a/Source/Cmlx/mlx/python/tests/test_quantized.py +++ b/Source/Cmlx/mlx/python/tests/test_quantized.py @@ -23,7 +23,9 @@ 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]: + if b == 1 and mx.cuda.is_available(): + continue 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) @@ -34,7 +36,9 @@ 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]: + if b == 1 and mx.cuda.is_available(): + continue 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)) @@ -224,6 +228,98 @@ def test_nvfp4_quantize_dequantize(self): ) self.assertTrue(mx.allclose(w, w_hat, rtol=1e-5, atol=1e-5)) + @unittest.skipIf(mx.cuda.is_available(), "1-bit quantization NYI on CUDA") + 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().item(), 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().item(), 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().item(), 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().item(), 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) + tol = 2e-5 if platform.system() == "Windows" else 1e-5 + self.assertLess((y_q - y_hat).abs().max().item(), tol) + def test_qqmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -276,6 +372,7 @@ def test_qqmm(self): self.skipTest("Not implemented for CPU") return + tol = 5e-2 if mx.cuda.is_available() else 1e-3 key = mx.random.key(0) k1, k2 = mx.random.split(key) tests = product( @@ -318,7 +415,7 @@ def test_qqmm(self): ) 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) + self.assertLess((y_q - y_hat).abs().max(), tol) @unittest.skipIf( not mx.metal.is_available(), "Global scale is only supported on Metal backend" @@ -372,13 +469,15 @@ 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 [True, False], # transposed ) for group_size, bits, M, N, K, transposed in tests: + if bits == 1 and mx.cuda.is_available(): + continue with self.subTest( shape=(M, N, K), group_size=group_size, @@ -558,6 +657,34 @@ def test_qmm_small_m_block(self): ) self.assertLess((y_ref - y).abs().max(), 1e-3) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmm_splitk_precision(self): + # Split-K sums the partitions in float32 and rounds the result to the + # input type once. Rounding each partition first gives two to three + # times the error of a single rounding at these shapes. + if mx.default_device() == mx.cpu: + self.skipTest("Covers GPU kernels only") + K, N = 2560, 1024 + # Keep M at or above get_qmv_batch_limit on all devices. Below it the + # qmv kernel runs, which dequantizes in float32 and not in the input type. + for mode, dtype, M in product( + ["affine", "mxfp4"], [mx.bfloat16, mx.float16], [25, 33, 65] + ): + with self.subTest(mode=mode, dtype=dtype, M=M): + k1, k2 = mx.random.split(mx.random.key(0)) + x = mx.random.normal((M, K), key=k1).astype(dtype) + w = (0.02 * mx.random.normal((N, K), key=k2)).astype(dtype) + group_size, bits = (64, 4) if mode == "affine" else (None, None) + wq = mx.quantize(w, group_size=group_size, bits=bits, mode=mode) + w_hat = mx.dequantize(*wq, group_size=group_size, bits=bits, mode=mode) + y_ref = x.astype(mx.float32) @ w_hat.astype(mx.float32).T + y = mx.quantized_matmul( + x, *wq, transpose=True, group_size=group_size, bits=bits, mode=mode + ) + error = (y.astype(mx.float32) - y_ref).abs().mean() + rounding = (y_ref.astype(dtype).astype(mx.float32) - y_ref).abs().mean() + self.assertLess(error.item(), 1.1 * rounding.item()) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -626,12 +753,13 @@ def test_qmm_shapes(self): w = mx.random.normal(shape=(32, 256), key=k2) w_q, scales, biases = mx.quantize(w, group_size, bits) w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) + tol = 2e-2 if mx.cuda.is_available() else 1e-3 for s in [(3, 256), (2, 1, 7, 256)]: x = mx.random.normal(shape=s, key=k1) y_q = mx.quantized_matmul(x, w_q, scales, biases, True, group_size, bits) y_hat = x @ w_hat.T self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) w = mx.random.normal(shape=(256, 256), key=k2) w_q, scales, biases = mx.quantize(w, group_size, bits) @@ -641,7 +769,7 @@ def test_qmm_shapes(self): y_q = mx.quantized_matmul(x, w_q, scales, biases, False, group_size, bits) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) def test_qmv(self): key = mx.random.key(0) @@ -670,9 +798,156 @@ def test_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmv_affine_bias_sum_precision(self): + # The bias contributes bias * sum(x), and a sum accumulated in the input + # type loses the small values: the sum of [1, 1 / denom, -1, 0] cancels + # to zero. + for n, k, gs, dtype, sign, bits in product( + [1, 4, 8, 12], + [64, 96, 512], + [32, 64], + [mx.bfloat16, mx.float16], + [-1, 1], + [2, 3, 4, 5, 6, 8], + ): + if k % gs: + continue + with self.subTest(n=n, k=k, gs=gs, dtype=dtype, sign=sign, bits=bits): + # 1 / denom is exactly half an ulp at 1.0, so 1 + 1 / denom + # rounds back to 1 (ties-to-even) in the input type. + denom = 2048 if dtype == mx.float16 else 256 + x = mx.tile(mx.array([1, 1 / denom, -1, 0], dtype), k // 4) + x = (sign * x).reshape(1, k) + # Zero scales and unit biases dequantize to all ones, so the + # product is exactly sum(x) = sign * k / (4 * denom). + q = mx.zeros((n, k * bits // 32), mx.uint32) + scales = mx.zeros((n, k // gs), dtype) + biases = mx.ones((n, k // gs), dtype) + y = mx.quantized_matmul(x, q, scales, biases, group_size=gs, bits=bits) + expected = mx.full((1, n), sign * k / (4 * denom), dtype) + self.assertTrue(mx.array_equal(y, expected).item()) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmv_fast_unaligned_n(self): + # Output sizes that are not a multiple of 8 with an aligned input size. + # "vector" is one input row, "batched" two weight batches, and "rows" + # is M = 3, which reaches this kernel only on GPUs before gen 15, since + # later ones route M >= 2 to qmv_wide. + for n, k, bits, gs, dtype, layout in product( + [1, 3, 5, 12, 17], + [512, 1024], + [2, 3, 4, 5, 6, 8], + [64], + [mx.float32, mx.bfloat16], + ["vector", "batched", "rows"], + ): + with self.subTest(n=n, k=k, bits=bits, gs=gs, dtype=dtype, layout=layout): + key = mx.random.key(n * 7 + k + bits) + if layout == "batched": + w_shape, x_shape = (2, n, k), (2, 1, k) + elif layout == "rows": + w_shape, x_shape = (n, k), (3, k) + else: + w_shape, x_shape = (n, k), (1, k) + w = mx.random.normal(w_shape, key=key) / k**0.5 + x = mx.random.normal(x_shape, key=mx.random.split(key)[0]) + x = (x / k**0.5).astype(dtype) + q, s, b = mx.quantize(w.astype(dtype), group_size=gs, bits=bits) + w_hat = mx.dequantize(q, s, b, group_size=gs, bits=bits) + y_hat = x.astype(mx.float32) @ mx.swapaxes( + w_hat.astype(mx.float32), -1, -2 + ) + y = mx.quantized_matmul(x, q, s, b, group_size=gs, bits=bits) + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 + self.assertEqual(y.shape, y_hat.shape) + self.assertLess((y - y_hat).abs().max().item(), tol) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_fp_qmv_fast_unaligned_n(self): + for n, k, mode, dtype, batched in product( + [1, 3, 5, 7, 8, 9, 12, 17], + [256, 512, 544, 1024], + ["mxfp4", "mxfp8", "nvfp4"], + [mx.float32, mx.float16, mx.bfloat16], + [False, True], + ): + with self.subTest(n=n, k=k, mode=mode, dtype=dtype, batched=batched): + key = mx.random.key(n * 7 + k) + w_shape = (2, n, k) if batched else (n, k) + x_shape = (2, 1, k) if batched else (1, k) + w = mx.random.normal(w_shape, key=key) / k**0.5 + x = mx.random.normal(x_shape, key=mx.random.split(key)[0]) + x = (x / k**0.5).astype(dtype) + q, s = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(q, s, mode=mode, dtype=mx.float32) + expected = x.astype(mx.float32) @ mx.swapaxes(w_hat, -1, -2) + actual = mx.quantized_matmul(x, q, s, mode=mode) + self.assertEqual(actual.shape, expected.shape) + self.assertLess((actual - expected).abs().max().item(), 1e-3) + + # Compare aligned inputs with the existing full-tile kernel. + alignment = 256 if mode == "mxfp8" else 512 + if k % alignment == 0: + padding = [(0, 0)] * q.ndim + padding[-2] = (0, (-n) % 8) + padded = mx.quantized_matmul( + x, mx.pad(q, padding), mx.pad(s, padding), mode=mode + ) + self.assertTrue(mx.array_equal(actual, padded[..., :n]).item()) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_fp_qmv_fast_unaligned_n_global_scale(self): + for n, k, m, dtype in product( + [1, 3, 5, 7, 8, 9, 12, 17], + [512, 544, 1024], + [1, 3], + [mx.float32, mx.float16, mx.bfloat16], + ): + with self.subTest(n=n, k=k, m=m, dtype=dtype): + key = mx.random.key(n * 7 + k) + w = mx.random.normal((n, k), key=key) / k**0.5 + x = mx.random.normal((m, k), key=mx.random.split(key)[0]) + x = (x / k**0.5).astype(dtype) + global_scale_x = mx.max(mx.abs(x)).astype(mx.float32) + global_scale_w = mx.max(mx.abs(w)).astype(mx.float32) + q, s = mx.quantize(w, mode="nvfp4", global_scale=global_scale_w) + w_hat = mx.dequantize( + q, s, mode="nvfp4", global_scale=global_scale_w, dtype=mx.float32 + ) + x_hat = mx.dequantize( + *mx.quantize(x, mode="nvfp4", global_scale=global_scale_x), + mode="nvfp4", + global_scale=global_scale_x, + dtype=dtype, + ) + expected = x_hat.astype(mx.float32) @ w_hat.T + actual = mx.qqmm( + x, + q, + s, + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + self.assertEqual(actual.shape, expected.shape) + self.assertLess((actual - expected).abs().max().item(), 1e-3) + + if k % 512 == 0: + padded = mx.qqmm( + x, + mx.pad(q, [(0, (-n) % 8), (0, 0)]), + mx.pad(s, [(0, (-n) % 8), (0, 0)]), + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + self.assertTrue(mx.array_equal(actual, padded[..., :n]).item()) + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) + tol = 2e-2 if mx.cuda.is_available() else 1e-3 tests = product( [256, 512, 67, 5, 7], # M -- 5, 7 exercise out_vec_size < 8 (#3762) [64, 256], # N @@ -697,7 +972,7 @@ def test_fp_qmv(self): ) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test multiple of 16 but not 32 M = 128 @@ -719,7 +994,7 @@ def test_fp_qmv(self): ) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) def test_fp_qmv_large_output(self): key = mx.random.key(0) @@ -757,6 +1032,7 @@ def test_fp_qmv_large_output(self): self.assertTrue(mx.allclose(y_q, y_hat, rtol=tol, atol=tol)) def test_qmv_wide(self): + tol = 5e-2 if mx.cuda.is_available() else 1e-3 # 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 # against a dequantize-then-matmul reference, with ragged M (token @@ -787,7 +1063,7 @@ def test_qmv_wide(self): ) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # FP modes (group_size and bits implied by the mode). for mode, K in product(["mxfp4", "nvfp4", "mxfp8"], [128, 512]): @@ -802,7 +1078,7 @@ def test_qmv_wide(self): y_q = mx.quantized_matmul(x, w_q, scales, transpose=True, mode=mode) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # 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)] @@ -831,7 +1107,7 @@ def test_qmv_wide(self): ) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) def test_qvm(self): key = mx.random.key(0) @@ -955,7 +1231,7 @@ def test_fp_qvm(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 2e-3) - FP_TAIL_TOLS = {mx.float32: 1e-3, mx.float16: 5e-3, mx.bfloat16: 4e-2} + FP_TAIL_TOLS = {mx.float32: 2e-3, mx.float16: 5e-3, mx.bfloat16: 4e-2} def _fp_tail_dtypes(self): # The half types only run on the GPU @@ -1198,6 +1474,7 @@ def test_throw(self): mx.eval(y) def test_small_matrix(self): + tol = 3e-2 if mx.cuda.is_available() else 1e-3 for w_shape in [(8, 256), (1, 8, 256), (3, 8, 256)]: with self.subTest(w_shape=w_shape): w = mx.random.normal(shape=(w_shape)) @@ -1210,30 +1487,31 @@ def test_small_matrix(self): y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm_t x = mx.random.normal(shape=(3, 10, 256)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qvm x = mx.random.normal(shape=(3, 1, 8)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm x = mx.random.normal(shape=(3, 10, 8)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) def test_non_multiples(self): + tol = 3e-2 if mx.cuda.is_available() else 1e-3 w = mx.random.normal(shape=(33, 256)) w_q, scales, biases = mx.quantize(w) w_hat = mx.dequantize(w_q, scales, biases) @@ -1242,28 +1520,28 @@ def test_non_multiples(self): x = mx.random.normal(shape=(1, 256)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ w_hat.T - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm_t x = mx.random.normal(shape=(10, 256)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ w_hat.T self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qvm x = mx.random.normal(shape=(1, 33)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm x = mx.random.normal(shape=(10, 33)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Smaller than 8 w = mx.random.normal(shape=(3, 256)) @@ -1274,28 +1552,28 @@ def test_non_multiples(self): x = mx.random.normal(shape=(1, 256)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ w_hat.T - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm_t x = mx.random.normal(shape=(10, 256)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ w_hat.T self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qvm x = mx.random.normal(shape=(1, 3)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test qmm x = mx.random.normal(shape=(10, 3)) y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) y_hat = x @ w_hat self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) # Test with larger than 128 unaligned sizes w = mx.random.normal(shape=(99, 256)) @@ -1305,7 +1583,7 @@ def test_non_multiples(self): y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) y_hat = x @ w_hat.T self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + self.assertLess((y_q - y_hat).abs().max(), tol) def test_qmv_small_non_multiples(self): # Test very small K and N dimensions (e.g., [MxK] x [NxK].T = [MxN]) @@ -1391,6 +1669,9 @@ def test_shape( group_size=None, bits=None, mode="affine", + noncontiguous_x=False, + rtol=1e-5, + atol=1e-4, ): with self.subTest( M=M, @@ -1405,6 +1686,9 @@ def test_shape( group_size=group_size, bits=bits, mode=mode, + noncontiguous_x=noncontiguous_x, + rtol=rtol, + atol=atol, ): x = mx.random.normal(shape=batch_A + (M, K)).astype(dtype) w = mx.random.normal( @@ -1412,6 +1696,9 @@ def test_shape( ).astype(dtype) w_hat, qw, s, b = quantize(w, transpose, group_size, bits, mode=mode) + if noncontiguous_x: + x = mx.concatenate([x, x], axis=-1)[..., : x.shape[-1]] + if lhs_indices is not None: lhs_indices = mx.array(lhs_indices) if rhs_indices is not None: @@ -1430,7 +1717,7 @@ def test_shape( bits=bits, mode=mode, ) - self.assertTrue(mx.allclose(c1, c2, atol=1e-4)) + self.assertTrue(mx.allclose(c1, c2, rtol=rtol, atol=atol)) inputs = ( { @@ -1507,6 +1794,71 @@ def test_shape( test_shape(32, 512, 32, transpose=False, **kwargs) test_shape(1, 512, 32, transpose=False, **kwargs) + # Long K with many gathered vectors uses the nvfp4 gather qmv kernel on + # the GPU. Low precision outputs need a tolerance of about one ulp. + on_gpu = mx.default_device() == mx.gpu + bf16 = dict(dtype=mx.bfloat16 if on_gpu else mx.float32, rtol=1e-2, atol=1e-2) + f16 = dict(dtype=mx.float16 if on_gpu else mx.float32, rtol=1e-3, atol=1e-3) + test_shape( + 1, + 32, + 1024, + batch_A=(1,), + batch_B=(3,), + lhs_indices=(0,), + rhs_indices=(2, 1, 0, 2, 1, 0, 2, 1, 0), + mode="nvfp4", + **bf16, + ) + test_shape( + 1, + 33, + 1024, + batch_A=(2, 2), + batch_B=(2,), + lhs_indices=((0,), (3,)), + rhs_indices=((1, 0, 1),), + mode="nvfp4", + noncontiguous_x=True, + **f16, + ) + # K = 1056 leaves one warp lane active in the second reduction step. + test_shape( + 1, + 32, + 1056, + batch_A=(1,), + batch_B=(3,), + lhs_indices=(0,), + rhs_indices=(2, 1, 0, 2, 1, 0, 2, 1, 0), + mode="nvfp4", + ) + for mode in ("mxfp4", "mxfp8"): + test_shape( + 1, + 33, + 1024, + batch_A=(1,), + batch_B=(3,), + lhs_indices=(0,), + rhs_indices=(2, 1, 0, 2, 1, 0, 2, 1, 0), + mode=mode, + **bf16, + ) + for mode in ("nvfp4", "mxfp4", "mxfp8"): + for K, tol in ((448, 1e-1), (512, 1e-2), (704, 1e-2)): + test_shape( + 1, + 33, + K, + batch_A=(1,), + batch_B=(3,), + lhs_indices=(0,), + rhs_indices=(2, 1, 0, 2, 1, 0, 2, 1, 0), + mode=mode, + **{**bf16, "rtol": tol, "atol": tol}, + ) + def test_gather_qqmm(self): if mx.default_device() == mx.cpu: self.skipTest("Not implemented for CPU") @@ -1655,7 +2007,8 @@ def check(x, rhs, lhs=None, sorted_indices=False): ) error = mx.abs(actual.astype(mx.float32) - expected.astype(mx.float32)) scale = mx.maximum(mx.abs(expected.astype(mx.float32)).max(), 1e-20) - self.assertLess(error.max() / scale, 3e-2) + tol = 1e-1 if mx.cuda.is_available() else 3e-2 + self.assertLess(error.max() / scale, tol) # Sorted indices select the RHS-grouped matrix kernel. x = mx.random.normal((24, 1, K), key=mx.random.key(20)).astype(mx.bfloat16) @@ -1808,8 +2161,10 @@ def scatter_unsort(x, inv_order, shape=None): k1, k2, k3 = mx.random.split(key, 3) dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 - for L, K, D, E, I, transpose, mode in parameters: - with self.subTest(L=L, K=K, D=D, E=E, I=I, transpose=transpose, mode=mode): + for L, K, D, E, num_indices, transpose, mode in parameters: + with self.subTest( + L=L, K=K, D=D, E=E, I=num_indices, transpose=transpose, mode=mode + ): if mode != "affine": group_size = None dtype = ( @@ -1822,7 +2177,7 @@ def scatter_unsort(x, inv_order, shape=None): ) K, D = (K, D) if transpose else (D, K) - ishape = (L, I) + ishape = (L, num_indices) xshape = (L, 1, 1, K) wshape = (E, D, K) if transpose else (E, K, D) @@ -1987,6 +2342,71 @@ def test_gather_qmm_sorted_unaligned_k(self): tol = 1e-3 if on_gpu else 1.5e-5 self.assertLess((y_sorted - y_unsorted).abs().max(), tol) + @unittest.skipIf(not mx.cuda.is_available(), "CUDA kernel path only") + def test_gather_qmm_sorted_rhs_cuda(self): + # Sorted rows with a shared LHS take the sorted-RHS matrix kernel. Row + # counts cover the 16/32/64-row tiles and a row tail; random routing + # gives runs of unequal length that cross tile edges. + mx.random.seed(0) + shapes = [(64, 4, 2816, 704)] + [ + (rows, 8, 256, 256) for rows in (16, 24, 40, 64, 100, 200, 1000) + ] + for rows, experts, n, k in shapes: + for dtype in (mx.bfloat16, mx.float16): + x = (mx.random.normal((rows, 1, 1, k)) / k**0.5).astype(dtype) + w = mx.random.normal((experts, n, k)).astype(dtype) + indices = mx.sort(mx.random.randint(0, experts, (rows,))).astype( + mx.uint32 + )[:, None] + for mode in ("nvfp4", "mxfp4", "mxfp8"): + with self.subTest(rows=rows, dtype=dtype, mode=mode): + qw, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(qw, scales, mode=mode) + reference = mx.gather_mm( + x, w_hat.swapaxes(-1, -2), rhs_indices=indices + ) + expected = mx.gather_qmm( + x, + qw, + scales, + rhs_indices=indices, + transpose=True, + mode=mode, + sorted_indices=False, + ) + mx.eval(reference, expected) + # Successive runs in one CTA reuse its shared memory; + # repeat to catch a race between them. + for _ in range(3): + actual = mx.gather_qmm( + x, + qw, + scales, + rhs_indices=indices, + transpose=True, + mode=mode, + sorted_indices=True, + ) + mx.eval(actual) + self.assertLess( + mx.max(mx.abs(expected - actual)).item(), 1e-2 + ) + self.assertTrue( + mx.allclose(reference, actual, rtol=2e-2, atol=2e-2) + ) + + # One x row broadcast over the sorted indices stays on the generic + # gathered kernel. + rows, experts, n, k = 16, 4, 256, 256 + x = mx.random.normal((1, 1, k)).astype(mx.bfloat16) + w = mx.random.normal((experts, n, k)).astype(mx.bfloat16) + qw, scales = mx.quantize(w, mode="mxfp8") + indices = mx.repeat(mx.arange(experts), rows // experts).astype(mx.uint32) + kwargs = dict(rhs_indices=indices, transpose=True, mode="mxfp8") + expected = mx.gather_qmm(x, qw, scales, sorted_indices=False, **kwargs) + actual = mx.gather_qmm(x, qw, scales, sorted_indices=True, **kwargs) + self.assertLess(mx.max(mx.abs(expected - actual)).item(), 1e-3) + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: @@ -2032,8 +2452,9 @@ def gather_qmm(x, w, s, b, lhs, rhs, trans, sort): self.assertLess((o1 - o2).abs().max(), 1e-4) self.assertTrue(mx.allclose(o1, o2, atol=1e-4)) self.assertTrue(mx.allclose(dx1, dx2, atol=1e-4)) - self.assertTrue(mx.allclose(ds1, ds2, atol=1e-3)) - self.assertTrue(mx.allclose(db1, db2, atol=1e-3)) + tol = 5e-1 if mx.cuda.is_available() else 1e-3 + self.assertTrue(mx.allclose(ds1, ds2, atol=tol)) + self.assertTrue(mx.allclose(db1, db2, atol=tol)) def test_vjp_scales_biases(self): mx.random.seed(0) @@ -2048,6 +2469,7 @@ def mm(sb, x, wq): dparams = mx.grad(mm)((s, b), x, wq) eps = 8e-3 + delta = 5e-1 if mx.cuda.is_available() else 2e-2 # numerical grad check with a few indices indices = [(0, 0), (11, 4), (22, 7)] for idx in indices: @@ -2058,7 +2480,7 @@ def mm(sb, x, wq): out_down = mm(params, x, wq) params[p][idx] += eps num_ds = (out_up - out_down) / (2 * eps) - self.assertAlmostEqual(dparams[p][idx], num_ds, delta=2e-2) + self.assertAlmostEqual(dparams[p][idx], num_ds, delta=delta) def test_fp_vjp_scales_throws(self): mx.random.seed(0) @@ -2072,7 +2494,7 @@ def mm(s, x, wq): # Should raise with self.assertRaises(ValueError): - ds = mx.grad(mm)(s, x, wq) + mx.grad(mm)(s, x, wq) rhs_indices = mx.array(0) with self.assertRaises(ValueError): @@ -2086,7 +2508,7 @@ def gmm(s, x, wq): mode=mode, ).sum() - ds = mx.grad(gmm)(s, x, wq) + _ds = mx.grad(gmm)(s, x, wq) @unittest.skipIf( not mx.is_available(mx.gpu), "Global scale is only supported on the GPU" @@ -2155,7 +2577,9 @@ def quantize_experts(w): ) out = mx.gather_qmm(x, wq, s, global_scale=gs, **kwargs) - tol = 1e-5 if dtype == mx.float32 else 3e-2 + tol = ( + 1e-5 if dtype == mx.float32 and not mx.cuda.is_available() else 3e-2 + ) self.assertLess(rel_err(out, expected), tol) # Each expert uses its own scale, not a neighbour's @@ -2197,6 +2621,28 @@ def test_quantize_strided(self): expected = mx.dequantize(w_q, mx.contiguous(scales), mode=mode) self.assertTrue(mx.allclose(w_hat, expected)) + def test_quantized_matmul_batched_split_k(self): + tol = 3e-2 if mx.cuda.is_available() else 1e-3 + K, N = 2048, 512 + w = mx.random.normal(shape=(K, N)) + w_q, scales, biases = mx.quantize(w, group_size=64, bits=4) + w_hat = mx.dequantize(w_q, scales, biases, group_size=64, bits=4) + + shapes = [ + (2, K), + (2, 1, K), + (3, 1, K), + (2, 1, 1, K), + ] + for shape in shapes: + x = mx.random.normal(shape=shape) + y_q = mx.quantized_matmul( + x, w_q, scales, biases, transpose=False, group_size=64, bits=4 + ) + y_hat = x @ w_hat + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), tol) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_random.py b/Source/Cmlx/mlx/python/tests/test_random.py index 96a789b60..c0c6d57cc 100644 --- a/Source/Cmlx/mlx/python/tests/test_random.py +++ b/Source/Cmlx/mlx/python/tests/test_random.py @@ -1,7 +1,6 @@ # Copyright © 2023 Apple Inc. import math -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_reduce.py b/Source/Cmlx/mlx/python/tests/test_reduce.py index 164e2dd80..b46dcee1e 100644 --- a/Source/Cmlx/mlx/python/tests/test_reduce.py +++ b/Source/Cmlx/mlx/python/tests/test_reduce.py @@ -1,5 +1,7 @@ # Copyright © 2023 Apple Inc. +import math +import unittest from itertools import combinations, permutations import mlx.core as mx @@ -56,6 +58,33 @@ def test_row_reduce_negative_stride(self): actual = getattr(mx, op)(x_mlx, axis=-1) self.assertTrue(np.allclose(expected, actual)) + def test_col_reduce_negative_stride(self): + # Exercises each Metal column reduction kernel on a negative-stride view. + cases = [ + ((2, 1024, 16), 1), + ((2, 2, 1024, 16), 2), + ((2, 512, 2, 2, 16), (1, 3)), + ((2, 512, 64), 1), + ((2, 64, 512), 1), + ((2, 16, 16), 1), + ] + for shape, axis in cases: + size = math.prod(shape) + x_npy = np.arange(1, size + 1).reshape(shape)[::-1] + x_mlx = mx.arange(1, size + 1).reshape(shape)[::-1] + for op in ["sum", "max", "min", "mean", "var"]: + with self.subTest(shape=shape, axis=axis, op=op): + expected = getattr(np, op)(x_npy, axis=axis) + actual = getattr(mx, op)(x_mlx, axis=axis) + self.assertTrue(np.allclose(expected, actual)) + x_ones = mx.ones(shape, dtype=mx.float32)[::-1] + with self.subTest(shape=shape, axis=axis, op="prod"): + self.assertTrue(mx.all(mx.prod(x_ones, axis=axis) == 1).item()) + x_bool = mx.ones(shape, dtype=mx.bool_)[::-1] + for op in ["all", "any"]: + with self.subTest(shape=shape, axis=axis, op=op): + self.assertTrue(mx.all(getattr(mx, op)(x_bool, axis=axis)).item()) + def test_dtypes(self): int_dtypes = [ "int8", @@ -159,6 +188,35 @@ def test_zero_size(self): with self.assertRaises(ValueError): getattr(mx, op)(a_mx, axis=axis) + # sum and prod have identities, so an empty reduction returns them for + # every dtype. The unsigned outputs (uint32, uint64) had no init kernel + # on Metal and aborted the process instead. + dtypes = [ + mx.bool_, + mx.uint8, + mx.uint16, + mx.uint32, + mx.uint64, + mx.int8, + mx.int16, + mx.int32, + mx.int64, + mx.float16, + mx.float32, + ] + for dtype in dtypes: + a = mx.zeros((0,), dtype=dtype) + for op, identity in [("sum", 0), ("prod", 1)]: + out = getattr(mx, op)(a) + mx.eval(out) + self.assertEqual(out.item(), identity, f"{op} of empty {dtype}") + # empty because of another axis, reduced over a non-empty one + b = mx.zeros((0, 3), dtype=dtype) + for op in ["sum", "prod"]: + out = getattr(mx, op)(b, axis=-1) + mx.eval(out) + self.assertEqual(out.shape, (0,)) + def test_sum_bool(self): x = np.random.uniform(0, 1, size=(10, 10, 10)) > 0.5 y = mx.array(x) @@ -256,7 +314,7 @@ def test_long_column(self): 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"]: + for dtype in ["float32", "float16", "float64", "complex64"]: with self.subTest(dtype=dtype): for values in [ [0.0, -0.0], @@ -282,6 +340,48 @@ def test_and_or_negative_zero(self): getattr(np, op)(x_np, axis=1).tolist(), ) + def test_large_offsets(self): + # Row r holds r % 251 and row 2**15 starts at offset 2**31, so any + # reduction that narrows offsets or sizes to 32 bits reads the wrong + # rows. The views below have fewer than 2**31 elements themselves. + rows, cols = 2**15 + 1, 2**16 + row_max = (mx.arange(rows) % 251).astype(mx.uint8) + x = mx.contiguous(mx.broadcast_to(row_max[:, None], (rows, cols))) + + self.assertTrue(mx.array_equal(x[:, 7:].max(axis=-1), row_max)) + self.assertTrue(mx.array_equal(x[:, 7:].min(axis=-1), row_max)) + y = x.reshape(rows, 256, 256)[:, 1:, 1:] + self.assertTrue(mx.array_equal(y.max(axis=(1, 2)), row_max)) + y = x.reshape(rows, 16, 16, 256)[:, 1:, :, 1:] + expected = mx.broadcast_to(row_max[:, None], (rows, 16)) + self.assertTrue(mx.array_equal(y.max(axis=(1, 3)), expected)) + + # Reducing all 2**31 + 2**16 elements at once + self.assertEqual(x.max().item(), 250) + self.assertTrue(x.any().item()) + self.assertFalse(x.all().item()) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_init_reduce_large_output(self): + out = mx.all(mx.zeros((0, 4, 2**30), dtype=mx.bool_), axis=0) + mx.eval(out) + self.assertTrue(mx.all(out).item()) + + def test_minmax(self): + x = mx.array([[3.0, 1.0, 4.0], [2.0, 6.0, 5.0]]) + out = mx.minmax(x) + self.assertIsInstance(out, tuple) + self.assertEqual(len(out), 2) + self.assertEqual(tuple(a.item() for a in out), (1.0, 6.0)) + for axis in (None, 0, (0, 1)): + for keepdims in (False, True): + with self.subTest(axis=axis, keepdims=keepdims): + out = mx.minmax(x, axis=axis, keepdims=keepdims) + for actual, op in zip(out, (mx.min, mx.max)): + expected = op(x, axis=axis, keepdims=keepdims) + self.assertEqual(actual.shape, expected.shape) + self.assertTrue(mx.array_equal(actual, expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True) diff --git a/Source/Cmlx/mlx/python/tests/test_threads.py b/Source/Cmlx/mlx/python/tests/test_threads.py index e939ca83a..3d2aa6519 100644 --- a/Source/Cmlx/mlx/python/tests/test_threads.py +++ b/Source/Cmlx/mlx/python/tests/test_threads.py @@ -1,7 +1,6 @@ # Copyright © 2026 Apple Inc. import threading -import unittest import mlx.core as mx import mlx_tests diff --git a/Source/Cmlx/mlx/python/tests/test_tree.py b/Source/Cmlx/mlx/python/tests/test_tree.py index 96184bd23..c7d2fadaa 100644 --- a/Source/Cmlx/mlx/python/tests/test_tree.py +++ b/Source/Cmlx/mlx/python/tests/test_tree.py @@ -1,6 +1,5 @@ # Copyright © 2023 Apple Inc. -import unittest import mlx.core as mx import mlx.nn as nn @@ -97,7 +96,8 @@ class Params(NamedTuple): params1 = Params(m=mx.array([0, 1]), b=mx.array(2)) dict1 = {"m": mx.array([0, 1]), "b": mx.array(2)} - add_one = lambda x: x + 1 + def add_one(x): + return x + 1 list2 = mlx.utils.tree_map(add_one, list1) tuple2 = mlx.utils.tree_map(add_one, tuple1) diff --git a/Source/Cmlx/mlx/python/tests/test_upsample.py b/Source/Cmlx/mlx/python/tests/test_upsample.py index 9a8ab9184..d71ff90ea 100644 --- a/Source/Cmlx/mlx/python/tests/test_upsample.py +++ b/Source/Cmlx/mlx/python/tests/test_upsample.py @@ -12,7 +12,7 @@ import torch.nn.functional as F has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -97,7 +97,9 @@ def run_upsample( @unittest.skipIf(not has_torch, "requires Torch") def test_torch_upsample_antialias(self): - """Test antialiased downsampling matches PyTorch F.interpolate(antialias=True).""" + """Test antialiased downsampling matches + PyTorch F.interpolate(antialias=True). + """ def run_antialias( N, @@ -146,7 +148,8 @@ def run_antialias( 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"antialias {mode} ac={align_corners}" + f" scale={scale_factor} max_diff=" f"{np.abs(out_pt - np.array(out_mx)).max():.2e}", ) diff --git a/Source/Cmlx/mlx/python/tests/test_vmap.py b/Source/Cmlx/mlx/python/tests/test_vmap.py index 7538b559d..570947d25 100644 --- a/Source/Cmlx/mlx/python/tests/test_vmap.py +++ b/Source/Cmlx/mlx/python/tests/test_vmap.py @@ -1,7 +1,6 @@ # Copyright © 2023-2024 Apple Inc. import gc -import unittest import mlx.core as mx import mlx_tests @@ -219,6 +218,19 @@ def test_vmap_indexing(self): ) self.assertTrue(mx.array_equal(out, expected)) + def test_vmap_strided_slice_single_element(self): + # A strided slice selecting exactly one element must vmap to the + # batched version of what the un-batched slice returns (regression: + # the stored stride was collapsed to 1 without narrowing stop, so + # re-deriving the region returned every element in the span). + x = mx.arange(24, dtype=mx.float32).reshape(4, 2, 3) + out = mx.vmap(lambda t: t[0::2])(x) + self.assertTrue(mx.array_equal(out, x[:, 0:1])) + + y = mx.arange(48, dtype=mx.float32).reshape(4, 4, 3) + out = mx.vmap(lambda t: t[0::2])(y) + self.assertTrue(mx.array_equal(out, y[:, 0:4:2])) + def test_vmap_reduce(self): a = mx.ones((5, 5), mx.int32) out = mx.vmap(lambda x: x.sum())(a) @@ -347,11 +359,11 @@ def test_mismatch_input_sizes(self): b = mx.ones((1, 1, 1, 5)) with self.assertRaises(ValueError): - out = mx.vmap(lambda x, y: x + y)(a, b) + mx.vmap(lambda x, y: x + y)(a, b) b = mx.ones((10, 5)) with self.assertRaises(ValueError): - out = mx.vmap(lambda x, y: x + y, in_axes=(0, 1))(a, b) + mx.vmap(lambda x, y: x + y, in_axes=(0, 1))(a, b) def test_vmap_matmul(self): a = mx.random.uniform(shape=(2, 3, 4)) @@ -395,10 +407,14 @@ def test_vmap_matmul(self): def test_vmap_svd(self): a = mx.random.uniform(shape=(3, 4, 2)) - cpu_svd_full = lambda x: mx.linalg.svd(x, compute_uv=True, stream=mx.cpu) - cpu_svd_singular = lambda x: mx.linalg.svd(x, compute_uv=False, stream=mx.cpu) + def cpu_svd_full(x): + return mx.linalg.svd(x, compute_uv=True, stream=mx.cpu) + + def cpu_svd_singular(x): + return mx.linalg.svd(x, compute_uv=False, stream=mx.cpu) - # Vmap over the first axis (this is already supported natively by the primitive). + # Vmap over the first axis (this is already supported natively by the + # primitive). Us, Ss, Vts = mx.vmap(cpu_svd_full, in_axes=(0,))(a) self.assertEqual(Us.shape, (a.shape[0], a.shape[1], a.shape[1])) self.assertEqual(Ss.shape, (a.shape[0], a.shape[2])) @@ -450,9 +466,11 @@ def test_vmap_inverse(self): mx.random.seed(42) a = mx.random.uniform(shape=(3, 4, 4)) - cpu_inv = lambda x: mx.linalg.inv(x, stream=mx.cpu) + def cpu_inv(x): + return mx.linalg.inv(x, stream=mx.cpu) - # Vmap over the first axis (this is already supported natively by the primitive). + # Vmap over the first axis (this is already supported natively by the + # primitive). invs = mx.vmap(cpu_inv, in_axes=(0,))(a) for i in range(a.shape[0]): @@ -547,6 +565,41 @@ def scatter_add(a): expected = mx.array([[2.0, 3.0, 4.0], [2.0, 3.0, 4.0]]) self.assertTrue(mx.allclose(out, expected)) + def test_vmap_scatter_higher_rank(self): + # The vmap axis becomes an extra scattered source axis, so the + # singleton added to the updates has to land at that axis rather than + # at the front of the source dims. Only shows up for a vmap axis >= 2, + # where the misplaced singleton actually reorders a non-unit dim. + def unstack(x, axis): + return [s.squeeze(axis) for s in mx.split(x, x.shape[axis], axis=axis)] + + for shape in [(2, 3), (2, 3, 4), (2, 3, 4, 5)]: + n = 1 + for d in shape: + n *= d + a = mx.arange(n, dtype=mx.float32).reshape(shape) + + fns = { + "add_derived": lambda x: x.at[mx.array([0])].add(x[:1]), + "add_const": lambda x: x.at[mx.array([0])].add( + mx.ones((1,) + tuple(x.shape[1:]), dtype=x.dtype) + ), + "add_dup": lambda x: x.at[mx.array([0, 0])].add( + mx.ones((2,) + tuple(x.shape[1:]), dtype=x.dtype) + ), + "maximum": lambda x: x.at[mx.array([0])].maximum(x[:1] + 1.0), + } + for name, fn in fns.items(): + for ax in range(len(shape)): + out = mx.vmap(fn, in_axes=ax, out_axes=ax)(a) + expected = mx.stack([fn(s) for s in unstack(a, ax)], axis=ax) + self.assertEqual( + out.shape, expected.shape, f"{name} ax{ax} {shape}" + ) + self.assertTrue( + mx.array_equal(out, expected), f"{name} ax{ax} {shape}" + ) + # Multiple indices def scatter(a): a[mx.array([0, 1]), mx.array([0, 1])] = mx.array((1.0, 1.0)) @@ -866,7 +919,7 @@ def transform_tuple(t): def transform_vector(t): return Vector([t[0] + 10, t[1] * 10]) - x = State(mx.array(1), mx.array(2)) + _x = State(mx.array(1), mx.array(2)) vmap_transform = mx.vmap(transform) vmap_transform_tuple = mx.vmap(transform_tuple) @@ -1037,6 +1090,17 @@ def test_vmap_argsort(self): out = mx.vmap(lambda x: mx.argsort(x, axis=axis), in_axes=1, out_axes=1)(a) self.assertTrue(mx.array_equal(out, expected)) + def test_vmap_view(self): + a = mx.arange(16, dtype=mx.uint8).reshape(4, 4) + b = mx.arange(16, dtype=mx.int32).reshape(4, 4) + for x, dtype in [(a, mx.int8), (a, mx.int16), (a, mx.int32), (b, mx.uint8)]: + for in_axes in [0, 1, -1]: + expected = mx.stack( + [mx.take(x, i, axis=in_axes).view(dtype) for i in range(4)] + ) + out = mx.vmap(lambda y: y.view(dtype), in_axes=in_axes)(x) + self.assertTrue(mx.array_equal(out, expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_zero_copy.py b/Source/Cmlx/mlx/python/tests/test_zero_copy.py index 248e93a76..296bfb1cf 100644 --- a/Source/Cmlx/mlx/python/tests/test_zero_copy.py +++ b/Source/Cmlx/mlx/python/tests/test_zero_copy.py @@ -2,6 +2,8 @@ import faulthandler import gc +import os +import tempfile import threading import unittest @@ -53,6 +55,35 @@ def test_copy_false_dtype_conversion_raises(self): with self.assertRaises(Exception): mx.asarray(a, dtype=mx.float32, copy=False) + def test_misaligned_source_copies(self): + # Pointer not aligned to the item size + raw = np.arange(64, dtype=np.uint8) + a = raw[6:22].view(np.int32) + x = mx.asarray(a) + self.assertTrue(np.array_equal(x + 1, a + 1)) + with self.assertRaises(ValueError): + mx.asarray(a, copy=False) + + def test_ops_do_not_write_to_source(self): + a = np.ones(4096, np.float32) + mx.eval(mx.exp(a)) + self.assertTrue(np.array_equal(a, np.ones(4096, np.float32))) + + b = np.ones(4096, np.float32).tobytes() + mx.eval(mx.exp(mx.asarray(np.frombuffer(b, dtype=np.float32)))) + self.assertEqual(b, np.ones(4096, np.float32).tobytes()) + + def test_ops_on_read_only_memmap(self): + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "a.npy") + np.save(path, np.ones(1 << 16, np.float32)) + a = np.load(path, mmap_mode="r") + y = mx.exp(mx.asarray(a)) + mx.eval(y) + self.assertTrue(np.allclose(np.array(y), np.exp(1.0))) + self.assertTrue(np.array_equal(np.load(path), np.ones(1 << 16))) + del a, y + def test_source_lifetime(self): if not mx.metal.is_available(): self.skipTest("copy=False requires Metal") diff --git a/Source/Cmlx/mlx/python/tests/xp-tests/skips.txt b/Source/Cmlx/mlx/python/tests/xp-tests/skips.txt new file mode 100644 index 000000000..2bd738bb7 --- /dev/null +++ b/Source/Cmlx/mlx/python/tests/xp-tests/skips.txt @@ -0,0 +1,7 @@ +# 64bit dtype not supported https://github.com/ml-explore/mlx/issues/4300 +array_api_tests/test_creation_functions.py::test_eye + +array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None] +array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1] +array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[None] +array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[1] \ No newline at end of file diff --git a/Source/Cmlx/mlx/python/tests/xp-tests/xfail.txt b/Source/Cmlx/mlx/python/tests/xp-tests/xfail.txt new file mode 100644 index 000000000..3e08d51b2 --- /dev/null +++ b/Source/Cmlx/mlx/python/tests/xp-tests/xfail.txt @@ -0,0 +1,420 @@ +# test_linalg.py +## Missing Extensions (linalg, fft) +array_api_tests/test_linalg.py::test_cholesky +array_api_tests/test_linalg.py::test_det +array_api_tests/test_linalg.py::test_diagonal +array_api_tests/test_linalg.py::test_eigh +array_api_tests/test_linalg.py::test_eigvalsh +array_api_tests/test_linalg.py::test_eig +array_api_tests/test_linalg.py::test_eigvals +array_api_tests/test_linalg.py::test_inv +array_api_tests/test_linalg.py::test_linalg_matmul +array_api_tests/test_linalg.py::test_matmul +array_api_tests/test_linalg.py::test_matrix_norm +array_api_tests/test_linalg.py::test_matrix_power +array_api_tests/test_linalg.py::test_matrix_rank +array_api_tests/test_linalg.py::test_linalg_matrix_transpose +array_api_tests/test_linalg.py::test_outer +array_api_tests/test_linalg.py::test_pinv +array_api_tests/test_linalg.py::test_qr +array_api_tests/test_linalg.py::test_slogdet +array_api_tests/test_linalg.py::test_solve +array_api_tests/test_linalg.py::test_svd +array_api_tests/test_linalg.py::test_svdvals +array_api_tests/test_linalg.py::test_linalg_tensordot +array_api_tests/test_linalg.py::test_tensordot +array_api_tests/test_linalg.py::test_trace +array_api_tests/test_linalg.py::test_linalg_vecdot +array_api_tests/test_linalg.py::test_vecdot +array_api_tests/test_linalg.py::test_vecdot_conj +array_api_tests/test_linalg.py::test_vector_norm +array_api_tests/test_linalg.py::test_eigvals + +# test_fft.py +## Missing Extensions (linalg, fft) +array_api_tests/test_fft.py::test_hfft +array_api_tests/test_fft.py::test_ihfft +array_api_tests/test_fft.py::test_fftfreq +array_api_tests/test_fft.py::test_rfftfreq + +# test_has_names.py +## Missing Functions +array_api_tests/test_has_names.py::test_has_names[linalg-diagonal] +array_api_tests/test_has_names.py::test_has_names[linalg-matmul] +array_api_tests/test_has_names.py::test_has_names[linalg-matrix_norm] +array_api_tests/test_has_names.py::test_has_names[linalg-matrix_power] +array_api_tests/test_has_names.py::test_has_names[linalg-matrix_rank] +array_api_tests/test_has_names.py::test_has_names[linalg-matrix_transpose] +array_api_tests/test_has_names.py::test_has_names[linalg-outer] +array_api_tests/test_has_names.py::test_has_names[linalg-svdvals] +array_api_tests/test_has_names.py::test_has_names[linalg-tensordot] +array_api_tests/test_has_names.py::test_has_names[linalg-trace] +array_api_tests/test_has_names.py::test_has_names[linalg-vecdot] +array_api_tests/test_has_names.py::test_has_names[linalg-vector_norm] +array_api_tests/test_has_names.py::test_has_names[fft-hfft] +array_api_tests/test_has_names.py::test_has_names[fft-ihfft] +array_api_tests/test_has_names.py::test_has_names[elementwise-copysign] +array_api_tests/test_has_names.py::test_has_names[elementwise-hypot] +array_api_tests/test_has_names.py::test_has_names[elementwise-nextafter] +array_api_tests/test_has_names.py::test_has_names[elementwise-signbit] +array_api_tests/test_has_names.py::test_has_names[array_method-to_device] +array_api_tests/test_has_names.py::test_has_names[array_attribute-device] +array_api_tests/test_has_names.py::test_has_names[set-isin] +array_api_tests/test_has_names.py::test_has_names[set-unique_all] +array_api_tests/test_has_names.py::test_has_names[set-unique_counts] +array_api_tests/test_has_names.py::test_has_names[set-unique_inverse] +array_api_tests/test_has_names.py::test_has_names[set-unique_values] +array_api_tests/test_has_names.py::test_has_names[searching-nonzero] + +# test_signatures.py +## Missing Extensions (linalg, fft) +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.cholesky] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.det] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.diagonal] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.eigh] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.eigvalsh] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.inv] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matmul] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_norm] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_power] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_rank] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_transpose] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.outer] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.pinv] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.qr] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.slogdet] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.solve] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.svd] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.svdvals] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.tensordot] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.trace] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vecdot] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vector_norm] +array_api_tests/test_signatures.py::test_extension_func_signature[fft.hfft] +array_api_tests/test_signatures.py::test_extension_func_signature[fft.ihfft] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.cholesky] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.det] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.eigh] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.eigvalsh] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.inv] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.slogdet] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.solve] + +## Missing functions +array_api_tests/test_signatures.py::test_array_method_signature[to_device] +array_api_tests/test_signatures.py::test_func_signature[signbit] +array_api_tests/test_signatures.py::test_func_signature[nextafter] +array_api_tests/test_signatures.py::test_func_signature[hypot] +array_api_tests/test_signatures.py::test_func_signature[copysign] +array_api_tests/test_signatures.py::test_func_signature[nonzero] +array_api_tests/test_signatures.py::test_func_signature[isin] +array_api_tests/test_signatures.py::test_func_signature[unique_all] +array_api_tests/test_signatures.py::test_func_signature[unique_counts] +array_api_tests/test_signatures.py::test_func_signature[unique_inverse] +array_api_tests/test_signatures.py::test_func_signature[unique_values] + +## Missing arguments +### Missing argument `include_initial` +array_api_tests/test_signatures.py::test_func_signature[cumulative_sum] +### Missing argument `include_initial` +array_api_tests/test_signatures.py::test_func_signature[cumulative_prod] +### Missing argument `dtype` +array_api_tests/test_signatures.py::test_func_signature[prod] +### Missing argument `correction` +array_api_tests/test_signatures.py::test_func_signature[std] +### Missing argument `dtype` +array_api_tests/test_signatures.py::test_func_signature[sum] +### Missing argument `correction` +array_api_tests/test_signatures.py::test_func_signature[var] +### Missing argument `sorter` +array_api_tests/test_signatures.py::test_func_signature[searchsorted] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[asarray] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[empty_like] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[from_dlpack] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[full_like] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[ones_like] +### Missing argument `device` +array_api_tests/test_signatures.py::test_func_signature[zeros_like] +### Missing arguments `prepend` `append` +array_api_tests/test_signatures.py::test_func_signature[diff] +### rejects an integer `axes` argument +array_api_tests/test_signatures.py::test_func_signature[permute_dims] +### Missing argument `copy` +array_api_tests/test_signatures.py::test_func_signature[reshape] +### Missing arguments `descending` `stable` +array_api_tests/test_signatures.py::test_func_signature[argsort] +### Missing arguments `descending` `stable` +array_api_tests/test_signatures.py::test_func_signature[sort] +### Missing `min` argument +array_api_tests/test_signatures.py::test_func_signature[clip] + + +# test_operators_and_elementwise_functions.py +## Missing functions +array_api_tests/test_operators_and_elementwise_functions.py::test_copysign +array_api_tests/test_operators_and_elementwise_functions.py::test_hypot +array_api_tests/test_operators_and_elementwise_functions.py::test_nextafter +array_api_tests/test_operators_and_elementwise_functions.py::test_signbit +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[nextafter] +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[hypot] +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[copysign] + +## Missing complex kernels +array_api_tests/test_operators_and_elementwise_functions.py::test_acosh +array_api_tests/test_operators_and_elementwise_functions.py::test_asinh +array_api_tests/test_operators_and_elementwise_functions.py::test_atanh +array_api_tests/test_operators_and_elementwise_functions.py::test_expm1 + +## Missing arguments +### edge case +array_api_tests/test_operators_and_elementwise_functions.py::test_clip + +## Loss of precision in the complex plane +array_api_tests/test_operators_and_elementwise_functions.py::test_acos +array_api_tests/test_operators_and_elementwise_functions.py::test_asin +array_api_tests/test_operators_and_elementwise_functions.py::test_atan + +## Other failures +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[bitwise_left_shift(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_divide[divide(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__truediv__(x1, x2)] +### wraps shift counts instead of returning zero for shifts >= the integer width +array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__itruediv__(x1, x2)] +### truncates toward zero instead of flooring negative results +array_api_tests/test_operators_and_elementwise_functions.py::test_log1p +### does not accept a scalar argument +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[atan2] + +# test_inspection_functions.py +## Missing arguments +### Missing argument `device` +array_api_tests/test_inspection_functions.py::test_array_namespace_info_dtypes + +## Other failures +### Missing attribute '__array_api_version__' +array_api_tests/test_inspection_functions.py::TestInspection::test_capabilities + +# test_statistical_functions.py +## Missing arguments +### Missing argument `correction` +array_api_tests/test_statistical_functions.py::test_std +### Missing argument `correction` +array_api_tests/test_statistical_functions.py::test_var +### Missing argument `include_initial` and does not promote uint8 results to uint32 +array_api_tests/test_statistical_functions.py::test_cumulative_sum +### Missing argument `include_initial` and does not promote uint8 results to uint32 +array_api_tests/test_statistical_functions.py::test_cumulative_prod +### Missing argument `include_initial` and does not promote uint8 results to uint32 +array_api_tests/test_statistical_functions.py::test_prod +### Missing argument `include_initial` and does not promote uint8 results to uint32 +array_api_tests/test_statistical_functions.py::test_sum + +# test_manipulation_functions.py +## Missing arguments +### Missing argument `copy` +array_api_tests/test_manipulation_functions.py::test_reshape + +## Other failures +### `repeats` only accepts int +array_api_tests/test_manipulation_functions.py::test_repeat +### `source` `destination` only accepts ints +array_api_tests/test_manipulation_functions.py::test_moveaxis + +# test_data_type_functions.py +## Missing arguments +### Missing argument `copy` +array_api_tests/test_data_type_functions.py::test_astype +### +array_api_tests/test_data_type_functions.py::test_can_cast +### +array_api_tests/test_data_type_functions.py::test_finfo[dtype0] +### +array_api_tests/test_data_type_functions.py::test_finfo[dtype1] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype0] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype1] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype2] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype3] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype4] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype5] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype6] +### Missing attribute `bits` +array_api_tests/test_data_type_functions.py::test_iinfo[dtype7] + +## Other failures +### MLX result_type() does not accept Python scalar values alongside dtypes +array_api_tests/test_data_type_functions.py::TestResultType::test_with_scalars +### returns a list instead of the required tuple +array_api_tests/test_data_type_functions.py::test_broadcast_arrays +### MLX broadcast_shapes() requires at least one shape +array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_empty + + +# test_searching_functions.py +## Missing arguments +### Missing argument `sorter` +array_api_tests/test_searching_functions.py::test_searchsorted +### Missing argument `sorter` +array_api_tests/test_searching_functions.py::test_searchsorted_with_scalars + +## Other failures +### returns uint32 indices instead of the required default signed integer dtype +array_api_tests/test_searching_functions.py::test_argmax +### returns uint32 indices instead of the required default signed integer dtype +array_api_tests/test_searching_functions.py::test_argmin + +# test_utility_functions.py +## Misisng arguments +### Missing argument `append` `prepend` +array_api_tests/test_utility_functions.py::test_diff_append_prepend + +# test_creation_functions.py +## Other failures +### copy=False is not supported, https://github.com/ml-explore/mlx/pull/4036 +array_api_tests/test_creation_functions.py::test_asarray_arrays + +# test_dlpack.py +## Other failures +array_api_tests/test_dlpack.py::test_from_dlpack + +# test_sorting_functions.py +## Missing attributes +### Missing attribute `stable` +array_api_tests/test_sorting_functions.py::test_argsort +### Missing attribute `stable` +array_api_tests/test_sorting_functions.py::test_sort + +# test_special_cases.py +## Special cases +array_api_tests/test_special_cases.py::test_unary[expm1((real(x_i) is +0 or real(x_i) == -0) and imag(x_i) is +0) -> 0 + 0j] +array_api_tests/test_special_cases.py::test_unary[expm1(isfinite(real(x_i)) and imag(x_i) is +infinity) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[expm1(isfinite(real(x_i)) and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is +infinity and imag(x_i) is +0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is -infinity and imag(x_i) is +infinity) -> -1 + 0j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is +infinity and imag(x_i) is +infinity) -> infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is -infinity and imag(x_i) is NaN) -> -1 + 0j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is +infinity and imag(x_i) is NaN) -> infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is NaN and imag(x_i) is +0) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is NaN and not imag(x_i) == 0) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[expm1(real(x_i) is NaN and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[log1p((real(x_i) is +infinity or real(x_i) == -infinity) and imag(x_i) is NaN) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[log1p(real(x_i) is NaN and imag(x_i) is +infinity) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[sign(x_i is NaN) -> NaN] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is +0) -> False] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is -0) -> True] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is +infinity) -> False] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is -infinity) -> True] +array_api_tests/test_special_cases.py::test_unary[signbit(isfinite(x_i) and x_i > 0) -> False] +array_api_tests/test_special_cases.py::test_unary[signbit(isfinite(x_i) and x_i < 0) -> True] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is +NaN) -> False] +array_api_tests/test_special_cases.py::test_unary[signbit(x_i is -NaN) -> True] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is +0 and imag(x_i) is +infinity) -> 0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is +0 and imag(x_i) is NaN) -> 0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is +infinity and imag(x_i) is +0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is +infinity and imag(x_i) is +infinity) -> infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is +infinity and imag(x_i) is NaN) -> infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[sinh(real(x_i) is NaN and imag(x_i) is +0) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[sqrt(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +infinity + 0 j] +array_api_tests/test_special_cases.py::test_unary[tanh(real(x_i) is +0 and imag(x_i) is +infinity) -> +0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[tanh(real(x_i) is +0 and imag(x_i) is NaN) -> +0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[tanh(real(x_i) is +infinity and imag(x_i) is +infinity) -> 1 + 0j] +array_api_tests/test_special_cases.py::test_unary[tanh(real(x_i) is +infinity and imag(x_i) is NaN) -> 1 + 0j] +array_api_tests/test_special_cases.py::test_unary[tanh(real(x_i) is NaN and imag(x_i) is +0) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_binary[copysign(x1_i is NaN and x2_i < 0) -> NaN] +array_api_tests/test_special_cases.py::test_binary[copysign(x1_i is NaN and x2_i is -0) -> NaN] +array_api_tests/test_special_cases.py::test_binary[copysign(x1_i is NaN and x2_i is +0) -> NaN] +array_api_tests/test_special_cases.py::test_binary[copysign(x1_i is NaN and x2_i > 0) -> NaN] +array_api_tests/test_special_cases.py::test_binary[nextafter(x1_i is NaN or x2_i is NaN) -> NaN] +array_api_tests/test_special_cases.py::test_binary[nextafter(x1_i is -0 and x2_i is +0) -> +0] +array_api_tests/test_special_cases.py::test_binary[nextafter(x1_i is +0 and x2_i is -0) -> -0] +array_api_tests/test_special_cases.py::test_binary[remainder(x1_i is -0 and x2_i > 0) -> +0] +array_api_tests/test_special_cases.py::test_binary[remainder(x1_i is +0 and x2_i < 0) -> -0] +array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is -0 and x2_i > 0) -> +0] +array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is +0 and x2_i < 0) -> -0] +array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is -0 and x2_i > 0) -> +0] +array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is +0 and x2_i < 0) -> -0] +array_api_tests/test_special_cases.py::test_unary[acos((real(x_i) is +0 or real(x_i) == -0) and imag(x_i) is NaN) -> \u03c0/2 + NaN j] +array_api_tests/test_special_cases.py::test_unary[acos(isfinite(real(x_i)) and imag(x_i) is +infinity) -> \u03c0/2] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is -infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> \u03c0] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +0] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is -infinity and imag(x_i) is +infinity) -> 3\u03c0/4] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is +infinity and imag(x_i) is +infinity) -> \u03c0/4] +array_api_tests/test_special_cases.py::test_unary[acos((real(x_i) is +infinity or real(x_i) == -infinity) and imag(x_i) is NaN) -> NaN \u00b1 infinity j] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is NaN and imag(x_i) is +infinity) -> NaN] +array_api_tests/test_special_cases.py::test_unary[acosh((real(x_i) is +0 or real(x_i) == -0) and imag(x_i) is +0) -> +0 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[acosh(isfinite(real(x_i)) and imag(x_i) is +infinity) -> +infinity + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[acosh(isfinite(real(x_i)) and real(x_i) != 0 and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +0 and imag(x_i) is NaN) -> NaN \u00b1 \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is -infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +infinity + \u03c0j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is -infinity and imag(x_i) is +infinity) -> +infinity + 3\u03c0j/4] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +infinity and imag(x_i) is +infinity) -> +infinity + \u03c0j/4] +array_api_tests/test_special_cases.py::test_unary[acosh((real(x_i) is +infinity or real(x_i) == -infinity) and imag(x_i) is NaN) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is NaN and isfinite(imag(x_i))) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is NaN and imag(x_i) is +infinity) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is NaN and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is +0 and imag(x_i) is +0) -> +0 + 0j] +array_api_tests/test_special_cases.py::test_unary[asinh(isfinite(real(x_i)) and real(x_i) > 0 and imag(x_i) is +infinity) -> +infinity + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[asinh(isfinite(real(x_i)) and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is +infinity and imag(x_i) is +infinity) -> +infinity + \u03c0j/4] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is NaN and imag(x_i) is +0) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is NaN and isfinite(imag(x_i)) and imag(x_i) != 0) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is NaN and imag(x_i) is +infinity) -> \u00b1infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is NaN and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[atan(x_i is -0) -> -0] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is +0 and imag(x_i) is +0) -> +0 + 0j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is +0 and imag(x_i) is NaN) -> +0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is 1 and imag(x_i) is +0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[atanh(isfinite(real(x_i)) and real(x_i) > 0 and imag(x_i) is +infinity) -> +0 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[atanh(isfinite(real(x_i)) and real(x_i) != 0 and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +0 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is +infinity and imag(x_i) is +infinity) -> +0 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is +infinity and imag(x_i) is NaN) -> +0 + NaN j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is NaN and isfinite(imag(x_i))) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is NaN and imag(x_i) is +infinity) -> \u00b10 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is NaN and imag(x_i) is NaN) -> NaN + NaN j] +array_api_tests/test_special_cases.py::test_unary[clip(x_i is NaN) -> NaN] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is +0 and imag(x_i) is +infinity) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is +0 and imag(x_i) is NaN) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is +infinity and imag(x_i) is +0) -> +infinity + 0j] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is +infinity and imag(x_i) is +infinity) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is +infinity and imag(x_i) is NaN) -> +infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[cosh(real(x_i) is NaN and (imag(x_i) is +0 or imag(x_i) == -0)) -> NaN + 0j] +array_api_tests/test_special_cases.py::test_unary[acos((real(x_i) is +infinity or real(x_i) == -infinity) and imag(x_i) is NaN) -> NaN \xb1 infinity j] +array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +0 and imag(x_i) is NaN) -> NaN \xb1 \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is +infinity and imag(x_i) is +infinity) -> \u03c0/4 - infinity j] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is -infinity and imag(x_i) is +infinity) -> 3\u03c0/4 - infinity j] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is NaN and imag(x_i) is +infinity) -> NaN - infinity j] +array_api_tests/test_special_cases.py::test_unary[asinh(real(x_i) is NaN and imag(x_i) is +infinity) -> \xb1infinity + NaN j] +array_api_tests/test_special_cases.py::test_unary[atanh(real(x_i) is NaN and imag(x_i) is +infinity) -> \xb10 + \u03c0j/2] +array_api_tests/test_special_cases.py::test_unary[acos(isfinite(real(x_i)) and imag(x_i) is +infinity) -> \u03c0/2 - infinity j] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is -infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> \u03c0 - infinity j] +array_api_tests/test_special_cases.py::test_unary[acos(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +0 - infinity j] +array_api_tests/test_special_cases.py::test_unary[sign((real(x_i) is -0 or real(x_i) == +0) and (imag(x_i) is -0 or imag(x_i) == +0)) -> 0 + 0j] +array_api_tests/test_special_cases.py::test_unary[ceil(x_i.is_integer()) -> x_i] +array_api_tests/test_special_cases.py::test_unary[floor(x_i.is_integer()) -> x_i] +array_api_tests/test_special_cases.py::test_unary[round(x_i.is_integer()) -> x_i] +array_api_tests/test_special_cases.py::test_unary[trunc(x_i.is_integer()) -> x_i] diff --git a/Source/Cmlx/mlx/setup.py b/Source/Cmlx/mlx/setup.py index ad258d8af..72cf59668 100644 --- a/Source/Cmlx/mlx/setup.py +++ b/Source/Cmlx/mlx/setup.py @@ -25,13 +25,13 @@ def cuda_toolkit_major_version(): def get_version(): with open("mlx/version.h", "r") as fid: - for l in fid: - if "#define MLX_VERSION_MAJOR" in l: - major = l.split()[-1] - if "#define MLX_VERSION_MINOR" in l: - minor = l.split()[-1] - if "#define MLX_VERSION_PATCH" in l: - patch = l.split()[-1] + for line in fid: + if "#define MLX_VERSION_MAJOR" in line: + major = line.split()[-1] + if "#define MLX_VERSION_MINOR" in line: + minor = line.split()[-1] + if "#define MLX_VERSION_PATCH" in line: + patch = line.split()[-1] version = f"{major}.{minor}.{patch}" pypi_release = int(os.environ.get("PYPI_RELEASE", 0)) dev_release = int(os.environ.get("DEV_RELEASE", 0)) @@ -57,7 +57,8 @@ def get_version(): # # Frontend package: # - Triggered with `MLX_BUILD_FRONTEND_PACKAGE=1` -# - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) +# - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, +# etc) # - Wheel has Python ABI and platform tags # - Wheel should be built for the cross-product of python version and platforms # - Package name is "mlx" and it depends on backend packages (e.g. mlx-metal, mlx-cuda) diff --git a/Source/Cmlx/mlx/tests/allocator_tests.cpp b/Source/Cmlx/mlx/tests/allocator_tests.cpp index d528f8ea2..1fff8d8cc 100644 --- a/Source/Cmlx/mlx/tests/allocator_tests.cpp +++ b/Source/Cmlx/mlx/tests/allocator_tests.cpp @@ -20,33 +20,33 @@ using namespace mlx::core; TEST_CASE("test simple allocations") { { - auto buffer = allocator::malloc(sizeof(float)); + auto buffer = allocator::allocator().malloc(sizeof(float)); auto fptr = static_cast(buffer.raw_ptr()); *fptr = 0.5f; CHECK_EQ(*fptr, 0.5f); - allocator::free(buffer); + allocator::allocator().free(buffer); } { - auto buffer = allocator::malloc(128 * sizeof(int)); + auto buffer = allocator::allocator().malloc(128 * sizeof(int)); int* ptr = static_cast(buffer.raw_ptr()); for (int i = 0; i < 128; ++i) { ptr[i] = i; } - allocator::free(buffer); + allocator::allocator().free(buffer); } { - auto buffer = allocator::malloc(0); - allocator::free(buffer); + auto buffer = allocator::allocator().malloc(0); + allocator::allocator().free(buffer); } } TEST_CASE("test large allocations") { size_t size = 1 << 30; for (int i = 0; i < 100; ++i) { - auto buffer = allocator::malloc(size); - allocator::free(buffer); + auto buffer = allocator::allocator().malloc(size); + allocator::allocator().free(buffer); } } @@ -54,14 +54,14 @@ TEST_CASE("test cached allocation keeps capacity") { auto old_limit = set_cache_limit(1 << 20); clear_cache(); - auto large = allocator::malloc(8192); - allocator::free(large); + auto large = allocator::allocator().malloc(8192); + allocator::allocator().free(large); auto cached = get_cache_memory(); CHECK_GE(cached, 8192); - auto small = allocator::malloc(6000); + auto small = allocator::allocator().malloc(6000); CHECK_GE(allocator::allocator().size(small), cached); - allocator::free(small); + allocator::allocator().free(small); CHECK_GE(get_cache_memory(), cached); clear_cache(); @@ -94,8 +94,8 @@ TEST_CASE("test clear cache synchronizes cpu streams") { auto old_limit = set_cache_limit(1 << 20); clear_cache(); - auto cached = allocator::malloc(8192); - allocator::free(cached); + auto cached = allocator::allocator().malloc(8192); + allocator::allocator().free(cached); CHECK_GE(get_cache_memory(), 8192); auto task_started = std::make_shared>(); diff --git a/Source/Cmlx/mlx/tests/array_tests.cpp b/Source/Cmlx/mlx/tests/array_tests.cpp index cc0cf5239..7adad04a2 100644 --- a/Source/Cmlx/mlx/tests/array_tests.cpp +++ b/Source/Cmlx/mlx/tests/array_tests.cpp @@ -581,7 +581,8 @@ TEST_CASE("test array shared buffer") { Shape shape = {2, 2}; auto n_elem = shape[0] * shape[1]; - allocator::Buffer buf_b = allocator::malloc(n_elem * sizeof(float)); + allocator::Buffer buf_b = + allocator::allocator().malloc(n_elem * sizeof(float)); void* buf_b_ptr = buf_b.raw_ptr(); float* float_buf_b = (float*)buf_b_ptr; @@ -594,11 +595,11 @@ TEST_CASE("test array shared buffer") { auto deleter = [float_buf_b](allocator::Buffer buf) { CHECK_EQ(float_buf_b, (float*)buf.raw_ptr()); CHECK_EQ(float_buf_b[0], ((float*)buf.raw_ptr())[0]); - allocator::free(buf); + allocator::allocator().free(buf); }; array a = ones(shape, float32); - array b = array(buf_b, shape, float32, deleter); + array b = array(allocator::Data(buf_b, deleter), shape, float32); eval(a + b); } @@ -647,6 +648,19 @@ TEST_CASE("test make array from user buffer") { CHECK_EQ(count, 1); } +TEST_CASE("test user buffer is not donated") { + int size = 4096; + std::vector buffer(size, 1.0f); + auto out = exp(array(buffer.data(), Shape{size}, float32, [](void*) {})); + eval(out); + CHECK_EQ(buffer[0], 1.0f); +} + +TEST_CASE("test null deleter is not owned") { + allocator::Data data(allocator::Buffer(nullptr), nullptr); + CHECK_FALSE(data.is_owned()); +} + TEST_CASE("test negative indexing for shape/strides") { // 2D array: shape = {2, 3} std::vector data(6, 1.0f); diff --git a/Source/Cmlx/mlx/tests/cuda_tests.cpp b/Source/Cmlx/mlx/tests/cuda_tests.cpp index aca39ca88..9afa980cd 100644 --- a/Source/Cmlx/mlx/tests/cuda_tests.cpp +++ b/Source/Cmlx/mlx/tests/cuda_tests.cpp @@ -109,3 +109,13 @@ TEST_CASE("test clear cache trims CUDA pool") { cudaSuccess); CHECK_LT(final_reserved, allocated_reserved); } + +TEST_CASE("test eval after cuda error") { + auto s = default_stream(Device::gpu); + // A failed allocation sets a CUDA error on the calling thread. + CHECK_THROWS(eval(zeros({1 << 20, 1 << 20}, float32, s))); + // The error must not resurface in later evaluations. + auto a = ones({4}, float32, s); + CHECK_NOTHROW(eval(a)); + CHECK(array_equal(a, ones({4})).item()); +} diff --git a/Source/Cmlx/mlx/tests/fft_tests.cpp b/Source/Cmlx/mlx/tests/fft_tests.cpp index 6adc4c1d4..fd7007852 100644 --- a/Source/Cmlx/mlx/tests/fft_tests.cpp +++ b/Source/Cmlx/mlx/tests/fft_tests.cpp @@ -396,3 +396,69 @@ TEST_CASE("test fftshift and ifftshift") { CHECK_THROWS_AS(fft::ifftshift(x, {3}), std::out_of_range); CHECK_THROWS_AS(fft::ifftshift(x, {-5}), std::out_of_range); } + +TEST_CASE("test stft and istft") { + auto win = ones({16}, float32); + + // Real, centered input -> (n_freq, n_frames). + auto x = astype(arange(64), float32); + auto z = fft::stft(x, 16, 8, 16, win, /* center = */ true); + CHECK_EQ(z.ndim(), 2); + CHECK_EQ(z.shape(0), 9); // n_fft / 2 + 1 + CHECK_EQ(z.shape(1), 9); // 1 + (64 + 16 - 16) / 8 + CHECK_EQ(z.dtype(), complex64); + + // Batched input -> (batch, n_freq, n_frames). + auto xb = reshape(astype(arange(2 * 64), float32), {2, 64}); + auto zb = fft::stft(xb, 16, 8, 16, win, true); + CHECK_EQ(zb.ndim(), 3); + CHECK_EQ(zb.shape(0), 2); + CHECK_EQ(zb.shape(1), 9); + + // Non-overlapping rectangular window reconstructs the signal exactly. + auto x2 = astype(arange(32), float32); + auto win2 = ones({8}, float32); + auto z2 = fft::stft(x2, 8, 8, 8, win2, /* center = */ false); + auto y2 = fft::istft( + z2, + /* n_fft = */ 8, + /* hop_length = */ 8, + /* win_length = */ 8, + win2, + /* center = */ false, + fft::FFTNorm::Backward, + /* onesided = */ true, + /* length = */ 32); + CHECK_EQ(y2.shape(0), 32); + CHECK(allclose(y2, x2, 1e-4, 1e-4).item()); + + // Overlapping window: round-trip recovers the fully covered interior. + auto x3 = astype(arange(64), float32); + auto z3 = fft::stft(x3, 16, 8, 16, win, /* center = */ false); + auto y3 = fft::istft(z3, 16, 8, 16, win, false, fft::FFTNorm::Backward, true); + CHECK_EQ(y3.shape(0), 64); // n_fft + (n_frames - 1) * hop = 16 + 6 * 8 + CHECK(allclose(slice(y3, {8}, {56}), slice(x3, {8}, {56}), 1e-4, 1e-4) + .item()); + + // n_fft not divisible by hop: length is n_fft + (n_frames - 1) * hop. + auto w7 = ones({7}, float32); + auto z4 = fft::stft(x3, 7, 3, 7, w7, /* center = */ false); + auto y4 = fft::istft(z4, 7, 3, 7, w7, false, fft::FFTNorm::Backward, true); + CHECK_EQ(y4.shape(0), 64); // n_frames = 20 -> 7 + 19 * 3 + + // Two-sided transform keeps all frequency bins. + auto zf = fft::stft( + x, 16, 8, 16, win, true, "reflect", fft::FFTNorm::Backward, false); + CHECK_EQ(zf.shape(0), 16); + + // Edge padding. + auto ze = fft::stft(x, 16, 8, 16, win, true, "edge"); + CHECK_EQ(ze.shape(0), 9); + + // Error cases. + CHECK_THROWS_AS(fft::stft(x, 0), std::invalid_argument); + CHECK_THROWS_AS( + fft::stft(x, 16, 8, 16, ones({32}, float32)), std::invalid_argument); + CHECK_THROWS_AS( + fft::stft(x, 16, 8, 16, win, true, "bogus"), std::invalid_argument); +} diff --git a/Source/Cmlx/mlx/tests/gpu_tests.cpp b/Source/Cmlx/mlx/tests/gpu_tests.cpp index 8f492e52f..a49f9a704 100644 --- a/Source/Cmlx/mlx/tests/gpu_tests.cpp +++ b/Source/Cmlx/mlx/tests/gpu_tests.cpp @@ -465,6 +465,59 @@ TEST_CASE("test gpu matmul") { } } +TEST_CASE("test gpu gather/scatter indices") { + // Run with METAL_DEVICE_WRAPPER_TYPE=1 to validate the metadata bindings. + + SUBCASE("scatter scalar indices") { + auto x = array({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, {3, 4}); + + SUBCASE("non-leading axis") { + auto out = take(x, array(2), 1, Device::gpu); + CHECK(array_equal(out, array({2, 6, 10}), Device::cpu).item()); + } + + SUBCASE("strided source") { + auto src = transpose(x, Device::gpu); + auto out = take(src, array(-1), 0, Device::gpu); + CHECK(array_equal(out, array({3, 7, 11}), Device::cpu).item()); + } + + SUBCASE("multiple scalar indices") { + auto out = gather(x, {array(1), array(-2)}, {0, 1}, {1, 1}, Device::gpu); + CHECK(array_equal(out, array({6}, {1, 1}), Device::cpu).item()); + } + } + + SUBCASE("gather without indices") { + SUBCASE("scalar source") { + auto out = gather(array(7), {}, std::vector{}, {}, Device::gpu); + CHECK_EQ(out.item(), 7); + } + + SUBCASE("matrix source") { + auto x = array({0, 1, 2, 3}, {2, 2}); + auto out = gather(x, {}, std::vector{}, {2, 2}, Device::gpu); + CHECK(array_equal(out, x, Device::cpu).item()); + } + } + + SUBCASE("scatter without indices") { + SUBCASE("scalar update") { + auto out = + scatter_max(array(1), {}, array(2), std::vector{}, Device::gpu); + CHECK_EQ(out.item(), 2); + } + + SUBCASE("strided update") { + auto x = array({0, 1, 2, 3}, {2, 2}); + auto updates = transpose(array({1, 2, 3, 4}, {2, 2}), Device::gpu); + auto out = scatter_add(x, {}, updates, std::vector{}, Device::gpu); + CHECK(array_equal(out, array({1, 4, 4, 7}, {2, 2}), Device::cpu) + .item()); + } + } +} + TEST_CASE("test gpu validation") { // Run this test with Metal validation enabled // METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./tests/tests \ diff --git a/Source/Cmlx/mlx/tests/ops_tests.cpp b/Source/Cmlx/mlx/tests/ops_tests.cpp index 9c95e9b5e..cda9523f3 100644 --- a/Source/Cmlx/mlx/tests/ops_tests.cpp +++ b/Source/Cmlx/mlx/tests/ops_tests.cpp @@ -9,6 +9,7 @@ #include "mlx/backend/cuda/cuda.h" #include "mlx/mlx.h" +#include "mlx/primitives.h" using namespace mlx::core; @@ -124,8 +125,9 @@ TEST_CASE("test reshape") { CHECK(x.flags().col_contiguous); y = reshape(x, {2, 5, 10}); eval(y); + // Strides (5, 1, 10) are neither row nor col major CHECK(!y.flags().row_contiguous); - CHECK(y.flags().col_contiguous); + CHECK(!y.flags().col_contiguous); y = reshape(x, {2, 50}); eval(y); CHECK(y.flags().row_contiguous); @@ -358,6 +360,22 @@ TEST_CASE("test slice") { out = slice(x, {3, 0}, {-5, 2}, {-1, 1}); eval(out); CHECK_EQ(out.data_size(), 8); + + x = as_strided(square(arange(16, float32)), {3, 2, 2}, {1, 1, 5}, 0); + out = slice(x, {0, 0, 0}, {2, 2, 2}); + eval(out); + CHECK_EQ(out.data_size(), 8); + CHECK_FALSE(out.flags().contiguous); + CHECK_EQ(sum(out).item(), 152.0f); + + x = reshape(arange(6), {2, 3}); + x = slice(x, {1, 0}, {-3, 3}, {-1, 1}); + out = slice(x, {1, 0}, {-3, 3}, {-1, 1}); + eval(out); + CHECK(out.flags().contiguous); + CHECK(out.flags().row_contiguous); + CHECK(array_equal(reshape(transpose(out), {-1}), array({0, 3, 1, 4, 2, 5})) + .item()); } TEST_CASE("test slice update") { @@ -606,6 +624,71 @@ TEST_CASE("test split") { CHECK(array_equal(out[1], array({})).item()); CHECK(array_equal(out[2], array({1})).item()); CHECK(array_equal(out[3], array({2, 3, 4})).item()); + + // data_size is the distance between the first and the last element, so a + // split that leaves holes reports more than the output has elements. + x = ones({4, 4}); + out = split(x, 2, 1); + eval(out); + CHECK_EQ(out[0].data_size(), 14); + CHECK_EQ(out[1].data_size(), 14); + + // Splitting the first axis leaves no holes. + out = split(x, 2, 0); + eval(out); + CHECK_EQ(out[0].data_size(), 8); + CHECK_EQ(out[1].data_size(), 8); + + x = ones({2, 3, 4}); + out = split(x, Shape{2}, 2); + eval(out); + CHECK_EQ(out[0].data_size(), 22); + CHECK_EQ(out[1].data_size(), 22); + + // A split output and the matching slice are the same view. + x = ones({1, 2, 8, 12}); + out = split(x, Shape{8}, 3); + auto sliced = slice(x, {0, 0, 0, 8}, {1, 2, 8, 12}); + eval(out); + eval(sliced); + CHECK_EQ(out[1].data_size(), sliced.data_size()); + CHECK_EQ(out[1].data_size(), 184); + + // (4, 4) array, but slice keeps the strides at (8, 1) + x = slice(reshape(arange(32), {4, 8}), {0, 0}, {4, 4}); + // out[0] and out[1] are (4, 2) + out = split(x, Shape{2}, 1); + // same out out[1] + auto holey = slice(x, {0, 2}, {4, 4}); + eval(out); + eval(holey); + CHECK_EQ(out[0].data_size(), 26); + CHECK_EQ(out[1].data_size(), 26); + CHECK_EQ(out[1].data_size(), holey.data_size()); + CHECK_EQ(out[0].strides(), x.strides()); + CHECK_EQ(out[1].strides(), holey.strides()); + CHECK(array_equal(out[1], holey).item()); + + // A negative stride travels backwards, so it bounds the span from below. + // strides (-4, 1) + x = flip(reshape(arange(24), {6, 4}), 0); + out = split(x, Shape{2}, 0); + eval(out); + // no holes + CHECK_EQ(out[0].data_size(), 8); + CHECK_EQ(out[1].data_size(), 16); + CHECK_EQ(out[0].strides(), x.strides()); + CHECK_EQ(out[1].strides(), x.strides()); + CHECK(array_equal(out[0], slice(x, {0, 0}, {2, 4})).item()); + CHECK(array_equal(out[1], slice(x, {2, 0}, {6, 4})).item()); + + // A zero-length axis empties every output. Splitting at 0 cannot get here, + // split() sends that to slice(). + x = ones({0, 4}); + out = split(x, Shape{2}, 1); + eval(out); + CHECK_EQ(out[0].data_size(), 0); + CHECK_EQ(out[1].data_size(), 0); } TEST_CASE("test flip") { @@ -896,6 +979,29 @@ TEST_CASE("test comparison ops") { } } +TEST_CASE("test float16 sigmoid mixed lanes") { + const float values[] = {-12, -4, -1, 0, 1, 4, 12, -2}; + for (int size : {1, 7, 8, 9, 15, 16, 17}) { + for (int offset = 0; offset < 8; ++offset) { + std::vector input(size); + for (int i = 0; i < size; ++i) { + input[i] = values[(i + offset) % 8]; + } + auto result = sigmoid(array(input.begin(), {size}, float16), Device::cpu); + result.eval(); + for (int i = 0; i < size; ++i) { + CAPTURE(size); + CAPTURE(offset); + CAPTURE(i); + auto expected = 1.0 / (1.0 + std::exp(-double(input[i]))); + CHECK_EQ( + static_cast(result.data()[i]), + doctest::Approx(expected).epsilon(0.002)); + } + } + } +} + TEST_CASE("test is nan") { array x(1.0f); CHECK_FALSE(isnan(x).item()); @@ -2006,9 +2112,10 @@ TEST_CASE("test arithmetic binary ops") { // Integer division by zero gives quotient 0 and remainder a. if (default_device() == Device::cpu) { for (auto dt : {int8, int16, int32, int64, uint8, uint16, uint32, uint64}) { - auto num = astype(array({7, 0, 5}, {3}), dt); - auto den = zeros({3}, dt); - CHECK(array_equal(floor_divide(num, den), zeros({3}, dt)).item()); + auto num = astype(array({7, 0, 5, -7}, {4}), dt); + auto den = zeros({4}, dt); + auto q = floor_divide(num, den); + CHECK(array_equal(q, zeros({4}, dt)).item()); CHECK(array_equal(remainder(num, den), num).item()); } } @@ -2893,6 +3000,13 @@ TEST_CASE("test as_strided op") { CHECK_EQ(y.data_size(), 10); CHECK_FALSE(y.flags().contiguous); + auto xf = square(arange(16, float32)); + y = as_strided(xf, {2, 2, 2}, {1, 1, 5}, 0); + eval(y); + CHECK_EQ(y.data_size(), 8); + CHECK_FALSE(y.flags().contiguous); + CHECK_EQ(sum(y).item(), 152.0f); + x = reshape(x, {2, 5}); // 0 1 2 3 ... x = transpose(x, {1, 0}); // 0 5 1 6 2 7 ... y = as_strided(x, {3, 3}, {2, 1}, 1); @@ -3631,6 +3745,25 @@ TEST_CASE("inner") { CHECK(array_equal(z, expected).item()); } +TEST_CASE("test floor divide") { + // Signed integers round the quotient towards minus infinity. + auto x = array({4, 5, -1, -6}); + auto y = array({-2, 3, 2, -3}); + CHECK(array_equal(floor_divide(x, y), array({-2, 1, -1, 2})).item()); + + // Unsigned integers use the plain quotient. + auto ux = astype(array({7, 0, 5}, {3}), uint32); + auto uy = astype(array({2, 2, 3}, {3}), uint32); + auto expected = astype(array({3, 0, 1}, {3}), uint32); + CHECK(array_equal(floor_divide(ux, uy), expected).item()); + + // The dtype is kept and the result is one primitive. + auto out = floor_divide(astype(x, int8), astype(array(3), int8)); + CHECK_EQ(out.dtype(), int8); + CHECK_EQ(std::string(out.primitive().name()), "FloorDivide"); + CHECK(out.siblings().empty()); +} + TEST_CASE("test divmod") { auto x = array({1, 2, 3}); auto y = array({1, 1, 1}); diff --git a/Source/Cmlx/mlx/tests/residency_tests.cpp b/Source/Cmlx/mlx/tests/residency_tests.cpp index 60cd7c90e..0edac9b6d 100644 --- a/Source/Cmlx/mlx/tests/residency_tests.cpp +++ b/Source/Cmlx/mlx/tests/residency_tests.cpp @@ -66,11 +66,11 @@ TEST_CASE("test residency set wires nothing when the wired limit is zero") { std::vector bufs; for (int i = 0; i < 4; ++i) { - bufs.push_back(allocator::malloc(4 * MB)); + bufs.push_back(allocator::allocator().malloc(4 * MB)); CHECK_EQ(residency().wired_size(), 0); } for (auto buf : bufs) { - allocator::free(buf); + allocator::allocator().free(buf); } CHECK_EQ(residency().wired_size(), 0); } @@ -89,14 +89,14 @@ TEST_CASE("test residency set never wires more than the wired 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)); + bufs.push_back(allocator::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); + allocator::allocator().free(buf); } // Everything of ours is unwired again, exactly. CHECK_EQ(residency().wired_size(), baseline); @@ -108,13 +108,13 @@ TEST_CASE("test raising the wired limit wires already-allocated buffers") { } LimitGuard guard(0, 0); - auto buf = allocator::malloc(8 * MB); + auto buf = allocator::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); + allocator::allocator().free(buf); } TEST_CASE("test lowering the wired limit unwires buffers") { @@ -124,7 +124,7 @@ TEST_CASE("test lowering the wired limit unwires buffers") { LimitGuard guard(64 * MB, 0); const size_t baseline_with_budget = residency().wired_size(); - auto buf = allocator::malloc(8 * MB); + auto buf = allocator::allocator().malloc(8 * MB); CHECK_GE(residency().wired_size(), baseline_with_budget + 8 * MB); set_wired_limit(0); @@ -134,7 +134,7 @@ TEST_CASE("test lowering the wired limit unwires buffers") { set_wired_limit(64 * MB); CHECK_GE(residency().wired_size(), baseline_with_budget + 8 * MB); - allocator::free(buf); + allocator::allocator().free(buf); } TEST_CASE("test residency accounting is exact across free/realloc cycles") { @@ -147,9 +147,9 @@ TEST_CASE("test residency accounting is exact across free/realloc cycles") { const size_t baseline = residency().wired_size(); for (int i = 0; i < 32; ++i) { - auto buf = allocator::malloc(8 * MB); + auto buf = allocator::allocator().malloc(8 * MB); CHECK_GE(residency().wired_size(), baseline + 8 * MB); - allocator::free(buf); + allocator::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); @@ -170,7 +170,7 @@ TEST_CASE("test wired allocations are spread across size-capped 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)); + bufs.push_back(allocator::allocator().malloc(5 * MB)); } CHECK_EQ(residency().wired_size(), baseline + 8 * 5 * MB); CHECK_GT(residency().num_sets(), baseline_sets); @@ -183,7 +183,7 @@ TEST_CASE("test wired allocations are spread across size-capped sets") { const size_t sets_before = residency().num_sets(); for (auto buf : bufs) { - allocator::free(buf); + allocator::allocator().free(buf); } CHECK_EQ(residency().wired_size(), baseline); @@ -192,11 +192,11 @@ TEST_CASE("test wired allocations are spread across size-capped sets") { for (int cycle = 0; cycle < 4; ++cycle) { std::vector again; for (int i = 0; i < 8; ++i) { - again.push_back(allocator::malloc(5 * MB)); + again.push_back(allocator::allocator().malloc(5 * MB)); } CHECK_EQ(residency().num_sets(), sets_before); for (auto buf : again) { - allocator::free(buf); + allocator::allocator().free(buf); } } CHECK_EQ(residency().wired_size(), baseline); @@ -215,7 +215,7 @@ TEST_CASE("test the set count stays within the command queue limit") { std::vector bufs; for (int i = 0; i < 64; ++i) { - bufs.push_back(allocator::malloc(2 * MB)); + bufs.push_back(allocator::allocator().malloc(2 * MB)); } CHECK_LE(residency().num_sets(), 32); CHECK_EQ(residency().wired_size(), baseline + 64 * 2 * MB); @@ -224,7 +224,7 @@ TEST_CASE("test the set count stays within the command queue limit") { check_gpu_work(); for (auto buf : bufs) { - allocator::free(buf); + allocator::allocator().free(buf); } CHECK_EQ(residency().wired_size(), baseline); } @@ -238,11 +238,11 @@ TEST_CASE("test an allocation larger than the set cap is still wired") { set_wired_limit(128 * MB); const size_t baseline = residency().wired_size(); - auto big = allocator::malloc(32 * MB); + auto big = allocator::allocator().malloc(32 * MB); CHECK_EQ(residency().wired_size(), baseline + 32 * MB); check_gpu_work(); - allocator::free(big); + allocator::allocator().free(big); CHECK_EQ(residency().wired_size(), baseline); } @@ -258,13 +258,13 @@ TEST_CASE("test a set cap of zero keeps everything in one set") { std::vector bufs; for (int i = 0; i < 16; ++i) { - bufs.push_back(allocator::malloc(4 * MB)); + bufs.push_back(allocator::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); + allocator::allocator().free(buf); } CHECK_EQ(residency().wired_size(), baseline); } @@ -280,17 +280,17 @@ TEST_CASE("test the wired limit is used up to its boundary") { const size_t baseline = residency().wired_size(); set_wired_limit(baseline + 4 * MB); - auto first = allocator::malloc(4 * MB); + auto first = allocator::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); + auto second = allocator::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); + allocator::allocator().free(first); CHECK_EQ(residency().wired_size(), baseline); - allocator::free(second); + allocator::allocator().free(second); CHECK_EQ(residency().wired_size(), baseline); } diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift index 34ded3f97..046260049 100644 --- a/Source/Encuda/encuda-compile.swift +++ b/Source/Encuda/encuda-compile.swift @@ -82,12 +82,37 @@ extension Encuda { guard status == 0 else { throw EncudaError.nvccFailed(status) } + writeStamp() #endif } + /// Inputs to the generated `.cpp` that are not themselves input *files*. + /// A modification-time comparison cannot see any of these change, so the + /// signature is stamped beside the output and compared on the next run. + private var configurationSignature: String { + [ + "nvcc=\(nvccPath ?? "")", + "ccbin=\(clangppPath ?? "")", + "std=\(std ?? "")", + "arch=\(ProcessInfo.processInfo.environment["CUDA_ARCH"] ?? "")", + "include=\(includeDirs.joined(separator: ":"))", + ].joined(separator: "\n") + } + + private var stampURL: URL { + URL(fileURLWithPath: output + ".encuda-stamp") + } + + private func writeStamp() { + try? configurationSignature.write(to: stampURL, atomically: true, encoding: .utf8) + } + private func isUpToDate() -> Bool { let fm = FileManager.default let outputURL = URL(fileURLWithPath: output) + guard let stamp = try? String(contentsOf: stampURL, encoding: .utf8), + stamp == configurationSignature + else { return false } guard fm.fileExists(atPath: output), let outputMod = (try? outputURL.resourceValues(forKeys: [.contentModificationDateKey]))? diff --git a/Source/MLX/Distributed.swift b/Source/MLX/Distributed.swift index 87d8f7659..47947636b 100644 --- a/Source/MLX/Distributed.swift +++ b/Source/MLX/Distributed.swift @@ -76,12 +76,16 @@ public enum MLXDistributed { /// The rank of this process in the group. public var rank: Int { - Int(mlx_distributed_group_rank(ctx)) + var result: Int32 = 0 + mlx_distributed_group_get_rank(&result, ctx) + return Int(result) } /// The number of processes in the group. public var size: Int { - Int(mlx_distributed_group_size(ctx)) + var result: Int32 = 0 + mlx_distributed_group_get_size(&result, ctx) + return Int(result) } /// Split the group into subgroups based on the provided color. diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md index 46b5ae120..c8e08c226 100644 --- a/Source/MLX/Documentation.docc/MLX.md +++ b/Source/MLX/Documentation.docc/MLX.md @@ -143,3 +143,7 @@ boundaries, take an evaluated, immutable snapshot. ### Distributed - ``MLXDistributed`` + +### Load Progress + +- ``LoadProgress`` diff --git a/Source/MLX/Documentation.docc/Organization/arithmetic.md b/Source/MLX/Documentation.docc/Organization/arithmetic.md index 2a5851293..bc9203bda 100644 --- a/Source/MLX/Documentation.docc/Organization/arithmetic.md +++ b/Source/MLX/Documentation.docc/Organization/arithmetic.md @@ -215,7 +215,7 @@ Note: the `-` and `/` operators are not able to be linked here. - ``blockMaskedMM(_:_:blockSize:maskOut:maskLHS:maskRHS:stream:)`` - ``addMM(_:_:_:alpha:beta:stream:)`` - ``quantizedMM(_:_:scales:biases:transpose:groupSize:bits:mode:stream:)`` -- ``gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:sortedIndices:stream:)`` +- ``gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:globalScale:sortedIndices:stream:)`` - ``quantizedQuantizedMM(_:_:scales:groupSize:bits:mode:globalScaleX:globalScaleW:stream:)`` - ``segmentedMM(_:_:segments:stream:)`` - ``inner(_:_:stream:)`` diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index e0afc1fe5..5fb01d0e9 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -120,9 +120,13 @@ operations as methods for convenience. - ``loadArray(url:stream:)`` - ``loadArrays(url:stream:)`` +- ``loadArrays(url:stream:progressHandler:)`` - ``loadArraysAndMetadata(url:stream:)`` +- ``loadArraysAndMetadata(url:stream:progressHandler:)`` - ``save(array:url:stream:)`` - ``save(arrays:metadata:url:stream:)`` +- ``withLoadProgressHandler(_:_:)-3ghip`` +- ``withLoadProgressHandler(_:_:)-8cm31`` ### Logical diff --git a/Source/MLX/IO+Concurrent.swift b/Source/MLX/IO+Concurrent.swift new file mode 100644 index 000000000..e3b2a7206 --- /dev/null +++ b/Source/MLX/IO+Concurrent.swift @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +import CoreFoundation +import Foundation + +// MARK: - Concurrent safetensors loading + +// Safetensors payload reads are lazy. Submit tensors in reverse file-offset order +// because MLX builds its evaluation tape in input order and consumes it from the back. +// Each Load primitive queues its read on MLX's I/O pool, which owns read concurrency; +// one checkedEval schedules the complete checkpoint and waits for all reads to finish. + +/// Match the limit enforced by MLX and the safetensors reference implementation. +private let maximumSafetensorHeaderBytes: UInt64 = 100_000_000 + +/// Decode a JSON integer without accepting booleans, floating-point values, or values +/// outside the supported signed 64-bit offset range. +private func safetensorOffset(_ value: Any) -> Int64? { + guard let number = value as? NSNumber else { return nil } + let numberType = String(cString: number.objCType) + guard + CFGetTypeID(number) != CFBooleanGetTypeID(), + numberType != "f", + numberType != "d", + numberType != "D" + else { + return nil + } + let offset = number.int64Value + return offset >= 0 ? offset : nil +} + +/// The tensor names in the safetensors file at `url`, ordered by position in the file. +/// +/// Reads the 8-byte header length and the JSON header only. Throws when the file is not a +/// well-formed safetensors file; callers fall back to loading the file whole. +func safetensorNamesInFileOrder(url: URL) throws -> [String] { + struct Malformed: Error {} + + let handle = try FileHandle(forReadingFrom: url) + defer { try? handle.close() } + + guard let lengthData = try handle.read(upToCount: 8), lengthData.count == 8 else { + throw Malformed() + } + let headerLength = lengthData.withUnsafeBytes { $0.loadUnaligned(as: UInt64.self) } + .littleEndian + guard headerLength > 0, headerLength < maximumSafetensorHeaderBytes else { + throw Malformed() + } + guard let headerData = try handle.read(upToCount: Int(headerLength)), + headerData.count == headerLength, + let header = try JSONSerialization.jsonObject(with: headerData) as? [String: Any] + else { + throw Malformed() + } + + var spans = [(name: String, begin: Int64, end: Int64)]() + for (name, value) in header { + guard name != "__metadata__" else { continue } + guard let entry = value as? [String: Any], + let offsets = entry["data_offsets"] as? [Any], offsets.count == 2, + let begin = safetensorOffset(offsets[0]), + let end = safetensorOffset(offsets[1]), + end >= begin + else { + throw Malformed() + } + spans.append((name, begin, end)) + } + // Safetensors requires the data buffer to be covered contiguously. Sorting by end as + // well handles zero-byte tensors that share the following tensor's begin offset. + spans.sort { + $0.begin == $1.begin ? $0.end < $1.end : $0.begin < $1.begin + } + var expectedBegin: Int64 = 0 + for span in spans { + guard span.begin == expectedBegin else { throw Malformed() } + expectedBegin = span.end + } + return spans.map(\.name) +} + +/// Roots in the order MLX must receive them to schedule reads in file-offset order. +func safetensorArraysInEvaluationOrder( + arrays: [[String: MLXArray]], names: [[String]?] +) -> [MLXArray] { + var orderedArrays = [MLXArray]() + for (fileArrays, fileNames) in zip(arrays, names) { + if let fileNames, + fileNames.count == fileArrays.count, + fileNames.allSatisfy({ fileArrays[$0] != nil }) + { + orderedArrays.append(contentsOf: fileNames.compactMap { fileArrays[$0] }) + } else { + // A planning failure or disagreement with MLX's header read must not leave + // any arrays out of the completion barrier. Their read order is unspecified. + orderedArrays.append(contentsOf: fileArrays.values) + } + } + // Reverse the whole checkpoint so MLX visits files in URL order and tensors within + // each file in ascending offset order. Keep duplicate names from different files. + return orderedArrays.reversed() +} + +/// Lock-guarded results of the parallel, header-only preparation pass. +private final class ConcurrentLoadPreparation: @unchecked Sendable { + private let lock = NSLock() + private var perFileArrays: [[String: MLXArray]] + private var perFileMetadata: [[String: String]] + private var namesPerFile: [[String]?] + private var firstError: Error? + + init(fileCount: Int) { + perFileArrays = Array(repeating: [:], count: fileCount) + perFileMetadata = Array(repeating: [:], count: fileCount) + namesPerFile = Array(repeating: nil, count: fileCount) + } + + func set( + file: Int, + arrays: [String: MLXArray], + metadata: [String: String], + names: [String]? + ) { + lock.lock() + defer { lock.unlock() } + perFileArrays[file] = arrays + perFileMetadata[file] = metadata + namesPerFile[file] = names + } + + func record(error: Error) { + lock.lock() + defer { lock.unlock() } + if firstError == nil { firstError = error } + } + + func result() throws -> ( + arrays: [[String: MLXArray]], + metadata: [[String: String]], + names: [[String]?] + ) { + lock.lock() + defer { lock.unlock() } + if let firstError { throw firstError } + return (perFileArrays, perFileMetadata, namesPerFile) + } +} + +/// Load a dictionary of ``MLXArray`` from several `safetensors` files at once, +/// scheduling payload reads in file-offset order and waiting for every array. +/// +/// The dictionaries are merged in file order: on a duplicate name the later file wins, +/// matching a serial load-and-merge loop over the same `urls`. Passing a single URL +/// also submits its tensors in file-offset order. MLX handles read concurrency. +/// +/// This function is synchronous and blocks until every file is loaded. From async code, +/// call it via a continuation on a `DispatchQueue` rather than directly on a Swift +/// concurrency cooperative thread. +/// +/// - Parameters: +/// - urls: URLs of the `safetensors` files to load +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``loadArrays(url:stream:)`` +/// - ``loadArraysAndMetadata(urls:stream:)`` +public func loadArrays(urls: [URL], stream: StreamOrDevice = .cpu) throws -> [String: MLXArray] { + try loadArraysAndMetadata(urls: urls, stream: stream).0 +} + +/// Load a dictionary of ``MLXArray`` and merged metadata `[String: String]` from several +/// `safetensors` files at once, scheduling payload reads in file-offset order. +/// +/// See ``loadArrays(urls:stream:)``. The metadata dictionaries are merged with the same +/// rule as the arrays: on a duplicate key the later file wins. +/// +/// - Parameters: +/// - urls: URLs of the `safetensors` files to load +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``loadArrays(urls:stream:)`` +/// - ``loadArraysAndMetadata(url:stream:)`` +public func loadArraysAndMetadata(urls: [URL], stream: StreamOrDevice = .cpu) throws -> ( + [String: MLXArray], [String: String] +) { + // Only header parsing happens on Dispatch worker threads. Payload reads stay lazy + // until checkedEval below, where MLX owns their concurrency. + let preparation = ConcurrentLoadPreparation(fileCount: urls.count) + DispatchQueue.concurrentPerform(iterations: urls.count) { file in + let url = urls[file] + let names = try? safetensorNamesInFileOrder(url: url) + do { + let (arrays, metadata) = try loadArraysAndMetadata(url: url, stream: stream) + preparation.set( + file: file, arrays: arrays, metadata: metadata, names: names) + } catch { + preparation.record(error: error) + } + } + + let prepared = try preparation.result() + let perFileArrays = prepared.arrays + let perFileMetadata = prepared.metadata + let evaluationOrder = safetensorArraysInEvaluationOrder( + arrays: perFileArrays, names: prepared.names) + + // checkedEval schedules the ordered roots under evalLock and waits outside it. + // Keep every root, including overwritten names, through this checked I/O barrier. + try checkedEval(evaluationOrder) + + var arrays = [String: MLXArray]() + var metadata = [String: String]() + for fileArrays in perFileArrays { + arrays.merge(fileArrays) { _, new in new } + } + for fileMetadata in perFileMetadata { + metadata.merge(fileMetadata) { _, new in new } + } + return (arrays, metadata) +} diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index 98347ade4..62bb3b075 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -3,6 +3,109 @@ import Cmlx import Foundation +/// Byte-level progress for loading arrays from disk. +/// +/// `completedUnitCount` and `totalUnitCount` are bytes and describe a single file. +/// Progress callbacks for each load are delivered in monotonically increasing order. +/// Loading the same file again starts a new count. Callbacks for different loads may +/// interleave -- use ``url`` to aggregate progress when loading a model made of several `safetensors` shards. +/// +/// Loading `safetensors` is lazy, which makes the reported progress approximate: +/// +/// - It may stop short of `totalUnitCount`. Only arrays that are actually evaluated are +/// read, so weights that are dropped before evaluation -- the ones a model's +/// `sanitize(weights:metadata:)` discards, for example -- are never read at all. Treat +/// completion of array evaluation as completion rather than waiting for +/// ``fractionCompleted`` to reach `1`. +/// - It counts bytes read rather than bytes of the file covered, so a region that happens +/// to be read more than once is counted more than once. ``completedUnitCount`` is +/// clamped to `totalUnitCount`, so ``fractionCompleted`` never exceeds `1`. +public struct LoadProgress: Sendable, Equatable { + /// The file being read. + public let url: URL + + /// Bytes read so far, clamped to ``totalUnitCount``. + public let completedUnitCount: Int64 + + /// Size of the file in bytes. + public let totalUnitCount: Int64 + + public var fractionCompleted: Double { + guard totalUnitCount > 0 else { return 0 } + return min(1, max(0, Double(completedUnitCount) / Double(totalUnitCount))) + } + + public init(url: URL, completedUnitCount: Int64, totalUnitCount: Int64) { + self.url = url + self.completedUnitCount = completedUnitCount + self.totalUnitCount = totalUnitCount + } +} + +/// Holder for the scoped ``withLoadProgressHandler(_:_:)-3ghip`` handler. +enum LoadProgressHandler { + + /// The stack of installed handlers -- the innermost scope wins. + @TaskLocal + static var handlers: [@Sendable (LoadProgress) -> Void] = [] + + static var current: (@Sendable (LoadProgress) -> Void)? { + handlers.last + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads. +/// +/// Any ``loadArrays(url:stream:)`` or ``loadArraysAndMetadata(url:stream:)`` performed +/// inside `body` reports byte progress to `handler`, without the call site having to pass +/// a progress handler explicitly. This makes it possible to drive a precise loading +/// progress bar for code -- such as a model loading library -- that you do not control: +/// +/// ```swift +/// let tracker = LoadProgressTracker(totalBytes: totalBytesOfSafetensors(in: directory)) +/// let model = try withLoadProgressHandler({ tracker.update($0) }) { +/// try loadModel(from: directory) +/// } +/// ``` +/// +/// Loading is lazy: the reader captures the handler when the file is opened and keeps +/// reporting to it when the arrays are evaluated, even after `body` returns. Include +/// evaluation in `body` to observe all progress before it returns. Arrays that are never +/// evaluated are never read, so progress may legitimately stop short of the file size. +/// +/// - Note: `handler` is called on the loading thread and MLX worker threads, potentially +/// concurrently for different loads, and is on the critical path of the read. It should be cheap and +/// must not call back into MLX loading or evaluation. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +/// +/// ### See Also +/// - ``loadArrays(url:stream:progressHandler:)`` +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () throws -> R +) rethrows -> R { + try LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try body() + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads (async). +/// +/// See ``withLoadProgressHandler(_:_:)-3ghip`` for details. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () async throws -> R +) async rethrows -> R { + try await LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try await body() + } +} + public enum LoadSaveError: Error { case unableToOpen(URL, String) case unknownExtension(String) @@ -117,6 +220,9 @@ public func loadArray(url: URL, stream: StreamOrDevice = .cpu) throws -> MLXArra /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArray(url:stream:)`` /// - ``loadArraysAndMetadata(url:stream:)`` @@ -128,6 +234,10 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArrays(url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -143,12 +253,35 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin } } +/// Load dictionary of ``MLXArray`` from a `safetensors` file, reporting byte progress as +/// lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. Called on the loading thread and MLX worker +/// threads, potentially concurrently across loads. Keep it cheap and do not call +/// MLX loading or evaluation from the callback. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. +public func loadArrays( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> [String: MLXArray] { + let (arrays, _) = try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + return arrays +} + /// Load dictionary of ``MLXArray`` and metadata `[String:String]` from a `safetensors` file. /// /// - Parameters: /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArrays(url:stream:)`` /// - ``loadArray(url:stream:)`` @@ -160,6 +293,11 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -175,6 +313,47 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw } } +/// Load dictionary of ``MLXArray`` and metadata from a `safetensors` file, reporting byte +/// progress as lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. Called on the loading thread and MLX worker +/// threads, potentially concurrently across loads. Keep it cheap and do not call +/// MLX loading or evaluation from the callback. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. +/// +/// ### See Also +/// - ``loadArraysAndMetadata(url:stream:)`` +/// - ``loadArrays(url:stream:progressHandler:)`` +public func loadArraysAndMetadata( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> ([String: MLXArray], [String: String]) { + precondition(url.isFileURL) + + switch url.pathExtension { + case "safetensors": + var r0 = mlx_map_string_to_array_new() + var r1 = mlx_map_string_to_string_new() + defer { mlx_map_string_to_array_free(r0) } + defer { mlx_map_string_to_string_free(r1) } + + let reader = try new_mlx_io_reader_fileIO(url, progressHandler: progressHandler) + defer { mlx_io_reader_free(reader) } + + _ = try withError { + mlx_load_safetensors_reader(&r0, &r1, reader, stream.ctx) + } + + return (mlx_map_array_values(r0), mlx_map_string_values(r1)) + default: + throw LoadSaveError.unknownExtension(url.pathExtension) + } +} + // MARK: - Memory I/O private class IOState { @@ -187,6 +366,190 @@ private class IOState { } } +private final class FileIOState { + private static let maximumReadChunkSize = 4 * 1024 * 1024 + + private let descriptor: CInt + private let lock = NSLock() + private let progressLock = NSLock() + private var offset: Int64 = 0 + private var completedUnitCount: Int64 = 0 + private var readError: String? + private let progressHandler: @Sendable (LoadProgress) -> Void + private let labelPointer: UnsafeMutablePointer + private let url: URL + + let totalUnitCount: Int64 + + init(url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void) throws { + let path = url.path(percentEncoded: false) + let descriptor = path.withCString { open($0, O_RDONLY) } + guard descriptor >= 0 else { + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + var statBuffer = stat() + guard fstat(descriptor, &statBuffer) == 0 else { + let message = String(cString: strerror(errno)) + close(descriptor) + throw LoadSaveError.unableToOpen(url, message) + } + + guard let labelPointer = strdup("file \(path)") else { + close(descriptor) + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + self.descriptor = descriptor + self.totalUnitCount = max(0, Int64(statBuffer.st_size)) + self.progressHandler = progressHandler + self.labelPointer = labelPointer + self.url = url + + progressHandler( + .init(url: url, completedUnitCount: 0, totalUnitCount: totalUnitCount)) + } + + deinit { + close(descriptor) + free(labelPointer) + } + + var isOpen: Bool { + descriptor >= 0 + } + + var good: Bool { + lock.withLock { + readError == nil + } + } + + var label: UnsafePointer { + UnsafePointer(labelPointer) + } + + func tell() -> Int { + lock.withLock { + Int(offset) + } + } + + /// Move the read position. + /// + /// Returns `0` on success and `-1` on failure; `CReader` turns a negative result into + /// a thrown error (ml-explore/mlx-c#130). + @discardableResult + func seek(offset newOffset: Int64, whence: Int32) -> Int32 { + lock.withLock { () -> Int32 in + let updated: Int64 + switch whence { + case SEEK_SET: + updated = newOffset + case SEEK_CUR: + let (position, overflow) = offset.addingReportingOverflow(newOffset) + guard !overflow else { return -1 } + updated = position + case SEEK_END: + // offset is relative to the end of the file, not the current position. + let (position, overflow) = totalUnitCount.addingReportingOverflow(newOffset) + guard !overflow else { return -1 } + updated = position + default: + return -1 + } + guard updated >= 0 else { return -1 } + offset = updated + return 0 + } + } + + /// Read `count` bytes at the current position and advance it. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int) -> Int { + guard let data else { return 0 } + + let readOffset = lock.withLock { + offset + } + let bytesRead = readBytes(to: data, count: count, offset: readOffset) + + lock.withLock { + offset += Int64(bytesRead) + } + + return bytesRead + } + + /// Read `count` bytes at an absolute `offset`, leaving the current position alone. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + /// This is the path used to materialize lazily loaded arrays, so it is called from + /// MLX worker threads. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int, offset readOffset: Int64) -> Int { + guard let data else { return 0 } + + return readBytes(to: data, count: count, offset: readOffset) + } + + private func readBytes( + to data: UnsafeMutablePointer, count: Int, offset readOffset: Int64 + ) -> Int { + guard count >= 0, readOffset >= 0, readOffset <= Int64.max - Int64(count) else { + return 0 + } + + var totalRead = 0 + while totalRead < count { + let chunkSize = min(count - totalRead, Self.maximumReadChunkSize) + let bytesRead = pread( + descriptor, + UnsafeMutableRawPointer(data.advanced(by: totalRead)), + chunkSize, + off_t(readOffset + Int64(totalRead))) + if bytesRead < 0 && errno == EINTR { continue } + guard bytesRead > 0 else { + recordReadError(bytesRead: bytesRead, requestedCount: count - totalRead) + break + } + totalRead += bytesRead + reportProgress(bytesRead: bytesRead) + } + + return totalRead + } + + private func recordReadError(bytesRead: Int, requestedCount: Int) { + let message: String + if bytesRead < 0 { + message = String(cString: strerror(errno)) + } else { + message = "unexpected end of file while reading \(requestedCount) bytes" + } + lock.withLock { + if readError == nil { + readError = message + } + } + } + + private func reportProgress(bytesRead: Int) { + guard bytesRead > 0 else { return } + + progressLock.withLock { + completedUnitCount += min(totalUnitCount - completedUnitCount, Int64(bytesRead)) + let progress = LoadProgress( + url: url, + completedUnitCount: completedUnitCount, + totalUnitCount: totalUnitCount) + progressHandler(progress) + } + } +} + private let label: StaticString = "\0" private func getData(_ writer: mlx_io_writer) -> Data { @@ -291,6 +654,51 @@ private func new_mlx_io_reader_dataIO(_ data: Data) -> mlx_io_reader { return mlx_io_reader_new(ptr, new_mlx_io_vtable_dataIO()) } +private func new_mlx_io_vtable_fileIO() -> mlx_io_vtable { + mlx_io_vtable { ptr in + guard let ptr else { return false } + return Unmanaged.fromOpaque(ptr).takeUnretainedValue().isOpen + } good: { ptr in + guard let ptr else { return false } + let state = Unmanaged.fromOpaque(ptr).takeUnretainedValue() + return state.isOpen && state.good + } tell: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.tell() + + } seek: { ptr, offset, whence in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.seek(offset: Int64(offset), whence: whence) + + } read: { ptr, data, n in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.read(to: data, count: n) + + } read_at_offset: { ptr, data, n, offset in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.read(to: data, count: n, offset: Int64(offset)) + + } write: { _, _, _ in + // read only -- reporting 0 bytes written makes CWriter throw + return 0 + + } label: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.label + + } free: { ptr in + Unmanaged.fromOpaque(ptr!).release() + } +} + +private func new_mlx_io_reader_fileIO( + _ url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> mlx_io_reader { + let ptr = Unmanaged.passRetained(try FileIOState(url: url, progressHandler: progressHandler)) + .toOpaque() + return mlx_io_reader_new(ptr, new_mlx_io_vtable_fileIO()) +} + private func new_mlx_io_writer_dataIO() -> mlx_io_writer { let ptr = Unmanaged.passRetained(IOState()).toOpaque() return mlx_io_writer_new(ptr, new_mlx_io_vtable_dataIO()) diff --git a/Source/MLX/MLXFast.swift b/Source/MLX/MLXFast.swift index 0834798a0..f4f697c7c 100644 --- a/Source/MLX/MLXFast.swift +++ b/Source/MLX/MLXFast.swift @@ -4,6 +4,42 @@ import Cmlx public enum MLXFast { + /// Chunked gated delta network forward pass. + /// + /// - Parameters: + /// - queries: shape (B, T, Hk, Dk) + /// - keys: shape (B, T, Hk, Dk) + /// - values: shape (B, T, Hv, Dv) + /// - gamma: decay rate in linear space (B, T, Hv) + /// - beta: delta update rates (B, T, Hv) + /// - initialState: initial hidden state (B, Hv, Dv, Dk), zeros if not given + /// - mask: optional boolean mask (B, T); masked steps leave the state unchanged + /// - stream: stream or device to evaluate on + /// - Returns: `output` (B, T, Hv, Dv) and the final `state` (B, Hv, Dv, Dk) + public static func gatedDeltaUpdate( + queries: MLXArray, + keys: MLXArray, + values: MLXArray, + gamma: MLXArray, + beta: MLXArray, + initialState: MLXArray? = nil, + mask: MLXArray? = nil, + stream: StreamOrDevice = .default + ) -> (output: MLXArray, state: MLXArray) { + var result = mlx_vector_array_new() + defer { mlx_vector_array_free(result) } + mlx_fast_gated_delta_update( + &result, + queries.ctx, keys.ctx, values.ctx, + gamma.ctx, beta.ctx, + (initialState ?? .mlxNone).ctx, + (mask ?? .mlxNone).ctx, + stream.ctx) + + let tuple = mlx_vector_array_values(result) + return (output: tuple[0], state: tuple[1]) + } + /// Optimized implementation of `NN.RoPE`. /// /// Used like this: diff --git a/Source/MLX/Memory.swift b/Source/MLX/Memory.swift index 7036633c0..ead36edfa 100644 --- a/Source/MLX/Memory.swift +++ b/Source/MLX/Memory.swift @@ -358,4 +358,21 @@ public enum Memory { mlx_clear_cache() } } + + /// Get the size of the buffers backing the given arrays in bytes. + /// + /// Each unique buffer is counted once. The full allocator size of each buffer + /// is used, which can exceed the logical size of its arrays. The arrays must be + /// evaluated before calling this function. + /// + /// - Parameter arrays: arrays to count + /// - Returns: size in bytes + public static func bufferSize(_ arrays: some Collection) -> Int { + let vector = new_mlx_vector_array(arrays) + defer { mlx_vector_array_free(vector) } + + var result = 0 + mlx_get_array_buffer_size(&result, vector) + return result + } } diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index a1cba744f..98b48b7e0 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -1928,6 +1928,7 @@ public func take(_ array: MLXArray, _ indices: MLXArray, stream: StreamOrDevice /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axis:stream:)`` /// - ``transposed(_:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` @@ -1941,6 +1942,19 @@ public func transposed( return MLXArray(result) } +/// Transpose the dimensions of the array. +/// +/// - Parameters: +/// - array: input array +/// - axes: Specifies the source axis for each axis in the new array +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - ``matrixTransposed(_:stream:)`` +/// - ``transposed(_:axis:stream:)`` +/// - ``transposed(_:stream:)`` +/// - ``MLXArray/transposed(axes:stream:)`` public func transposed(_ array: MLXArray, _ axes: Int..., stream: StreamOrDevice = .default) -> MLXArray { @@ -1955,6 +1969,7 @@ public func transposed(_ array: MLXArray, _ axes: Int..., stream: StreamOrDevice /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axes:stream:)`` /// - ``transposed(_:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` @@ -1971,6 +1986,7 @@ public func transposed(_ array: MLXArray, axis: Int, stream: StreamOrDevice = .d /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axes:stream:)`` /// - ``transposed(_:axis:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index 5a0da9e1a..f653918a0 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -1491,7 +1491,7 @@ public func gatherMM( @available( *, deprecated, renamed: - "gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:sortedIndices:stream:)" + "gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:globalScale:sortedIndices:stream:)" ) public func gatherQuantizedMatmul( _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray?, @@ -1528,6 +1528,8 @@ public func gatherQuantizedMatmul( /// - groupSize: The size of the group in `w` that shares a scale and bias. Default is `64` /// - bits: The number of bits occupied by each element in `w`. Default is `4` /// - mode: The quantization mode. Default is `.affine` +/// - globalScale: The per-input float32 scale used for +/// `nvfp4` quantization of `w`. Only supported on the GPU. /// - sortedIndices: Whether the indices are sorted. Default is `false` /// - stream: Stream or device to evaluate on /// @@ -1535,10 +1537,11 @@ public func gatherQuantizedMatmul( /// - /// - ``quantizedMM(_:_:scales:biases:transpose:groupSize:bits:mode:stream:)`` public func gatherQuantizedMM( - _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray?, + _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray? = nil, lhsIndices: MLXArray? = nil, rhsIndices: MLXArray? = nil, transpose: Bool = true, groupSize: Int? = nil, bits: Int? = nil, mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil, sortedIndices: Bool = false, stream: StreamOrDevice = .default ) -> MLXArray { @@ -1551,7 +1554,10 @@ public func gatherQuantizedMM( &result, x.ctx, w.ctx, scales.ctx, (biases ?? .mlxNone).ctx, (lhsIndices ?? .mlxNone).ctx, (rhsIndices ?? .mlxNone).ctx, transpose, - gs, bits, mode.rawValue, sortedIndices, + gs, bits, + mode.rawValue, + (globalScale ?? .mlxNone).ctx, + sortedIndices, stream.ctx) return MLXArray(result) @@ -2043,6 +2049,22 @@ public func meshGrid( return mlx_vector_array_values(vec) } +/// Transpose the last two dimensions of the array. +/// +/// ### See Also +/// - ``transposed(_:stream:)`` +/// +/// - Parameters: +/// - array: array to transpose +/// - stream: stream or device to evaluate on +public func matrixTransposed( + _ array: MLXArray, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_matrix_transpose(&result, array.ctx, stream.ctx) + return MLXArray(result) +} + /// Element-wise maximum. /// /// Take the element-wise max of two arrays with @@ -2999,6 +3021,163 @@ public func unstacked( return mlx_vector_array_values(result) } +/// Returns the sorted unique elements of the flattened array. +/// +/// The shape of the output does not depend on the values in `array`, so +/// `array` is not evaluated. Unlike NumPy, the size is never inferred from +/// the values and must be provided. +/// +/// Entries past the last unique element hold `fillValue`. +/// A truncated output keeps the smallest `size` unique elements. +/// +/// This op builds on ``sorted(_:stream:)``, so a `Bool` input needs the CPU +/// stream. +/// +/// ### See Also +/// - ``uniqued(_:size:fillValue:returning:stream:)`` +/// +/// - Parameters: +/// - array: input array +/// - size: The size of the output. If the size is smaller than +/// the number of unique elements of `array`, the output is truncated. +/// If it is larger, the output is padded with `fillValue`. +/// - fillValue: optional value to fill empty positions in the result. If +/// not specified this defaults to the first of the sorted unique elements. +/// - stream: stream or device to evaluate on +/// - Returns: sorted unique elements +public func uniqued( + _ array: MLXArray, + size: Int, + fillValue: MLXArray? = nil, + stream: StreamOrDevice = .default +) -> MLXArray { + var vector = mlx_vector_array_new() + defer { mlx_vector_array_free(vector) } + + mlx_unique( + &vector, + array.ctx, size.int32, + false, false, false, + (fillValue ?? .mlxNone).ctx, + stream.ctx) + + let arrays = mlx_vector_array_values(vector) + return arrays[0] +} + +/// Options for `uniqued()`. +/// +/// See ``uniqued(_:size:fillValue:returning:stream:)``. +public struct UniqueOptions: OptionSet, Sendable { + public let rawValue: Int + public init(rawValue: Int) { self.rawValue = rawValue } + + /// Requests the index in the flattened input of the first + /// occurence of each unique element. + public static let index = UniqueOptions(rawValue: 1 << 0) + + /// Requests the indices of the unique array that rebuild + /// the input to `uniqued()`. + public static let inverse = UniqueOptions(rawValue: 1 << 1) + + /// Requests the number of times each unique element occurs in the + /// input. + public static let counts = UniqueOptions(rawValue: 1 << 2) +} + +/// Return value for `uniqued()`. +/// +/// See ``uniqued(_:size:fillValue:returning:stream:)``. +public struct UniqueResult { + + /// Array of unique values of the given size. + public var values: MLXArray + public var index: MLXArray? + public var inverse: MLXArray? + public var counts: MLXArray? + + public init( + values: MLXArray, + index: MLXArray? = nil, inverse: MLXArray? = nil, counts: MLXArray? = nil + ) { + self.values = values + self.index = index + self.inverse = inverse + self.counts = counts + } +} + +/// Returns the sorted unique elements of the flattened array. +/// +/// The shape of the output does not depend on the values in `array`, so +/// `array` is not evaluated. Unlike NumPy, the size is never inferred from +/// the values and must be provided. +/// +/// Entries past the last unique element hold `fillValue` and their +/// count is `0`, so `sum(result.counts > 0)` gives the number of unique +/// elements as long as the output is not truncated. +/// +/// A truncated output keeps the smallest `size` unique elements, and +/// `result.index`, `result.inverse` and `result.counts` +/// then stop describing all of `array`, +/// and `result.inverse` and `result.counts` stop agreeing with each other: +/// the counts of the dropped elements are gone, while their +/// indices in `result.inverse` are clamped to the last entry. +/// +/// This op builds on ``sorted(_:stream:)``, so a `Bool` input needs the CPU +/// stream. +/// +/// ### See Also +/// - ``uniqued(_:size:fillValue:stream:)`` +/// +/// - Parameters: +/// - array: input array +/// - size: The size of the output. If the size is smaller than +/// the number of unique elements of `array`, the output is truncated. +/// If it is larger, the output is padded with `fillValue`. +/// - fillValue: optional value to fill empty positions in the result. If +/// not specified this defaults to the first of the sorted unique elements. +/// - returning: Additional values to return +/// - stream: stream or device to evaluate on +/// - Returns: sorted unique elements +public func uniqued( + _ array: MLXArray, + size: Int, + fillValue: MLXArray? = nil, + returning: UniqueOptions, + stream: StreamOrDevice = .default +) -> UniqueResult { + var vector = mlx_vector_array_new() + defer { mlx_vector_array_free(vector) } + + mlx_unique( + &vector, + array.ctx, size.int32, + returning.contains(.index), + returning.contains(.inverse), + returning.contains(.counts), + (fillValue ?? .mlxNone).ctx, + stream.ctx) + + let arrays = mlx_vector_array_values(vector) + var result = UniqueResult(values: arrays[0]) + + var readIndex = 1 + if returning.contains(.index) { + result.index = arrays[readIndex] + readIndex += 1 + } + if returning.contains(.inverse) { + result.inverse = arrays[readIndex] + readIndex += 1 + } + if returning.contains(.counts) { + result.counts = arrays[readIndex] + readIndex += 1 + } + return result +} + /// Stop gradients from being computed. /// ///The operation is the identity but it prevents gradients from flowing diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index bf3bb5ed5..d9ce19d42 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -3,17 +3,48 @@ import Cmlx import Foundation -// Note: this is all immutable state -- the `id` property is only set at init time -- -// except for ``CompiledFunction/tracedConstants``, which is guarded by -// ``CompiledFunction/lock``. -// Also: this is marked as Sendable even though it may capture non-Sendable state -// in inputs/outputs. The functions that create it will not return a Sendable -// type in that case -- see compile(inputs:) +/// A weak handle to MLX's thread-local cache. Retaining this wrapper does not keep +/// the backend cache alive after its thread exits. +private final class ThreadCompilationCache { + private let ctx: mlx_compile_cache + + private init() { + var cache = mlx_compile_cache_new() + mlx_detail_compile_cache(&cache) + self.ctx = cache + } + + deinit { + mlx_compile_cache_free(ctx) + } + + static var current: ThreadCompilationCache { + let key = "MLX.ThreadCompilationCache" + let dictionary = Thread.current.threadDictionary + if let cache = dictionary[key] as? ThreadCompilationCache { + return cache + } + let cache = ThreadCompilationCache() + dictionary[key] = cache + return cache + } + + func erase(_ functionID: UInt) { + mlx_detail_compile_erase(ctx, functionID) + } +} + +// Mutable cache bookkeeping is guarded by evalLock. The stateful overloads also +// hold the instance lock while replacing observed inputs and outputs. +// Only the overloads accepting Sendable captures return a Sendable closure. final class CompiledFunction: @unchecked (Sendable) { /// unique (for the lifetime of the object) identifier for the compiled function private var id: UInt! + /// Every backend cache this function has used, guarded by evalLock. + private var caches: [ObjectIdentifier: ThreadCompilationCache] = [:] + /// guards the observed state (``inputs`` / ``outputs``) -- see #226 /// /// This is always acquired *inside* `evalLock`; see ``call(_:)`` for the @@ -75,77 +106,24 @@ final class CompiledFunction: @unchecked (Sendable) { self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque()) } - /// Remove the compiled structure from the back end. - /// - /// ### Why this takes `evalLock` - /// - /// `mlx_detail_compile_erase` is `compiler_cache().erase(fun_id)`, i.e. - /// `std::unordered_map::erase` on the *same* process-global - /// `mlx::core::detail::CompilerCache` that ``innerCall(_:)`` reads and - /// inserts into via `CompilerCache::find` (`cache_[fun_id]`, which itself - /// can insert and rehash, followed by `entries.push_back`). That cache has - /// no synchronization of its own, and `find` hands back a *reference* to an - /// entry that the caller then fills in and reads from across the whole - /// trace. An unlocked erase can therefore rehash the map out from under a - /// concurrent `find`, or destroy the very entry another thread is compiling - /// into. - /// - /// Every other use of that cache happens under `evalLock`, so the erase has - /// to as well; otherwise the lock does not actually make the cache - /// single-threaded, it only makes most of it single-threaded. - /// - /// ### Why blocking here is safe - /// - /// `deinit` runs on whichever thread drops the last reference: - /// - /// - A thread that holds no MLX lock simply waits for the current eval or - /// compiled call to finish. Bounded, since `evalLock` is only ever held - /// for the duration of one call. - /// - A thread already inside a compiled call or an `eval` -- for instance a - /// traced body that drops the last reference to some other compiled - /// function -- already holds `evalLock`, and it is an `NSRecursiveLock`, - /// so re-entry is free. - /// - A thread holding some other ``CompiledFunction``'s ``lock`` also holds - /// `evalLock` already (see ``call(_:)``), so that is the same case; no new - /// edge is added to the lock graph. - /// - /// Within this library the only releases of a ``CompiledFunction`` are the - /// caller's own (the closure returned by ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` - /// holds the sole strong reference) and the `mlx_closure_free` of the - /// closure built in ``innerCall(_:)``, which runs on the calling thread with - /// `evalLock` already held. Nothing hands a reference to a back-end worker - /// thread, so the erase can never be the thing an `evalLock` holder is - /// waiting for. - /// - /// The erase stays synchronous rather than being deferred to a queue on - /// purpose: ``id`` is the object's own address, so once this object is - /// freed a new ``CompiledFunction`` can be allocated at the same address and - /// receive the same id. A deferred erase could delete that new function's - /// cache entries. + // Erase every calling thread's cache before this address can be reused as an ID. deinit { - let functionID = id! withEvalLock { #if DEBUG - // Trivially true two lines below the acquisition -- it is here to - // catch a future change that moves the erase back out from under - // the lock, the same way `withInstanceLock` guards `call(_:)`. EvalLockOwnership.requireHeldByCurrentThread( "the compiler cache was erased without holding evalLock; " + "see CompiledFunction.deinit") EvalLockOwnership.recordCompileErase() #endif - Self.eraseCache(functionID) + eraseCaches() } } - /// Remove the compiled structure for `functionID` from the back end. - /// - /// - Precondition: `evalLock` is held -- see ``deinit``. - private static func eraseCache(_ functionID: UInt) { - var cache = mlx_compile_cache_new() - mlx_detail_compile_cache(&cache) - defer { mlx_compile_cache_free(cache) } - mlx_detail_compile_erase(cache, functionID) + /// Must be called under evalLock, including when sealed state is rebound. + private func eraseCaches() { + for cache in caches.values { + cache.erase(id) + } } /// Evaluate the compiled function. @@ -155,10 +133,8 @@ final class CompiledFunction: @unchecked (Sendable) { /// `evalLock` is acquired **outermost**, before the per-instance ``lock``. /// That order is not a preference, it is forced: /// - /// - Tracing mutates process-global state that has no synchronization of its - /// own -- `mlx::core::detail::CompilerCache` is an unguarded singleton and - /// `detail::InTracing::trace_stack()` is a plain function-local `static`, - /// not `thread_local` -- so `evalLock` has to span the whole trace (#339). + /// - A trace can call back into Swift and mutate observed state, so + /// `evalLock` spans the whole trace and cache invalidation (#339). /// - The trace calls back into Swift, and the traced body routinely calls /// *other* compiled functions: every `MLXNN` activation is a file-scope /// `CompiledFunction`, so e.g. `silu` inside a traced body enters @@ -182,16 +158,7 @@ final class CompiledFunction: @unchecked (Sendable) { /// closure creation at the top of ``innerCall(_:)``. func call(_ arguments: [MLXArray]) -> [MLXArray] { withEvalLock { - // ``lock`` exists only to guard the observed state (#226): the - // `_updateInternal` swaps in `innerCall` are the sole mutation of - // anything reachable from `self` -- every other stored property is - // `let` except ``tracedConstants`` (only written when there is - // state), and `id` is written once at init. With no state there is - // nothing to guard, and `evalLock` above already serializes the - // compiler cache and every other caller. - // - // This is the common case: every compiled function in `MLXNN` is - // stateless. + // Stateless functions need only evalLock for cache bookkeeping. if inputs.isEmpty && outputs.isEmpty { return innerCall(arguments) } @@ -220,6 +187,9 @@ final class CompiledFunction: @unchecked (Sendable) { /// too unless this function has no observed state; see ``call(_:)``, /// which is the only caller. func innerCall(_ arguments: [MLXArray]) -> [MLXArray] { + let cache = ThreadCompilationCache.current + caches[ObjectIdentifier(cache)] = cache + // Sealed arrays (MaterializedArray, the parameters of a sealed module) // can never change, and must not be mutated, so they are not threaded // through the tracer swap below -- they are simply captured as @@ -239,7 +209,7 @@ final class CompiledFunction: @unchecked (Sendable) { traced.count == constants.count && zip(traced, constants).allSatisfy { $0 === $1 } if !same { - Self.eraseCache(id) + eraseCaches() tracedConstants = constants } } else { diff --git a/Source/MLXNN/Losses.swift b/Source/MLXNN/Losses.swift index bf8d8542e..b6c7c69fe 100644 --- a/Source/MLXNN/Losses.swift +++ b/Source/MLXNN/Losses.swift @@ -51,8 +51,8 @@ public func crossEntropy( if targets_as_probs { score = sum(logits * targets, axis: axis) } else { - score = takeAlong(logits, targets.expandedDimensions(axis: -1), axis: axis).squeezed( - axis: -1) + score = takeAlong(logits, targets.expandedDimensions(axis: axis), axis: axis).squeezed( + axis: axis) } let logSumExpLogits = logSumExp(logits, axis: axis) @@ -164,8 +164,8 @@ public func mseLoss(predictions: MLXArray, targets: MLXArray, reduction: LossRed public func nllLoss( inputs: MLXArray, targets: MLXArray, axis: Int = -1, reduction: LossReduction = .none ) -> MLXArray { - let loss = -takeAlong(inputs, targets.expandedDimensions(axis: -1), axis: axis).squeezed( - axis: -1) + let loss = -takeAlong(inputs, targets.expandedDimensions(axis: axis), axis: axis).squeezed( + axis: axis) return reduction.reduce(loss: loss) } diff --git a/Source/MLXNN/Module.swift b/Source/MLXNN/Module.swift index a2c0a0305..f1f7ea0af 100644 --- a/Source/MLXNN/Module.swift +++ b/Source/MLXNN/Module.swift @@ -509,20 +509,37 @@ open class Module { modulePath: [String] = [] ) throws -> Self { _checkNotSealed() - try update(parameters: parameters, verify: verify, path: path, modulePath: modulePath) { - m, k, a, v in - a._updateInternal(v) - } + try update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath, + mutate: nil) return self } + /// Applies `parameters` to this module and recurses into its children. + /// + /// With a `nil` `mutate`, arrays update in place and children recurse through the + /// open `update(parameters:verify:path:modulePath:)` so subclass overrides run. + /// A custom `mutate` (see `materialize()`) recurses directly and skips overrides. func update( parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [], modulePath: [String] = [], - mutate: (Module, String, MLXArray, MLXArray) -> Void + mutate: ((Module, String, MLXArray, MLXArray) -> Void)? ) throws { let modulePath = modulePath + [describeType(self)] + func updateChild( + _ module: Module, _ parameters: ModuleParameters, path: [String] + ) throws { + if let mutate { + try module.update( + parameters: parameters, verify: verify, path: path, + modulePath: modulePath, mutate: mutate) + } else { + try module.update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath) + } + } + func apply( key: String, path: [String], _ item: ModuleItem, _ value: NestedItem ) throws { @@ -542,7 +559,11 @@ open class Module { path: path, modules: modulePath, expectedShape: p.shape, actualShape: newArray.shape) } - mutate(self, key, p, newArray) + if let mutate { + mutate(self, key, p, newArray) + } else { + p._updateInternal(newArray) + } case (.value(.parameters), .none): if Self.parameterIsValid(key) { @@ -586,14 +607,10 @@ open class Module { } case (.value(.module(let module)), .dictionary(let values)): - try module.update( - parameters: NestedDictionary(values: values), verify: verify, path: path, - modulePath: modulePath, mutate: mutate) + try updateChild(module, NestedDictionary(values: values), path: path) case (.value(.module(let module)), .none): - try module.update( - parameters: NestedDictionary(), verify: verify, path: path, - modulePath: modulePath, mutate: mutate) + try updateChild(module, NestedDictionary(), path: path) case (.none, .none), (.value(.none), .none), (.value(.other(_)), .none): break diff --git a/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md b/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md index 02d7345b1..23ca9faa2 100644 --- a/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md +++ b/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md @@ -29,6 +29,21 @@ for _ in 0 ..< epochs { } ``` +## Checkpoints + +``Optimizer/state()`` is the buffers training has accumulated (momentum, moments, +and per-parameter step counts), in the same shape as the model parameters. +``Optimizer/update(parameters:)`` writes that tree back. The learning rate and +other hyperparameters stay on the optimizer you construct: + +```swift +eval(optimizer) +let state = optimizer.state() + +let resumed = Adam(learningRate: 1e-3) +resumed.update(parameters: state) +``` + ## Other MLX Packages - [MLX](mlx) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index 42f247517..b66671b5b 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -24,6 +24,71 @@ public protocol Optimizer: Updatable, Evaluatable { /// Apply the gradients to the parameters of the model and update the model with the new parameters. func update(model: Module, gradients: ModuleParameters) + + /// Buffers accumulated while training, as a `ModuleParameters` tree. + /// + /// The tree follows the model parameters. Optimizers that keep one buffer per + /// parameter (``SGD``, ``RMSprop``, ``AdaGrad``, ``Lion``, ``Muon``) store that + /// array at the parameter. Optimizers that keep several store a dictionary of + /// named buffers there: + /// + /// - ``Adam`` and ``AdamW``: `m`, `v`, and `step` + /// - ``Adamax``: `m` and `v` + /// - ``AdaDelta``: `v` and `u` + /// - ``Adafactor``: `step`, plus the running averages that parameter uses + /// (`expAvgSq`, or `expAvgSqRow` and `expAvgSqCol`, and `expAvg` when + /// `beta1` is set) + /// + /// ``MultiOptimizer`` nests each sub-optimizer under `states`, indexed + /// `"0"`, `"1"`, ... in optimizer order. + /// `step` is counted per parameter. Hyperparameters such as the learning rate are + /// not part of this tree — build the resumed optimizer with the same values. + /// + /// The result is empty until the first ``update(model:gradients:)`` or + /// ``update(parameters:)``. `flattened()` is the dictionary + /// `save(arrays:metadata:url:)` writes. + /// + /// ```swift + /// eval(optimizer) + /// let checkpoint = optimizer.state() + /// + /// let resumed = Adam(learningRate: 1e-3) + /// resumed.update(parameters: checkpoint) + /// ``` + func state() -> ModuleParameters + + /// Replace buffers with a tree produced by ``state()``. + /// + /// Parameters present in the tree replace the stored buffers. Parameters left + /// out are unchanged, as with `Module.update(parameters:)`. + /// + /// - Parameters: + /// - parameters: buffers from ``state()``, or the same tree read back from + /// `loadArrays(url:)` + func update(parameters: ModuleParameters) +} + +extension Optimizer { + // External `Optimizer` conformers predate this API. The built-in optimizers + // replace both methods; a custom one keeps compiling and returns no buffers + // until it implements them. + public func state() -> ModuleParameters { + ModuleParameters() + } + + public func update(parameters: ModuleParameters) { + } +} + +private func requiredState( + _ arrays: [(String, MLXArray)], keys: [String], owner: Any +) -> [String: MLXArray] { + let byKey = Dictionary(arrays, uniquingKeysWith: { _, latest in latest }) + let missing = keys.filter { byKey[$0] == nil } + precondition( + missing.isEmpty, + "\(type(of: owner)) state is missing \(missing.joined(separator: ", "))") + return byKey } /// An optimizer that delegates to several sub-optimizers, routing each parameter to the first @@ -88,6 +153,38 @@ open class MultiOptimizer: Optimizer { public func innerState() -> [MLXArray] { optimizers.flatMap { $0.innerState() } } + + /// Buffers for every sub-optimizer, under `states.0`, `states.1`, ... + public func state() -> ModuleParameters { + var flat: [(String, MLXArray)] = [] + // `states` matches Python. A bare leading index unflattens as an array, + // which ModuleParameters cannot hold at the root. + for (index, optimizer) in optimizers.enumerated() { + for (key, value) in optimizer.state().flattened() { + flat.append(("states.\(index).\(key)", value)) + } + } + return ModuleParameters.unflattened(flat) + } + + public func update(parameters: ModuleParameters) { + var grouped: [Int: [(String, MLXArray)]] = [:] + for (key, value) in parameters.flattened() { + // "states..", leaving any further dots intact. + let parts = key.split(separator: ".", maxSplits: 2).map(String.init) + guard parts.count == 3, parts[0] == "states", + let index = Int(parts[1]), + optimizers.indices.contains(index) + else { + preconditionFailure( + "MultiOptimizer state key \"\(key)\" is missing a sub-optimizer index") + } + grouped[index, default: []].append((parts[2], value)) + } + for (index, pairs) in grouped { + optimizers[index].update(parameters: ModuleParameters.unflattened(pairs)) + } + } } /// The base class for all optimizers. It allows us to implement an optimizer on a per-parameter basis @@ -124,6 +221,128 @@ open class OptimizerBase: Optimizer { .flatMap { $0.innerState() } } + /// Names written by ``exportState(_:)`` for one parameter. + /// + /// `""` stores the only buffer directly on the parameter. Any other name stores + /// a dictionary of buffers under the parameter (`m`, `v`, `step`, ...). + /// ``stateSuffixes()`` must list every name ``exportState(_:)`` can produce. + open func stateSuffixes() -> [String] { [""] } + + /// Encode one parameter's buffers. + /// + /// Return `[("", array)]` for a single buffer. Return `[(name, array)]` for + /// several, using the names from ``/MLXOptimizers/OptimizerBase/stateSuffixes()``. + open func exportState(_ state: State) -> [(String, MLXArray)] { + fatalError("exportState() not implemented \(type(of: self))") + } + + /// Rebuild one parameter's buffers from the pairs ``exportState(_:)`` produced. + open func importState(_ arrays: [(String, MLXArray)]) -> State { + fatalError("importState() not implemented \(type(of: self))") + } + + open func state() -> ModuleParameters { + ModuleParameters(item: exportTree(stateStorage.asItem())) + } + + open func update(parameters: ModuleParameters) { + let merged = mergeTree(stateStorage.asItem(), parameters.asItem()) + guard case .dictionary(let values) = merged else { + preconditionFailure("optimizer state must be a dictionary") + } + stateStorage = NestedDictionary(values: values) + } + + /// A dictionary of arrays whose keys are all buffer names is one parameter. + /// A dictionary of further dictionaries is the module tree, so a parameter + /// that is itself named `m` or `v` is not mistaken for a buffer. + private func isStateDictionary(_ updates: [String: NestedItem]) -> Bool { + let names = Set(stateSuffixes().filter { !$0.isEmpty }) + guard !names.isEmpty, !updates.isEmpty else { return false } + guard updates.keys.allSatisfy({ names.contains($0) }) else { return false } + for value in updates.values { + guard case .value = value else { return false } + } + return true + } + + private func exportTree(_ item: NestedItem) -> NestedItem { + switch item { + case .none: + return .none + case .value(let state): + return exportValue(state) + case .array(let items): + return .array(items.map { exportTree($0) }) + case .dictionary(let dictionary): + return .dictionary(dictionary.mapValues { exportTree($0) }) + } + } + + private func exportValue(_ state: State) -> NestedItem { + let pairs = exportState(state) + precondition(!pairs.isEmpty, "\(type(of: self)) exportState() returned no buffers") + if pairs.count == 1, pairs[0].0.isEmpty { + return .value(pairs[0].1) + } + + var exported: [String: NestedItem] = [:] + for (name, array) in pairs { + precondition( + !name.isEmpty, "\(type(of: self)) mixed an unnamed buffer with named buffers") + exported[name] = .value(array) + } + return .dictionary(exported) + } + + private func mergeTree( + _ existing: NestedItem, _ update: NestedItem + ) -> NestedItem { + switch update { + case .none: + return existing + case .value(let array): + return .value(importState([("", array)])) + case .array(let updates): + var merged: [NestedItem] + if case .array(let items) = existing { + merged = items + } else { + merged = [] + } + if merged.count < updates.count { + merged.append( + contentsOf: Array(repeating: .none, count: updates.count - merged.count)) + } + for (index, item) in updates.enumerated() { + merged[index] = mergeTree(merged[index], item) + } + return .array(merged) + case .dictionary(let updates): + if isStateDictionary(updates) { + return .value( + importState( + updates.map { name, value in + guard case .value(let array) = value else { + preconditionFailure( + "optimizer state buffer \(name) is not an array") + } + return (name, array) + })) + } + var merged: [String: NestedItem] + if case .dictionary(let items) = existing { + merged = items + } else { + merged = [:] + } + for (key, item) in updates { + merged[key] = mergeTree(merged[key] ?? .none, item) + } + return .dictionary(merged) + } + } + final public func update(model: Module, gradients: ModuleParameters) { model.update(parameters: apply(gradients: gradients, modelParameters: model.parameters())) } @@ -159,6 +378,17 @@ open class OptimizerBaseArrayState: OptimizerBase { override open func newState(parameter: MLXArray) -> MLXArray { MLXArray.zeros(like: parameter) } + + override open func exportState(_ state: MLXArray) -> [(String, MLXArray)] { + [("", state)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> MLXArray { + guard let array = arrays.first(where: { $0.0.isEmpty })?.1 else { + preconditionFailure("\(type(of: self)) state is missing the parameter buffer") + } + return array + } } /// State container for ``OptimizerBase`` holding a tuple of `MLXArray`. @@ -388,6 +618,17 @@ open class AdaDelta: OptimizerBase { TupleState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["v", "u"] } + + override open func exportState(_ state: TupleState) -> [(String, MLXArray)] { + [("v", state.values.0), ("u", state.values.1)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> TupleState { + let byKey = requiredState(arrays, keys: ["v", "u"], owner: self) + return TupleState(byKey["v"]!, byKey["u"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( MLXArray, TupleState ) { @@ -441,6 +682,17 @@ open class Adam: OptimizerBase { AdamState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["m", "v", "step"] } + + override open func exportState(_ state: AdamState) -> [(String, MLXArray)] { + [("m", state.values.0), ("v", state.values.1), ("step", state.step)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> AdamState { + let byKey = requiredState(arrays, keys: ["m", "v", "step"], owner: self) + return AdamState(byKey["m"]!, byKey["v"]!, step: byKey["step"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: AdamState) -> ( MLXArray, AdamState ) { @@ -541,6 +793,17 @@ open class Adamax: OptimizerBase { TupleState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["m", "v"] } + + override open func exportState(_ state: TupleState) -> [(String, MLXArray)] { + [("m", state.values.0), ("v", state.values.1)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> TupleState { + let byKey = requiredState(arrays, keys: ["m", "v"], owner: self) + return TupleState(byKey["m"]!, byKey["v"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( MLXArray, TupleState ) { @@ -643,7 +906,10 @@ open class Adafactor: OptimizerBase { var expAvg: MLXArray? = nil public func innerState() -> [MLXArray] { - [expAvgSqRow, expAvgSqCol, expAvgSq, expAvg].compactMap { $0 } + // `compile` and `eval` only see these arrays. `step` has to be + // here or a compiled run keeps the count from the first trace, + // and `state()` saves that stuck value. + [expAvgSqRow, expAvgSqCol, expAvgSq, expAvg].compactMap { $0 } + [step] } } @@ -696,6 +962,30 @@ open class Adafactor: OptimizerBase { return s } + override open func stateSuffixes() -> [String] { + ["step", "expAvgSqRow", "expAvgSqCol", "expAvgSq", "expAvg"] + } + + override open func exportState(_ state: State) -> [(String, MLXArray)] { + var exported: [(String, MLXArray)] = [("step", state.step)] + if let value = state.expAvgSqRow { exported.append(("expAvgSqRow", value)) } + if let value = state.expAvgSqCol { exported.append(("expAvgSqCol", value)) } + if let value = state.expAvgSq { exported.append(("expAvgSq", value)) } + if let value = state.expAvg { exported.append(("expAvg", value)) } + return exported + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> State { + let byKey = requiredState(arrays, keys: ["step"], owner: self) + var state = State() + state.step = byKey["step"]! + state.expAvgSqRow = byKey["expAvgSqRow"] + state.expAvgSqCol = byKey["expAvgSqCol"] + state.expAvgSq = byKey["expAvgSq"] + state.expAvg = byKey["expAvg"] + return state + } + func rms(_ inputs: MLXArray) -> MLXArray { sqrt(mean(square(inputs))) } diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift index 6a0102a4a..ac560dad1 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 7 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 15 import Foundation import MLX @@ -198,4 +200,272 @@ struct GeneratedFastTests { tolerance: .float32) } } + + @Test("gatedDeltaUpdate/fallback/output") + func test_gatedDeltaUpdate_fallback_output() throws { + // Hv != Hk, shape not supported by the kernel + try withIntegrationState(seed: 53947) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: -0.018037354573607445, + minimum: -0.5729220509529114, + maximum: 0.4151778519153595, + absoluteSum: 30.585939407348633, + positionChecksum: 17.2472412109375, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -0.05421087518334389, -0.12156698107719421, -0.15238893032073975, + -0.14626511931419373, 0.2740327715873718, -0.10090052336454391, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/state") + func test_gatedDeltaUpdate_fallback_state() throws { + // Hv != Hk, shape not supported by the kernel + try withIntegrationState(seed: 93688) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: 0.006113756913691759, + minimum: -0.15103420615196228, + maximum: 0.1642816811800003, + absoluteSum: 18.600929260253906, + positionChecksum: 9.371909141540527, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.07339409738779068, 0.03697947785258293, 0.03930574283003807, + 0.003275439143180847, 0.005969321355223656, 0.0017838235944509506, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/initialState/output") + func test_gatedDeltaUpdate_fallback_initialState_output() throws { + try withIntegrationState(seed: 34247) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: 0.05453367158770561, + minimum: -3.9586031436920166, + maximum: 7.431026935577393, + absoluteSum: 279.72869873046875, + positionChecksum: 135.665087890625, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -1.2578498125076294, 1.2448211908340454, 0.4136521518230438, + -0.41495442390441895, 0.16835975646972656, -0.2699488699436188, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/initialState/state") + func test_gatedDeltaUpdate_fallback_initialState_state() throws { + try withIntegrationState(seed: 95025) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: -0.009089105762541294, + minimum: -0.5524362325668335, + maximum: 0.5001603960990906, + absoluteSum: 72.35923767089844, + positionChecksum: 33.11988830566406, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.20271210372447968, -0.2317522168159485, 0.268993079662323, + 0.03319134935736656, 0.012991426512598991, 0.024516910314559937, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/mask/output") + func test_gatedDeltaUpdate_fallback_mask_output() throws { + // masked steps leave the state unchanged and output zeros + try withIntegrationState(seed: 92033) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MLXArray( + [true, true, true, true, true, true, true, true, false, false], [2, 5]) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState, mask: mask + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: 0.052943795919418335, + minimum: -6.441189765930176, + maximum: 6.602780342102051, + absoluteSum: 320.8973693847656, + positionChecksum: 117.5300537109375, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -5.145998001098633, -0.6544278860092163, 0.9613724946975708, + 0.406800240278244, -0.6966778635978699, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/mask/state") + func test_gatedDeltaUpdate_fallback_mask_state() throws { + // masked steps leave the state unchanged and output zeros + try withIntegrationState(seed: 40800) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MLXArray( + [true, true, true, true, true, true, true, true, false, false], [2, 5]) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState, mask: mask + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: -0.005854031536728144, + minimum: -1.9470419883728027, + maximum: 2.252192974090576, + absoluteSum: 151.9166259765625, + positionChecksum: 87.9920883178711, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.21978959441184998, -0.0349111407995224, 0.3030872046947479, + 0.8860476016998291, -0.3838585615158081, 0.10365576297044754, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/kernel/output") + func test_gatedDeltaUpdate_kernel_output() throws { + // a shape the metal kernel supports + try withIntegrationState(seed: 91400) { + let queries = MLXRandom.normal([1, 3, 16, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [1, 3, 16, 128], dtype: .float32) + let values = MLXRandom.normal([1, 3, 32, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [1, 3, 32], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [1, 3, 32], dtype: .float32) + let initialState = MLXRandom.normal( + [1, 32, 128, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).output + expectSummary( + result, + ArraySummary( + shape: [1, 3, 32, 128], + dtype: .float32, + mean: -0.07715722918510437, + minimum: -33.80734634399414, + maximum: 35.589317321777344, + absoluteSum: 67446.203125, + positionChecksum: 29519.645833333332, + sampleIndices: [0, 2457, 4915, 7372, 9830, 12287], + samples: [ + 6.1729230880737305, -1.2932864427566528, 4.83355188369751, + 9.386662483215332, 3.525575637817383, -4.2895827293396, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/kernel/state") + func test_gatedDeltaUpdate_kernel_state() throws { + // a shape the metal kernel supports + try withIntegrationState(seed: 81342) { + let queries = MLXRandom.normal([1, 3, 16, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [1, 3, 16, 128], dtype: .float32) + let values = MLXRandom.normal([1, 3, 32, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [1, 3, 32], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [1, 3, 32], dtype: .float32) + let initialState = MLXRandom.normal( + [1, 32, 128, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).state + expectSummary( + result, + ArraySummary( + shape: [1, 32, 128, 128], + dtype: .float32, + mean: -0.0014557009562849998, + minimum: -3.1630849838256836, + maximum: 2.9891040325164795, + absoluteSum: 172419.046875, + positionChecksum: 84804.59375, + sampleIndices: [0, 104857, 209715, 314572, 419430, 524287], + samples: [ + 0.1790539026260376, -0.17156574130058289, -0.12500464916229248, + -0.21596695482730865, 0.10746431350708008, -0.5037912130355835, + ]), + tolerance: .float32) + } + } } diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift index 11ce7fffa..40a3feb4f 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 22 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 25 import Foundation import MLX @@ -568,4 +570,112 @@ struct GeneratedQuantizationTests { tolerance: .float32) } } + + @Test("gatherQuantizedMM/nvfp4/globalScale") + func test_gatherQuantizedMM_nvfp4_globalScale() throws { + // global scale is gpu only and needs indices + try withIntegrationState(seed: 48073) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 64], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([2, 0, 3] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, globalScale: globalScale) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 32], + dtype: .float32, + mean: -0.9819647669792175, + minimum: -165.96966552734375, + maximum: 144.31866455078125, + absoluteSum: 5297.81396484375, + positionChecksum: 3304.0732421875, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + -55.56662368774414, -11.178827285766602, -5.531567096710205, + 0.3299040198326111, -16.427204132080078, 57.386314392089844, + ]), + tolerance: .float32) + } + } + + @Test("gatherQuantizedMM/nvfp4/globalScale/transpose") + func test_gatherQuantizedMM_nvfp4_globalScale_transpose() throws { + try withIntegrationState(seed: 33918) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 32], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([2, 0, 3] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, transpose: false, mode: .nvfp4, + globalScale: globalScale) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 64], + dtype: .float32, + mean: -1.6645965576171875, + minimum: -119.096923828125, + maximum: 92.36067199707031, + absoluteSum: 7135.9306640625, + positionChecksum: 4085.1412760416665, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -43.466453552246094, 8.144684791564941, -3.2288098335266113, + 6.519178867340088, 6.189074516296387, 76.47562408447266, + ]), + tolerance: .float32) + } + } + + @Test("gatherQuantizedMM/nvfp4/globalScale/sorted") + func test_gatherQuantizedMM_nvfp4_globalScale_sorted() throws { + try withIntegrationState(seed: 24933) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 64], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([0, 2, 2] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, globalScale: globalScale, + sortedIndices: true) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 32], + dtype: .float32, + mean: -1.1882468461990356, + minimum: -89.71018981933594, + maximum: 76.8901138305664, + absoluteSum: 3490.58984375, + positionChecksum: 2165.3154296875, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + 2.895949125289917, -2.178328037261963, 52.326904296875, 13.967162132263184, + -27.429550170898438, -89.71018981933594, + ]), + tolerance: .float32) + } + } } diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift index bed682645..f5ca87e1e 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 49 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 69 import Foundation import MLX @@ -1169,4 +1171,440 @@ struct GeneratedShapeTests { tolerance: .float32) } } + + @Test("matrixTransposed") + func test_matrixTransposed() throws { + try withIntegrationState(seed: 94372) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matrixTransposed(a) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.33135515451431274, + minimum: -1.0816655158996582, + maximum: 1.7783619165420532, + absoluteSum: 8.355605125427246, + positionChecksum: 4.310164133707683, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.028264261782169342, 0.2258906066417694, 1.7783619165420532, + 0.21972355246543884, 1.0169203281402588, -0.8231726884841919, + ]), + tolerance: .float32) + } + } + + @Test("matrixTransposed/4d") + func test_matrixTransposed_4d() throws { + // only the last two axes move, unlike transposed(a) which reverses all of them + try withIntegrationState(seed: 89528) { + let a = MLXRandom.normal([2, 3, 4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matrixTransposed(a) + expectSummary( + result, + ArraySummary( + shape: [2, 3, 5, 4], + dtype: .float32, + mean: -0.06181011348962784, + minimum: -2.80637788772583, + maximum: 2.895658016204834, + absoluteSum: 101.28848266601562, + positionChecksum: 53.0994384765625, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + -1.9200016260147095, 0.419328898191452, -0.08306543529033661, + -1.4899338483810425, -1.7050701379776, 0.0025647925212979317, + ]), + tolerance: .float32) + } + } + + @Test("uniqued") + func test_uniqued() throws { + try withIntegrationState(seed: 57697) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 6) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .int32, + mean: 2.5, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 15.0, + positionChecksum: 11.666666666666666, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), + tolerance: .exact) + } + } + + @Test("uniqued/padded") + func test_uniqued_padded() throws { + // no fillValue: pins down the default padding + try withIntegrationState(seed: 24389) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.5, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 15.0, + positionChecksum: 7.0, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/fillValue") + func test_uniqued_fillValue() throws { + try withIntegrationState(seed: 96023) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, fillValue: MLXArray(-1)) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.100000023841858, + minimum: -1.0, + maximum: 5.0, + absoluteSum: 19.0, + positionChecksum: 10.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, -1.0, -1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/truncated") + func test_uniqued_truncated() throws { + // keeps the smallest `size` unique elements + try withIntegrationState(seed: 55506) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3) + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .int32, + mean: 1.0, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 3.0, + positionChecksum: 2.6666666666666665, + sampleIndices: [0, 1, 2], + samples: [0.0, 1.0, 2.0]), + tolerance: .exact) + } + } + + @Test("uniqued/float") + func test_uniqued_float() throws { + try withIntegrationState(seed: 21512) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.uniqued(a, size: 12) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: -0.4238990545272827, + minimum: -2.4105277061462402, + maximum: 1.042361855506897, + absoluteSum: 8.695816040039062, + positionChecksum: 3.2035303115844727, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.4105277061462402, -0.8730838298797607, -0.5980460047721863, + -0.14297938346862793, 0.14883598685264587, 1.042361855506897, + ]), + tolerance: .float32) + } + } + + @Test("uniqued/bool") + func test_uniqued_bool() throws { + // bool needs the cpu stream (builds on sort) + try withIntegrationState(seed: 40846) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.uniqued(a, size: 2, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [2], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0, 1], + samples: [0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/values") + func test_uniqued_returning_values() throws { + try withIntegrationState(seed: 21630) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued( + a, size: 10, fillValue: MLXArray(-1), returning: [.index, .inverse, .counts] + ).values + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.100000023841858, + minimum: -1.0, + maximum: 5.0, + absoluteSum: 19.0, + positionChecksum: 10.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, -1.0, -1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/index") + func test_uniqued_returning_index() throws { + try withIntegrationState(seed: 59710) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .index).index! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .uint32, + mean: 3.9000000953674316, + minimum: 0.0, + maximum: 10.0, + absoluteSum: 39.0, + positionChecksum: 20.2, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 7.0, 0.0, 10.0, 3.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/index/truncated") + func test_uniqued_returning_index_truncated() throws { + try withIntegrationState(seed: 97801) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .index).index! + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .uint32, + mean: 2.0, + minimum: 1.0, + maximum: 3.0, + absoluteSum: 6.0, + positionChecksum: 4.333333333333333, + sampleIndices: [0, 1, 2], + samples: [2.0, 1.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverse") + func test_uniqued_returning_inverse() throws { + try withIntegrationState(seed: 34157) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .inverse).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.899999976158142, + minimum: 0.0, + maximum: 4.0, + absoluteSum: 38.0, + positionChecksum: 22.35, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [2.0, 1.0, 0.0, 4.0, 2.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverse/truncated") + func test_uniqued_returning_inverse_truncated() throws { + try withIntegrationState(seed: 75343) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .inverse).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.5, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 30.0, + positionChecksum: 15.55, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [0.0, 1.0, 2.0, 2.0, 2.0, 2.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/counts") + func test_uniqued_returning_counts() throws { + try withIntegrationState(seed: 13539) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .counts).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 6.0, + absoluteSum: 20.0, + positionChecksum: 6.9, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 5.0, 5.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/counts/truncated") + func test_uniqued_returning_counts_truncated() throws { + try withIntegrationState(seed: 70850) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .counts).counts! + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .int32, + mean: 4.0, + minimum: 2.0, + maximum: 6.0, + absoluteSum: 12.0, + positionChecksum: 8.666666666666666, + sampleIndices: [0, 1, 2], + samples: [4.0, 2.0, 6.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/index") + func test_uniqued_returning_all_index() throws { + try withIntegrationState(seed: 80422) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).index! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .uint32, + mean: 6.700000286102295, + minimum: 0.0, + maximum: 12.0, + absoluteSum: 67.0, + positionChecksum: 37.1, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [8.0, 0.0, 3.0, 1.0, 8.0, 8.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/inverse") + func test_uniqued_returning_all_inverse() throws { + try withIntegrationState(seed: 44680) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.600000023841858, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 32.0, + positionChecksum: 14.9, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [1.0, 2.0, 2.0, 1.0, 1.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/counts") + func test_uniqued_returning_all_counts() throws { + try withIntegrationState(seed: 302) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 9.0, + absoluteSum: 20.0, + positionChecksum: 6.2, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [2.0, 9.0, 3.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/indexCounts") + func test_uniqued_returning_indexCounts() throws { + // skips .inverse: counts must be read from slot 2, not 3 + try withIntegrationState(seed: 95441) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .counts]).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 20.0, + positionChecksum: 7.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 5.0, 2.0, 4.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverseCounts") + func test_uniqued_returning_inverseCounts() throws { + try withIntegrationState(seed: 29675) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.inverse, .counts]).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 2.6500000953674316, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 53.0, + positionChecksum: 28.6, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [5.0, 0.0, 5.0, 2.0, 5.0, 4.0]), + tolerance: .exact) + } + } } diff --git a/Tests/MLXTests/ConcurrentLoadTests.swift b/Tests/MLXTests/ConcurrentLoadTests.swift new file mode 100644 index 000000000..a9e27a7bc --- /dev/null +++ b/Tests/MLXTests/ConcurrentLoadTests.swift @@ -0,0 +1,266 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import XCTest + +@testable import MLX + +final class ConcurrentLoadTests: XCTestCase { + + let temporaryPath = FileManager.default.temporaryDirectory.appending( + path: UUID().uuidString, + directoryHint: .isDirectory + ) + + override func setUpWithError() throws { + try FileManager.default.createDirectory( + at: temporaryPath, + withIntermediateDirectories: false + ) + } + + override func tearDownWithError() throws { + try FileManager.default.removeItem(at: temporaryPath) + } + + private func write( + arrays: [String: MLXArray], metadata: [String: String] = [:], name: String + ) throws -> URL { + let url = temporaryPath.appending(path: name, directoryHint: .notDirectory) + try save(arrays: arrays, metadata: metadata, url: url) + return url + } + + private func writeRawSafetensor( + header: String, payload: Data = Data(), name: String + ) throws -> URL { + let url = temporaryPath.appending(path: name, directoryHint: .notDirectory) + var headerLength = UInt64(header.utf8.count).littleEndian + var data = withUnsafeBytes(of: &headerLength) { Data($0) } + data.append(contentsOf: header.utf8) + data.append(payload) + try data.write(to: url) + return url + } + + func testLoadArraysMatchesSerialLoads() throws { + var urls = [URL]() + var expected = [String: MLXArray]() + for file in 0 ..< 3 { + var arrays = [String: MLXArray]() + // More roots than MLX's default BFS width exercises a larger evaluation tape. + for tensor in 0 ..< 32 { + let name = "file\(file).tensor\(tensor)" + arrays[name] = MLXArray(Int32(0) ..< 1024).reshaped(32, 32) + Int32(tensor) + } + urls.append(try write(arrays: arrays, name: "shard\(file).safetensors")) + expected.merge(arrays) { _, new in new } + } + + let loaded = try loadArrays(urls: urls) + + XCTAssertEqual(loaded.keys.sorted(), expected.keys.sorted()) + XCTAssertTrue(loaded.values.allSatisfy { isEvaluated($0) }) + for (name, array) in expected { + assertEqual(try XCTUnwrap(loaded[name]), array) + } + } + + func testLoadArraysSingleFile() throws { + let arrays = [ + "a": MLXArray(converting: [1.5, 2.5, 3.5]), + "b": MLX.ones([8, 8]), + ] + let url = try write(arrays: arrays, name: "single.safetensors") + + let loaded = try loadArrays(urls: [url]) + XCTAssertEqual(loaded.keys.sorted(), arrays.keys.sorted()) + for (name, array) in arrays { + assertEqual(try XCTUnwrap(loaded[name]), array) + } + } + + func testLoadArraysMaterializesBeforeReturning() throws { + let firstExpected = MLXArray(Int32(0) ..< 4096).reshaped(64, 64) + let secondExpected = MLXArray(Int32(4096) ..< 8192).reshaped(64, 64) + let firstURL = try write( + arrays: ["first": firstExpected], name: "first-materialized.safetensors") + let secondURL = try write( + arrays: ["second": secondExpected], name: "second-materialized.safetensors") + + let loaded = try loadArrays(urls: [firstURL, secondURL]) + for url in [firstURL, secondURL] { + let handle = try FileHandle(forWritingTo: url) + try handle.truncate(atOffset: 0) + try handle.close() + } + + // The checkedEval barrier covers every file. Neither returned array may attempt + // deferred I/O after both are truncated. + assertEqual(try XCTUnwrap(loaded["first"]), firstExpected) + assertEqual(try XCTUnwrap(loaded["second"]), secondExpected) + } + + func testLoadArraysEmpty() throws { + let loaded = try loadArrays(urls: []) + XCTAssertTrue(loaded.isEmpty) + } + + func testLoadArraysMetadataOnlyFile() throws { + let url = try writeRawSafetensor( + header: #"{"__metadata__":{"format":"test"}}"#, + name: "metadata-only.safetensors") + + let (arrays, metadata) = try loadArraysAndMetadata(urls: [url]) + XCTAssertTrue(arrays.isEmpty) + XCTAssertEqual(metadata, ["format": "test"]) + } + + func testLoadArraysFallbackMaterializesEveryTensor() throws { + // MLX accepts gaps, but our ordering parser rejects them and falls back. + let url = try writeRawSafetensor( + header: + #"{"a":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"b":{"dtype":"U8","shape":[1],"data_offsets":[2,3]}}"#, + payload: Data([7, 0, 9]), + name: "fallback.safetensors") + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + + let loaded = try loadArrays(urls: [url]) + XCTAssertTrue(loaded.values.allSatisfy { isEvaluated($0) }) + let handle = try FileHandle(forWritingTo: url) + try handle.truncate(atOffset: 0) + try handle.close() + + XCTAssertEqual(try XCTUnwrap(loaded["a"]).item(UInt8.self), 7) + XCTAssertEqual(try XCTUnwrap(loaded["b"]).item(UInt8.self), 9) + } + + func testDuplicateNameLaterFileWins() throws { + let first = try write(arrays: ["x": MLX.zeros([4])], name: "first.safetensors") + let second = try write(arrays: ["x": MLX.ones([4])], name: "second.safetensors") + + let loaded = try loadArrays(urls: [first, second]) + assertEqual(try XCTUnwrap(loaded["x"]), MLX.ones([4])) + } + + func testMetadataMergedInFileOrder() throws { + let first = try write( + arrays: ["a": MLX.zeros([2])], + metadata: ["shared": "first", "only-first": "1"], + name: "first.safetensors") + let second = try write( + arrays: ["b": MLX.ones([2])], + metadata: ["shared": "second", "only-second": "2"], + name: "second.safetensors") + + let (arrays, metadata) = try loadArraysAndMetadata(urls: [first, second]) + XCTAssertEqual(arrays.keys.sorted(), ["a", "b"]) + XCTAssertEqual( + metadata, ["shared": "second", "only-first": "1", "only-second": "2"]) + } + + func testMissingFileThrows() throws { + let good = try write(arrays: ["a": MLX.zeros([2])], name: "good.safetensors") + let missing = temporaryPath.appending( + path: "missing.safetensors", directoryHint: .notDirectory) + + XCTAssertThrowsError(try loadArrays(urls: [good, missing])) + } + + func testUnknownExtensionThrows() throws { + let bogus = temporaryPath.appending(path: "weights.bin", directoryHint: .notDirectory) + try Data([0, 1, 2, 3]).write(to: bogus) + + XCTAssertThrowsError(try loadArrays(urls: [bogus])) + } + + // MARK: - internals + + func testSafetensorNamesInFileOrder() throws { + let arrays = [ + "big": MLX.zeros([64, 64]), + "small": MLX.zeros([4]), + "medium": MLX.zeros([16, 16]), + ] + let url = try write(arrays: arrays, name: "spans.safetensors") + + let names = try safetensorNamesInFileOrder(url: url) + XCTAssertEqual(names.sorted(), ["big", "medium", "small"]) + } + + func testSafetensorNamesFollowOffsetsRatherThanNameOrHeaderOrder() throws { + let url = try writeRawSafetensor( + header: + #"{"middle":{"dtype":"U8","shape":[2],"data_offsets":[1,3]},"alpha":{"dtype":"U8","shape":[1],"data_offsets":[3,4]},"zeta":{"dtype":"U8","shape":[1],"data_offsets":[0,1]}}"#, + payload: Data([1, 2, 3, 4]), + name: "offset-order.safetensors") + + XCTAssertEqual(try safetensorNamesInFileOrder(url: url), ["zeta", "middle", "alpha"]) + } + + func testSafetensorNamesRejectsMalformedFile() throws { + let url = temporaryPath.appending(path: "bad.safetensors", directoryHint: .notDirectory) + try Data([1, 2, 3]).write(to: url) + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesRejectsOverflowingOffsets() throws { + let url = try writeRawSafetensor( + header: + #"{"bad":{"dtype":"U8","shape":[1],"data_offsets":[-9223372036854775808,9223372036854775807]}}"#, + name: "overflow.safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesRejectsNonintegerOffsets() throws { + for (name, offset) in [("fractional", "1.5"), ("boolean", "true")] { + let url = try writeRawSafetensor( + header: + #"{"bad":{"dtype":"U8","shape":[1],"data_offsets":[0,"# + + offset + #"]}}"#, + name: "\(name).safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + } + + func testSafetensorNamesRejectsNoncontiguousOffsets() throws { + let url = try writeRawSafetensor( + header: + #"{"a":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"b":{"dtype":"U8","shape":[1],"data_offsets":[2,3]}}"#, + payload: Data([0, 0, 0]), + name: "noncontiguous.safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesOrdersEmptyTensorBeforeSharedOffset() throws { + let url = try writeRawSafetensor( + header: + #"{"value":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"empty":{"dtype":"U8","shape":[0],"data_offsets":[0,0]}}"#, + payload: Data([42]), + name: "empty.safetensors") + + XCTAssertEqual(try safetensorNamesInFileOrder(url: url), ["empty", "value"]) + } + + func testEvaluationOrderReversesFilesAndOffsetsWithoutDroppingDuplicateNames() { + let first = ["shared": MLXArray(Int32(1)), "last": MLXArray(Int32(2))] + let second = ["shared": MLXArray(Int32(3)), "last": MLXArray(Int32(4))] + let roots = safetensorArraysInEvaluationOrder( + arrays: [first, second], names: [["shared", "last"], ["shared", "last"]]) + + // eval_impl pops its tape from the back: all four independent Load roots must + // be retained in reverse order, including the name overwritten by the merge. + XCTAssertEqual(roots.map { $0.item(Int32.self) }, [4, 3, 2, 1]) + } + + func testEvaluationOrderIncludesEveryArrayWhenPlanningFailsOrNamesDisagree() { + let arrays = ["a": MLXArray(Int32(1)), "b": MLXArray(Int32(2))] + for names: [String]? in [nil, [], ["a"], ["a", "missing"]] { + let roots = safetensorArraysInEvaluationOrder(arrays: [arrays], names: [names]) + XCTAssertEqual(roots.map { $0.item(Int32.self) }.sorted(), [1, 2]) + } + } +} diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift index 427e976e5..9df64346c 100644 --- a/Tests/MLXTests/FP8Tests.swift +++ b/Tests/MLXTests/FP8Tests.swift @@ -9,24 +9,29 @@ class FP8Tests: XCTestCase { override class func setUp() { } - func testFromFP8DecodesE4M3Bytes() { - // 0x7F and 0xFF are NaN in E4M3FN (mlx #4376); every other byte is finite. - let nanBytes: Set = [0x7F, 0xFF] - let allBytes = (0 ... 255).map { UInt8($0) } + func testFromFP8DecodesE4M3BytesPositive() { + // 127 and 255 are NaN and don't compare equal, so skip them + let allBytes = (0 ... 126).map { UInt8($0) } let bytes = MLXArray(allBytes) let decoded = fromFP8(bytes, dtype: .float32) XCTAssertEqual(decoded.dtype, .float32) - XCTAssertEqual(decoded.shape, [256]) - - let values = decoded.asArray(Float.self) - for byte in allBytes { - if nanBytes.contains(byte) { - XCTAssertTrue(values[Int(byte)].isNaN, "byte \(byte) should decode to NaN") - } else { - XCTAssertEqual(values[Int(byte)], e4m3Value(byte), "byte \(byte)") - } - } + XCTAssertEqual(decoded.shape, [127]) + + let expected = MLXArray(allBytes.map { e4m3Value($0) }) + assertEqual(decoded, expected) + } + + func testFromFP8DecodesE4M3BytesNegative() { + let allBytes = (128 ... 254).map { UInt8($0) } + let bytes = MLXArray(allBytes) + let decoded = fromFP8(bytes, dtype: .float32) + + XCTAssertEqual(decoded.dtype, .float32) + XCTAssertEqual(decoded.shape, [127]) + + let expected = MLXArray(allBytes.map { e4m3Value($0) }) + assertEqual(decoded, expected) } func testToFP8EncodesE4M3Bytes() { @@ -46,6 +51,7 @@ class FP8Tests: XCTestCase { } private func e4m3Value(_ value: UInt8) -> Float { + if (value & 0x7F) == 0x7F { return .nan } let sign: Float = (value & 0x80) == 0 ? 1 : -1 let exponent = Int((value >> 3) & 0x0F) let mantissa = Int(value & 0x07) diff --git a/Tests/MLXTests/LossTests.swift b/Tests/MLXTests/LossTests.swift index 8aba886af..566166570 100644 --- a/Tests/MLXTests/LossTests.swift +++ b/Tests/MLXTests/LossTests.swift @@ -30,6 +30,27 @@ class LossTests: XCTestCase { XCTAssertTrue(all(isNaN(loss2)).item()) } + func testCrossEntropyAxis() { + // The class axis need not be the last one. Expected values from Python mlx: + // nn.losses.cross_entropy(logits, targets, axis=0) -> [0.169846, 0.169846]. + let logits = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2]) + let targets = MLXArray([0, 1]) + let expected = MLXArray(converting: [0.169846, 0.169846]) + let loss = crossEntropy(logits: logits, targets: targets, axis: 0, reduction: .none) + XCTAssertEqual(loss.shape, [2]) + XCTAssertTrue(allClose(loss, expected, atol: 1e-5).item()) + } + + func testNLLLossAxis() { + // nn.losses.nll_loss(inputs, targets, axis=0) -> [-2, -2] + let inputs = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2]) + let targets = MLXArray([0, 1]) + let expected = MLXArray(converting: [-2.0, -2.0]) + let loss = nllLoss(inputs: inputs, targets: targets, axis: 0, reduction: .none) + XCTAssertEqual(loss.shape, [2]) + XCTAssertTrue(allClose(loss, expected).item()) + } + func testTripletLoss() { // This mainly tests that tripletLoss now works for arbitrary values of p diff --git a/Tests/MLXTests/MaterializedHazardTests.swift b/Tests/MLXTests/MaterializedHazardTests.swift index 5bc940211..644bb9c64 100644 --- a/Tests/MLXTests/MaterializedHazardTests.swift +++ b/Tests/MLXTests/MaterializedHazardTests.swift @@ -1,5 +1,6 @@ // Copyright © 2026 Apple Inc. +import Cmlx import Foundation @_spi(MaterializedModule) import MLXNN import Testing @@ -336,9 +337,121 @@ private func compileInput() -> MLXArray { MLXArray((0 ..< 8).map { Float($0) }, [2, 4]) } +// Access to the mutable state is handed between threads by the test's semaphores. +private final class ReboundCompileState: @unchecked Sendable { + let holder: ArrayHolder + let compiled: (MLXArray) -> MLXArray + var workerBefore: Float = 0 + var workerAfter: Float = 0 + + init() { + let holder = ArrayHolder(MLXArray(Float(1)).materialized()) + self.holder = holder + self.compiled = compile(inputs: [holder]) { (x: MLXArray) in x * holder.value } + } + + func value() -> Float { + compiled(MLXArray(Float(2))).item(Float.self) + } +} + +private final class ReleasedCompileState: @unchecked Sendable { + var compiled: CompiledFunction? + let functionID: UInt + var before: Float = 0 + var after: Float = 0 + + init() { + let scale = MLXArray(Float(2)).materialized() + let compiled = CompiledFunction(shapeless: false) { [$0[0] * scale] } + self.compiled = compiled + self.functionID = UInt(bitPattern: Unmanaged.passUnretained(compiled).toOpaque()) + } + + // Reuse the old address-derived ID without depending on allocator reuse. + func replacementValue() -> Float { + let result = withEvalLock { + let body = new_mlx_closure { [$0[0] * 3] } + defer { mlx_closure_free(body) } + var replacement = mlx_closure_new() + defer { mlx_closure_free(replacement) } + mlx_detail_compile(&replacement, body, functionID, false, [], 0) + let input = new_mlx_vector_array([MLXArray(Float(2))]) + defer { mlx_vector_array_free(input) } + var output = mlx_vector_array_new() + defer { mlx_vector_array_free(output) } + mlx_closure_apply(&output, replacement, input) + return mlx_vector_array_values(output)[0] + } + return result.item(Float.self) + } +} + @Suite("compile with rebound materialized state") struct CompileReboundMaterializedStateTests { + @Test func deinitErasesCacheOnTheThreadThatCompiled() { + let state = ReleasedCompileState() + let primed = DispatchSemaphore(value: 0) + let released = DispatchSemaphore(value: 0) + let finished = DispatchSemaphore(value: 0) + let worker = Thread { + state.before = state.compiled!.call([MLXArray(Float(2))])[0].item(Float.self) + primed.signal() + if released.wait(timeout: .now() + 30) == .success { + state.after = state.replacementValue() + } + finished.signal() + } + worker.start() + guard primed.wait(timeout: .now() + 30) == .success else { + released.signal() + Issue.record("worker did not populate its compilation cache") + return + } + state.compiled = nil + released.signal() + guard finished.wait(timeout: .now() + 30) == .success else { + Issue.record("worker did not finish the replacement call") + return + } + #expect(state.before == 4) + #expect(state.after == 6) + } + + @Test func rebindingInvalidatesEveryCallingThreadCache() { + let state = ReboundCompileState() + let primed = DispatchSemaphore(value: 0) + let rebound = DispatchSemaphore(value: 0) + let finished = DispatchSemaphore(value: 0) + + #expect(state.value() == 2) + let worker = Thread { + state.workerBefore = state.value() + primed.signal() + if rebound.wait(timeout: .now() + 30) == .success { + state.workerAfter = state.value() + } + finished.signal() + } + worker.start() + guard primed.wait(timeout: .now() + 30) == .success else { + rebound.signal() + Issue.record("worker did not populate its compilation cache") + return + } + + state.holder.value = MLXArray(Float(3)).materialized() + #expect(state.value() == 6) + rebound.signal() + guard finished.wait(timeout: .now() + 30) == .success else { + Issue.record("worker did not finish its cached call") + return + } + #expect(state.workerBefore == 2) + #expect(state.workerAfter == 6) + } + /// Control: plain modules rebound between calls are threaded through as /// inputs and the compiled function follows the rebinding. @Test func plainModuleRebound() { @@ -546,8 +659,40 @@ struct SendableTransformsMaterializedTests { let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in x * scale } - let input = compileInput() - #expect(compiled(input).allClose(input * 3).item(Bool.self)) + for offset in 0 ..< 3 { + let input = compileInput() + Float(offset) + #expect(compiled(input).allClose(input * 3).item(Bool.self)) + #expect(isEvaluated(scale)) + #expect(!hasGraph(scale)) + #expect(scale.item(Float.self) == 3) + } + } + + @Test func sendableCompileCapturingMaterializedArrayConcurrently() async { + let scale = MLXArray(Float(3)).materialized() + let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in + x * scale + } + + await withTaskGroup(of: Bool.self) { group in + for worker in 0 ..< 8 { + group.addTask { + for iteration in 0 ..< 25 { + let input = compileInput() + Float(worker * 25 + iteration) + if !compiled(input).allClose(input * 3).item(Bool.self) { + return false + } + } + return true + } + } + for await ok in group { + #expect(ok) + } + } + #expect(isEvaluated(scale)) + #expect(!hasGraph(scale)) + #expect(scale.item(Float.self) == 3) } /// `CustomFunction` closures may capture a MaterializedArray, including diff --git a/Tests/MLXTests/MemoryTests.swift b/Tests/MLXTests/MemoryTests.swift index 42485e76e..dc16ca9b7 100644 --- a/Tests/MLXTests/MemoryTests.swift +++ b/Tests/MLXTests/MemoryTests.swift @@ -27,4 +27,54 @@ class MemoryTests: XCTestCase { XCTAssertEqual(again.maxBufferSize, info.maxBufferSize) XCTAssertEqual(again.maxRecommendedWorkingSetSize, info.maxRecommendedWorkingSetSize) } + + // `Memory.bufferSize` reports the allocator length of each buffer, and the + // allocator may hand out a cached buffer up to ~2x the requested size, so the + // exact value depends on the allocation history of the process (and other + // tests running in parallel). These tests only check the properties that hold + // regardless of the cache. + + func testBufferSizeAtLeastLogicalSize() { + let a = MLXRandom.normal([64, 32]) + let b = MLXArray.zeros([10, 10], type: UInt8.self) + eval(a, b) + + XCTAssertGreaterThanOrEqual(Memory.bufferSize([a]), a.nbytes) + XCTAssertGreaterThanOrEqual(Memory.bufferSize([b]), b.nbytes) + } + + func testBufferSizeEmpty() { + XCTAssertEqual(Memory.bufferSize([MLXArray]()), 0) + } + + func testBufferSizeCountsDuplicatesOnce() { + let a = MLXRandom.normal([64, 32]) + eval(a) + + XCTAssertEqual(Memory.bufferSize([a, a, a]), Memory.bufferSize([a])) + } + + func testBufferSizeCountsViewsOnce() { + let a = MLXRandom.normal([64, 32]) + let slice = a[0 ..< 8] + let transposed = a.T + eval(a, slice, transposed) + + // views share the buffer of `a` + let size = Memory.bufferSize([a]) + XCTAssertEqual(Memory.bufferSize([a, slice]), size) + XCTAssertEqual(Memory.bufferSize([a, transposed]), size) + XCTAssertEqual(Memory.bufferSize([slice, transposed]), size) + } + + func testBufferSizeSumsDistinctBuffers() { + let a = MLXRandom.normal([64, 32]) + let b = MLXRandom.normal([5]) + let c = MLXRandom.randInt(low: 0, high: 4, [100], type: UInt32.self) + eval(a, b, c) + + XCTAssertEqual( + Memory.bufferSize([a, b, c]), + Memory.bufferSize([a]) + Memory.bufferSize([b]) + Memory.bufferSize([c])) + } } diff --git a/Tests/MLXTests/ModuleTests.swift b/Tests/MLXTests/ModuleTests.swift index 07a771b3d..c314addf8 100644 --- a/Tests/MLXTests/ModuleTests.swift +++ b/Tests/MLXTests/ModuleTests.swift @@ -444,6 +444,40 @@ class ModuleTests: XCTestCase { XCTAssertEqual(collect["interior2.children.1"], "Linear") } + func testUpdateParametersCallsChildOverride() throws { + // a child's update(parameters:verify:path:modulePath:) override runs + // when the update starts at an ancestor + class RecordingLinear: Linear { + var updateCount = 0 + + override func update( + parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [], + modulePath: [String] = [] + ) throws -> Self { + updateCount += 1 + return try super.update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath) + } + } + + class Parent: Module { + let child = RecordingLinear(2, 2) + let list = [RecordingLinear(2, 2)] + } + + let m = Parent() + let weight = MLXArray.ones([2, 2]) + m.update( + parameters: .unflattened([ + "child.weight": weight, "list.0.weight": weight, + ])) + + XCTAssertEqual(m.child.updateCount, 1) + XCTAssertEqual(m.list[0].updateCount, 1) + assertEqual(m.child.weight, weight) + assertEqual(m.list[0].weight, weight) + } + func testUpdateModuleSameType() throws { // set a single module with same type let m = newStructureModule() diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index 6bc06ff66..08d7d87c0 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -436,4 +436,233 @@ class OptimizerTests: XCTestCase { XCTAssertGreaterThan(abs(p2).sum().item(Float.self), 0) } + // MARK: - optimizer state checkpoints + + class StateBlock: Module { + let weight = MLXArray.zeros([3]) + } + + /// Nested parameters, a sibling vector, and an array of parameters. + /// The dotted and indexed keys are what a checkpoint has to round-trip. + class StateNet: Module { + let block = StateBlock() + let bias = MLXArray.zeros([3]) + let heads = [MLXArray.zeros([2]), MLXArray.zeros([2])] + } + + class RankNet: Module { + let vector = MLXArray.zeros([4]) + let matrix = MLXArray.zeros([3, 5]) + } + + func step(_ optimizer: Optimizer, _ model: Module, times: Int) { + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + for _ in 0 ..< times { + optimizer.update(model: model, gradients: gradients) + } + eval(model, optimizer) + } + + /// Value copy. `parameters()` returns the arrays the module will mutate in place. + func detached(_ parameters: ModuleParameters) -> ModuleParameters { + parameters.mapValues { MLXArray($0.asArray(Float.self), $0.shape) } + } + + func parameter(_ parameters: ModuleParameters, _ key: String) -> MLXArray { + parameters.flattened().first { $0.0 == key }!.1 + } + + func assertParametersEqual( + _ lhs: ModuleParameters, _ rhs: ModuleParameters, file: StaticString = #filePath, + line: UInt = #line + ) { + let left = lhs.flattened() + let right = rhs.flattened() + XCTAssertEqual(left.map(\.0), right.map(\.0), file: file, line: line) + for (lhs, rhs) in zip(left, right) { + assertEqual(lhs.1, rhs.1, file: file, line: line) + } + } + + /// Train `stepsBefore + stepsAfter` on one optimizer, and the same trajectory + /// again with the state saved through the flattened dictionary `save` writes. + func assertResumes( + _ make: () -> T, model: () -> Module, stepsBefore: Int = 3, stepsAfter: Int = 1, + file: StaticString = #filePath, line: UInt = #line + ) { + let reference = model() + step(make(), reference, times: stepsBefore + stepsAfter) + let continued = detached(reference.parameters()) + + let trained = model() + let optimizer = make() + step(optimizer, trained, times: stepsBefore) + let encoded = Dictionary(uniqueKeysWithValues: optimizer.state().flattened()) + let decoded = ModuleParameters.unflattened(encoded) + + let resumedModel = model() + resumedModel.update(parameters: detached(trained.parameters())) + let resumed = make() + resumed.update(parameters: decoded) + step(resumed, resumedModel, times: stepsAfter) + + assertParametersEqual( + continued, detached(resumedModel.parameters()), file: file, line: line) + } + + func testStateIsEmptyBeforeTheFirstStep() { + XCTAssertTrue(Adam(learningRate: 0.1).state().isEmpty) + XCTAssertTrue(SGD(learningRate: 0.1, momentum: 0.9).state().isEmpty) + } + + func testStateRoundTripSGD() { + assertResumes({ SGD(learningRate: 0.1, momentum: 0.9) }, model: { StateNet() }) + + // A fresh optimizer from the same weights is not the same trajectory. + // Otherwise the comparison above would pass with the buffers ignored. + let model = StateNet() + step(SGD(learningRate: 0.1, momentum: 0.9), model, times: 3) + let weights = detached(model.parameters()) + let continued = StateNet() + step(SGD(learningRate: 0.1, momentum: 0.9), continued, times: 4) + + let freshModel = StateNet() + freshModel.update(parameters: weights) + step(SGD(learningRate: 0.1, momentum: 0.9), freshModel, times: 1) + assertNotEqual( + parameter(detached(freshModel.parameters()), "bias"), + parameter(detached(continued.parameters()), "bias")) + } + + func testStateRoundTripAdam() { + assertResumes( + { Adam(learningRate: 0.1, biasCorrection: true) }, model: { StateNet() }) + + let model = StateNet() + let optimizer = Adam(learningRate: 0.1, biasCorrection: true) + step(optimizer, model, times: 3) + + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("block.weight.m")) + XCTAssertTrue(keys.contains("block.weight.v")) + XCTAssertTrue(keys.contains("block.weight.step")) + XCTAssertTrue(keys.contains("bias.m")) + XCTAssertTrue(keys.contains("heads.0.m")) + XCTAssertTrue(keys.contains("heads.1.v")) + XCTAssertTrue(keys.contains("heads.1.step")) + + for (key, step) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(step.item(Int32.self), 3, key) + } + + // `state()` on the existential is the real tree, not the empty default. + let existential: Optimizer = optimizer + XCTAssertEqual( + Set(existential.state().flattened().map(\.0)), keys) + + // Omitting `bias` leaves it alone and replaces `block.weight`. + let biasBefore = optimizer.state().flattened().filter { $0.0.hasPrefix("bias.") } + let zeroWeight = optimizer.state().flattened().filter { + $0.0.hasPrefix("block.weight.") + }.map { ($0.0, MLXArray.zeros(like: $0.1)) } + optimizer.update(parameters: ModuleParameters.unflattened(zeroWeight)) + let after = optimizer.state().flattened() + for (key, value) in biasBefore { + assertEqual(after.first { $0.0 == key }!.1, value) + } + for (key, _) in zeroWeight { + let value = after.first { $0.0 == key }!.1 + assertEqual(value, MLXArray.zeros(like: value)) + } + } + + func testStateRoundTripAdaDelta() { + assertResumes({ AdaDelta(learningRate: 0.1) }, model: { StateNet() }) + } + + func testStateRoundTripAdamax() { + assertResumes({ Adamax(learningRate: 0.1) }, model: { StateNet() }) + + let model = TwoParameterModel() + let optimizer = Adamax(learningRate: 0.1) + step(optimizer, model, times: 1) + XCTAssertEqual( + Set(optimizer.state().flattened().map(\.0)), + ["bias.m", "bias.v", "weight.m", "weight.v"]) + } + + func testStateRoundTripAdafactor() { + assertResumes( + { Adafactor(learningRate: 0.1, relativeStep: false) }, model: { RankNet() }) + + let model = RankNet() + let optimizer = Adafactor(learningRate: 0.1, beta1: 0.9, relativeStep: false) + step(optimizer, model, times: 2) + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("vector.expAvgSq")) + XCTAssertTrue(keys.contains("vector.expAvg")) + XCTAssertTrue(keys.contains("vector.step")) + XCTAssertTrue(keys.contains("matrix.expAvgSqRow")) + XCTAssertTrue(keys.contains("matrix.expAvgSqCol")) + XCTAssertFalse(keys.contains("matrix.expAvgSq")) + for (key, step) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(step.item(Int32.self), 2, key) + } + + assertResumes( + { Adafactor(learningRate: 0.1, beta1: 0.9, relativeStep: false) }, + model: { RankNet() }) + } + + /// `step` changes beta2 on every update. compile only writes back arrays + /// from `innerState()`, so a checkpoint after a compiled run must see the + /// same count an eager run would. + func testAdafactorCompiledStepAdvances() { + func parameters(compiled: Bool) -> ModuleParameters { + let model = RankNet() + let optimizer = Adafactor(learningRate: 0.1, warmupInit: true) + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + let token = MLXArray(0) + func once(_ token: MLXArray) -> MLXArray { + optimizer.update(model: model, gradients: gradients) + return token + } + let step = + compiled + ? compile(inputs: [model, optimizer], outputs: [model, optimizer], once) + : once + for _ in 0 ..< 4 { + _ = step(token) + } + eval(model, optimizer) + for (key, value) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(value.item(Int32.self), 4, key) + } + return detached(model.parameters()) + } + + assertParametersEqual(parameters(compiled: false), parameters(compiled: true)) + } + + func testStateRoundTripMultiOptimizer() { + func make() -> MultiOptimizer { + MultiOptimizer( + optimizers: [ + SGD(learningRate: 1.0, momentum: 0.5), + Adam(learningRate: 0.1, biasCorrection: true), + ], + filters: [{ key, _ in key == "bias" }]) + } + assertResumes(make, model: { TwoParameterModel() }) + + let model = TwoParameterModel() + let optimizer = make() + step(optimizer, model, times: 2) + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("states.0.bias")) + XCTAssertTrue(keys.contains("states.1.weight.m")) + XCTAssertTrue(keys.contains("states.1.weight.v")) + XCTAssertTrue(keys.contains("states.1.weight.step")) + } + } diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 0ef8e1d55..273757f7d 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -8,6 +8,44 @@ import MLX import XCTest +// All access to the recorded callbacks is protected by the lock. +private final class ProgressRecorder: @unchecked Sendable { + private let lock = NSLock() + private var progress = [LoadProgress]() + + func record(_ progress: LoadProgress) { + lock.withLock { + self.progress.append(progress) + } + } + + var reported: [LoadProgress] { + lock.withLock { progress } + } + + var values: [Double] { + reported.map { $0.fractionCompleted } + } + + /// Fractions reported for a single file, in order. + func values(for url: URL) -> [Double] { + reported.filter { $0.url == url }.map { $0.fractionCompleted } + } + + /// Aggregate fraction across every file seen, by bytes. + var aggregateFraction: Double { + var completed = [URL: Int64]() + var total = [URL: Int64]() + for progress in reported { + completed[progress.url] = progress.completedUnitCount + total[progress.url] = progress.totalUnitCount + } + let totalBytes = total.values.reduce(0, +) + guard totalBytes > 0 else { return 0 } + return Double(completed.values.reduce(0, +)) / Double(totalBytes) + } +} + final class SaveTests: XCTestCase { let temporaryPath = FileManager.default.temporaryDirectory.appending( @@ -46,6 +84,289 @@ final class SaveTests: XCTestCase { assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) } + public func testLoadArraysProgressReportsThroughEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "arrays.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays( + url: safetensorsPath + ) { @Sendable progress in + recorder.record(progress) + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + } + } + + public func testLoadArraysProgressFailsOnTruncatedTensorData() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + var data = try Data(contentsOf: safetensorsPath) + data.removeLast(32) + try data.write(to: safetensorsPath) + + // A truncated file has to be reported either eagerly, while the header is + // parsed (mlx >= 0.32.1 validates the tensor data offsets against the size of + // the file), or lazily, when the arrays are evaluated and the short read is + // turned into an error (ml-explore/mlx-c#130). + do { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } + try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("a truncated safetensors file must not load successfully") + } catch { + // expected + } + } + } + + /// mlx validates the tensor data offsets against the size of the file while it parses + /// the header, so a file that is _already_ truncated fails eagerly. A file that is + /// truncated after the header is parsed can only be caught when the lazy arrays are + /// evaluated and the read comes up short -- `mlx_io_reader` turns that into an error + /// rather than leaving the destination buffer uninitialized (ml-explore/mlx-c#130). + public func testLoadFailsWhenFileIsTruncatedBeforeEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated-late.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + // the header is parsed and validated against the size of the file here, but + // the tensor data is only read once the lazy arrays are evaluated + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { + @Sendable in recorder.record($0) + } + let size = try XCTUnwrap(recorder.reported.first?.totalUnitCount) + + // ... so truncate the file out from under them. `truncate()` shortens the + // inode the reader already has open, unlike a rewrite which may replace it. + XCTAssertEqual( + truncate(safetensorsPath.path(percentEncoded: false), off_t(size - 32)), 0) + + do { + try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("evaluating arrays read from a truncated file must fail") + } catch { + // expected + } + + // the bytes that were read are still accounted for, monotonically, and the + // aggregate stops short of the original size of the file + let fractions = recorder.values + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(fractions.first, 0) + XCTAssertLessThan(try XCTUnwrap(fractions.last), 1) + } + } + + public func testScopedLoadProgressHandler() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "scoped.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + + // note: the plain loadArrays(url:) -- no progress handler passed at the call site + let loadedArrays = try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + return loadedArrays + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(Set(recorder.reported.map(\.url)), [safetensorsPath]) + } + } + + public func testScopedLoadProgressHandlerIsScoped() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "unscoped.safetensors", + directoryHint: .notDirectory + ) + try MLX.save(arrays: ["foo": MLX.ones([128, 128])], url: safetensorsPath) + + let recorder = ProgressRecorder() + withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + } + + // outside the scope nothing is reported + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + + XCTAssertTrue(recorder.reported.isEmpty) + } + } + + public func testScopedLoadProgressAggregatesAcrossFiles() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let shards = try (0 ..< 3).map { index -> URL in + let url = temporaryPath.appending( + path: "shard-\(index).safetensors", + directoryHint: .notDirectory + ) + try MLX.save(arrays: ["w\(index)": MLX.ones([64, 128])], url: url) + return url + } + + let recorder = ProgressRecorder() + try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + // this mimics a model loader: several shards loaded lazily, then evaluated + var weights = [String: MLXArray]() + for url in shards { + let (w, _) = try MLX.loadArraysAndMetadata(url: url) + weights.merge(w) { _, new in new } + } + MLX.eval(Array(weights.values)) + } + + XCTAssertEqual(Set(recorder.reported.map(\.url)), Set(shards)) + for url in shards { + XCTAssertEqual(recorder.values(for: url).last, 1) + } + XCTAssertEqual(recorder.aggregateFraction, 1, accuracy: 1e-9) + } + } + + public func testProgressChunksLargeTensorReads() throws { + let url = temporaryPath.appending(path: "large.safetensors") + let array = MLXArray(Array(repeating: Float(3), count: 3 * 1024 * 1024)) + try MLX.save(arrays: ["large": array], metadata: ["format": "test"], url: url) + + let recorder = ProgressRecorder() + let (loaded, metadata) = try MLX.loadArraysAndMetadata(url: url) { + recorder.record($0) + } + let headerBytes = try XCTUnwrap(recorder.reported.last?.completedUnitCount) + XCTAssertLessThan(headerBytes, Int64(array.nbytes)) + XCTAssertEqual(metadata, ["format": "test"]) + XCTAssertEqual( + try XCTUnwrap(loaded["large"]).asArray(Float.self), array.asArray(Float.self)) + + let reports = recorder.reported + XCTAssertEqual(reports.first?.completedUnitCount, 0) + XCTAssertEqual(reports.last?.fractionCompleted, 1) + XCTAssertEqual( + reports.map(\.completedUnitCount), reports.map(\.completedUnitCount).sorted()) + let tensorReports = reports.filter { $0.completedUnitCount > headerBytes } + XCTAssertGreaterThanOrEqual(tensorReports.count, 3) + for (previous, current) in zip(reports, reports.dropFirst()) { + XCTAssertLessThanOrEqual( + current.completedUnitCount - previous.completedUnitCount, 4 * 1024 * 1024) + } + } + + public func testProgressRetainsHandlerAndSkipsUnevaluatedArrays() throws { + let url = temporaryPath.appending(path: "partial.safetensors") + let values = Array(repeating: Float(7), count: 1024) + try MLX.save(arrays: ["used": MLXArray(values), "unused": MLXArray(values)], url: url) + + let recorder = ProgressRecorder() + let arrays = try withLoadProgressHandler({ recorder.record($0) }) { + try MLX.loadArrays(url: url) + } + let headerBytes = try XCTUnwrap(recorder.reported.last?.completedUnitCount) + XCTAssertEqual(try XCTUnwrap(arrays["used"]).asArray(Float.self), values) + let partial = try XCTUnwrap(recorder.reported.last) + XCTAssertEqual( + partial.completedUnitCount, headerBytes + Int64(values.count * MemoryLayout.size) + ) + XCTAssertLessThan(partial.fractionCompleted, 1) + XCTAssertEqual(try XCTUnwrap(arrays["unused"]).asArray(Float.self), values) + XCTAssertEqual(recorder.reported.last?.fractionCompleted, 1) + } + + public func testNestedProgressHandlersRestoreAfterThrowAndExplicitHandlerWins() throws { + enum TestError: Error { case expected } + let url = temporaryPath.appending(path: "nested.safetensors") + try MLX.save(arrays: ["weight": MLXArray([Float(1), 2])], url: url) + let outer = ProgressRecorder() + let inner = ProgressRecorder() + let explicit = ProgressRecorder() + + try withLoadProgressHandler({ outer.record($0) }) { + XCTAssertThrowsError( + try withLoadProgressHandler({ inner.record($0) }) { + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + throw TestError.expected + }) + XCTAssertTrue(outer.reported.isEmpty) + XCTAssertEqual(inner.reported.last?.fractionCompleted, 1) + + let explicitlyLoaded = try MLX.loadArrays(url: url) { explicit.record($0) } + try checkedEval(Array(explicitlyLoaded.values) as [Any]) + XCTAssertTrue(outer.reported.isEmpty) + XCTAssertEqual(explicit.reported.last?.fractionCompleted, 1) + + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + } + XCTAssertEqual(outer.reported.last?.fractionCompleted, 1) + } + + public func testAsyncProgressScopeInheritsHandler() async throws { + let url = temporaryPath.appending(path: "async.safetensors") + try MLX.save(arrays: ["weight": MLXArray([Float(1), 2])], url: url) + let recorder = ProgressRecorder() + try await withLoadProgressHandler({ recorder.record($0) }) { + try await Task { + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + }.value + } + XCTAssertEqual(recorder.reported.first?.fractionCompleted, 0) + XCTAssertEqual(recorder.reported.last?.fractionCompleted, 1) + } + public func testSaveArray() throws { // single array npy file let path = temporaryPath.appending( @@ -111,4 +432,27 @@ final class SaveTests: XCTestCase { assertEqual(try XCTUnwrap(loaded["big"]), try XCTUnwrap(arrays["big"])) } + /// A truncated in-memory buffer must be reported as an error. The in-memory reader + /// reports a read that runs past the end of the buffer as zero bytes so that + /// `mlx_io_reader` throws instead of leaving the destination uninitialized + /// (ml-explore/mlx-c#130). + public func testLoadFromTruncatedDataFails() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let arrays: [String: MLXArray] = [ + "big": MLX.ones([64, 64]) + ] + + var data = try saveToData(arrays: arrays) + data.removeLast(32) + + do { + let loaded = try loadArrays(data: data) + try checkedEval(Array(loaded.values) as [Any]) + XCTFail("a truncated safetensors buffer must not load successfully") + } catch { + // expected + } + } + } + } diff --git a/cmake/mlx.patch b/cmake/mlx.patch deleted file mode 100644 index de99b7ebc..000000000 --- a/cmake/mlx.patch +++ /dev/null @@ -1,34 +0,0 @@ -diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h -index f8a234ee6..3b7a79f27 100644 ---- a/mlx/backend/cuda/cuda_utils.h -+++ b/mlx/backend/cuda/cuda_utils.h -@@ -30,7 +30,11 @@ class CudaHandle { - if (cudaPeekAtLastError() != cudaSuccess) { - return; - } -- reset(); -+ // Inline reset() without throwing -- handle shutdown cases. -+ if (handle_ != nullptr) { -+ Destroy(handle_); -+ handle_ = nullptr; -+ } - } - - CudaHandle(const CudaHandle&) = delete; -diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp -index 472b3d99f..e7d8f2620 100644 ---- a/mlx/backend/cuda/device.cpp -+++ b/mlx/backend/cuda/device.cpp -@@ -617,8 +617,10 @@ std::unordered_map& get_command_encoders() { - } - - std::unordered_map& get_global_command_encoders() { -- static std::unordered_map encoders; -- return encoders; -+ // encoders are leaked intentionally as they would synchronize on process -+ // shutdown -+ static auto* encoders = new std::unordered_map(); -+ return *encoders; - } - - } // namespace mlx::core::cu diff --git a/tools/fix-metal-includes.sh b/tools/fix-metal-includes.sh index 3493d8470..c32b82014 100755 --- a/tools/fix-metal-includes.sh +++ b/tools/fix-metal-includes.sh @@ -30,6 +30,9 @@ rms_norm.metal \ rope.metal \ searchsorted.metal \ scaled_dot_product_attention.metal \ +scaled_dot_product_attention_vjp.metal \ +gated_delta_update.metal \ +gated_delta_update_vjp.metal \ steel/attn/kernels/steel_attention.metal" # We fixup all the header files AND the listed kernel files diff --git a/tools/integration_tests/README.md b/tools/integration_tests/README.md index 77eec885d..567fc2f2b 100644 --- a/tools/integration_tests/README.md +++ b/tools/integration_tests/README.md @@ -35,14 +35,14 @@ pair (`MAINTENANCE.md` step 10). | `Elementwise` | unary math, predicates, dtype conversions, complex parts | | `Binary` | binary ops, operators, comparisons, logical/bitwise, matmul family | | `Reduction` | reductions, cumulative ops, argMin/argMax, softmax | -| `Shape` | shape/structure ops, sort/partition, take/put, trace, hadamard | +| `Shape` | shape/structure ops, sort/partition, take/put, trace, hadamard, `matrixTransposed`, `uniqued` | | `Indexing` | subscript forms (ranges, `.ellipsis`, `.newAxis`, strides, arrays) | | `Factory` | `arange`, `linspace`, `zeros`/`ones`/`full`, `eye`, windows | | `Convolution` | `convolve`, `conv1d/2d/3d`, `convTransposed*`, `convGeneral` | -| `Quantization` | `quantized`/`dequantized`/`quantizedMM` per bits, gather/masked MM, fp8 | +| `Quantization` | `quantized`/`dequantized`/`quantizedMM` per bits, gather/masked MM, fp8, nvfp4 `globalScale` | | `FFT` | the whole `mx.fft` surface, including complex results | | `Linalg` | single-value linalg (`inv`, `cholesky`, `solve`, `det`, ...) | -| `Fast` | `rmsNorm`, `layerNorm`, `RoPE`, `scaledDotProductAttention` | +| `Fast` | `rmsNorm`, `layerNorm`, `RoPE`, `scaledDotProductAttention`, `gatedDeltaUpdate` | | `Random` | `gumbel`, `laplace`, `truncatedNormal`, `categorical`, `permutation`, ... | | `Activations` | the `MLXNN` activation free functions | | `Losses` | the `MLXNN` loss functions, per `reduction` | @@ -381,6 +381,11 @@ and exits before generating. bugs -- `MLX.nanToNum` defaulted `posInf`/`negInf` to `0` where python uses the dtype max/min (contradicting its own doc comment), and `MLX.tensordot` defaulted `axes` to `1` where python/numpy use `2`. Both are fixed. +- `Memory.bufferSize` / `mx.get_array_buffer_size` cannot be compared against + python: it reports the allocator length of each buffer, and the buffer cache + reuses any cached buffer smaller than `min(2 * size, size + 2 * page)`, so the + value depends on the allocation history of the process. It has property tests + in `Tests/MLXTests/MemoryTests.swift` instead. - `MLXRandom.laplace` requires `dtype:` where python defaults it, so the case spells it out. - `linspace` used to take its dtype from the literal type (`Int` -> `int64`, diff --git a/tools/integration_tests/cases.py b/tools/integration_tests/cases.py index be7ad47da..b7beb6d08 100644 --- a/tools/integration_tests/cases.py +++ b/tools/integration_tests/cases.py @@ -977,6 +977,132 @@ def binary(name: str, py_name: str, swift_name: str, lhs=None, rhs=None) -> None "mx.hadamard_transform(a, scale=0.25)", "MLX.hadamardTransform(a, scale: 0.25)", ) +case( + "matrixTransposed", + "Shape", + {"a": Normal(S)}, + "mx.matrix_transpose(a)", + "MLX.matrixTransposed(a)", +) +case( + "matrixTransposed/4d", + "Shape", + {"a": Normal((2, 3, 4, 5))}, + "mx.matrix_transpose(a)", + "MLX.matrixTransposed(a)", + note="only the last two axes move, unlike transposed(a) which reverses all of them", +) + +# `uniqued` has a fixed output `size`; inputs draw from a small range so there +# are repeats (20 draws from 6 values). Multiple return values are compared +# one field at a time; the combined `returning:` cases check that the fields +# are read back in the right order. +UNIQUE_A = RandInt(0, 6, (4, 5)) +case( + "uniqued", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=6)", + "MLX.uniqued(a, size: 6)", +) +case( + "uniqued/padded", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10)", + "MLX.uniqued(a, size: 10)", + note="no fillValue: pins down the default padding", +) +case( + "uniqued/fillValue", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, fill_value=mx.array(-1))", + "MLX.uniqued(a, size: 10, fillValue: MLXArray(-1))", +) +case( + "uniqued/truncated", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=3)", + "MLX.uniqued(a, size: 3)", + note="keeps the smallest `size` unique elements", +) +case( + "uniqued/float", + "Shape", + {"a": Normal(S)}, + "mx.unique(a, size=12)", + "MLX.uniqued(a, size: 12)", +) +case( + "uniqued/bool", + "Shape", + {"a": Bernoulli(S)}, + "mx.unique(a, size=2, stream=mx.cpu)", + "MLX.uniqued(a, size: 2, stream: .cpu)", + note="bool needs the cpu stream (builds on sort)", +) +case( + "uniqued/returning/values", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, " + "return_counts=True, fill_value=mx.array(-1))[0]", + "MLX.uniqued(a, size: 10, fillValue: MLXArray(-1), " + "returning: [.index, .inverse, .counts]).values", +) +for _field, _flag in (("index", "return_index"), ("inverse", "return_inverse"), ("counts", "return_counts")): + case( + f"uniqued/returning/{_field}", + "Shape", + {"a": UNIQUE_A}, + f"mx.unique(a, size=10, {_flag}=True)[1]", + f"MLX.uniqued(a, size: 10, returning: .{_field}).{_field}!", + ) + case( + f"uniqued/returning/{_field}/truncated", + "Shape", + {"a": UNIQUE_A}, + f"mx.unique(a, size=3, {_flag}=True)[1]", + f"MLX.uniqued(a, size: 3, returning: .{_field}).{_field}!", + ) +case( + "uniqued/returning/all/index", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[1]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).index!", +) +case( + "uniqued/returning/all/inverse", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[2]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).inverse!", +) +case( + "uniqued/returning/all/counts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[3]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).counts!", +) +case( + "uniqued/returning/indexCounts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_counts=True)[2]", + "MLX.uniqued(a, size: 10, returning: [.index, .counts]).counts!", + note="skips .inverse: counts must be read from slot 2, not 3", +) +case( + "uniqued/returning/inverseCounts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_inverse=True, return_counts=True)[1]", + "MLX.uniqued(a, size: 10, returning: [.inverse, .counts]).inverse!", +) # ------------------------------------------------------------------- factory # @@ -1335,6 +1461,63 @@ def binary(name: str, py_name: str, swift_name: str, lhs=None, rhs=None) -> None "MLX.fromFP8(encoded, dtype: .float32)", ) +# nvfp4 with one global scale per expert: `quantize` takes a single scale, so +# each expert is quantized on its own and the results stacked. The experts get +# different magnitudes so a scale applied to the wrong expert is visible. +NVFP4_EXPERTS = { + "w": Expression( + py="mx.random.normal((4, 32, 64)) * mx.array([1.0, 2.0, 4.0, 8.0]).reshape(4, 1, 1)", + swift="MLXRandom.normal([4, 32, 64]) * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1)", + ), + "globalScale": Expression( + py="mx.max(mx.abs(w), axis=(1, 2))", + swift="MLX.abs(w).max(axes: [1, 2])", + ), + "q": Expression( + py='[mx.quantize(w[e], mode="nvfp4", global_scale=globalScale[e]) for e in range(4)]', + swift="(0 ..< 4).map { MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) }", + array=False, + ), + "wq": Expression(py="mx.stack([p[0] for p in q])", swift="MLX.stacked(q.map { $0.wq })"), + "scales": Expression( + py="mx.stack([p[1] for p in q])", swift="MLX.stacked(q.map { $0.scales })" + ), + "x": Normal((3, 2, 64)), + "indices": Expression(py="mx.array([2, 0, 3], dtype=mx.uint32)", swift="MLXArray([2, 0, 3] as [UInt32])"), +} +case( + "gatherQuantizedMM/nvfp4/globalScale", + "Quantization", + NVFP4_EXPERTS, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, mode="nvfp4", global_scale=globalScale)', + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, " + "globalScale: globalScale)", + note="global scale is gpu only and needs indices", +) +case( + "gatherQuantizedMM/nvfp4/globalScale/transpose", + "Quantization", + {**NVFP4_EXPERTS, "x": Normal((3, 2, 32))}, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, transpose=False, mode="nvfp4", ' + "global_scale=globalScale)", + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, transpose: false, " + "mode: .nvfp4, globalScale: globalScale)", +) +case( + "gatherQuantizedMM/nvfp4/globalScale/sorted", + "Quantization", + { + **NVFP4_EXPERTS, + "indices": Expression( + py="mx.array([0, 2, 2], dtype=mx.uint32)", swift="MLXArray([0, 2, 2] as [UInt32])" + ), + }, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, mode="nvfp4", global_scale=globalScale, ' + "sorted_indices=True)", + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, " + "globalScale: globalScale, sortedIndices: true)", +) + # ----------------------------------------------------------------------- fft FFT_1D = Normal((100,)) @@ -1643,6 +1826,78 @@ def fft_case(name: str, py: str, swift: str, *, complex_input=True, shape=(100,) "scale: 0.35, mask: mask)", ) + +def gated_delta_inputs(B, T, Hk, Hv, Dk, Dv, *, initial_state=True, mask=None) -> dict: + """inputs for `gated_delta_update`; small `keys` keep the delta rule stable""" + inputs = { + "queries": Normal((B, T, Hk, Dk)), + "keys": Uniform((B, T, Hk, Dk), low=-0.1, high=0.1), + "values": Normal((B, T, Hv, Dv)), + "gamma": Uniform((B, T, Hv), low=0.5, high=1.0), + "beta": Uniform((B, T, Hv), low=0.1, high=0.9), + } + if initial_state: + inputs["initialState"] = Normal((B, Hv, Dv, Dk)) + if mask is not None: + inputs["mask"] = mask + return inputs + + +# (name, dims, extra inputs, python kwargs, swift arguments, note) +GATED_DELTA = [ + # the generic (ops) path: any shape the metal kernel does not support + ("fallback", (2, 5, 2, 4, 8, 8), {}, "", "", "Hv != Hk, shape not supported by the kernel"), + ( + "fallback/initialState", + (2, 5, 2, 4, 8, 8), + {"initial_state": True}, + ", initial_state=initialState", + ", initialState: initialState", + None, + ), + ( + "fallback/mask", + (2, 5, 2, 4, 8, 8), + { + "initial_state": True, + "mask": Expression( + py="mx.array([[True] * 5, [True, True, True, False, False]])", + swift="MLXArray([true, true, true, true, true, true, true, true, false, false], [2, 5])", + ), + }, + ", initial_state=initialState, mask=mask", + ", initialState: initialState, mask: mask", + "masked steps leave the state unchanged and output zeros", + ), + # Dk == Dv == 128 with these head counts runs the metal kernel + ( + "kernel", + (1, 3, 16, 32, 128, 128), + {"initial_state": True}, + ", initial_state=initialState", + ", initialState: initialState", + "a shape the metal kernel supports", + ), +] +for _name, _dims, _extra, _py_kwargs, _swift_args, _note in GATED_DELTA: + _extra = {"initial_state": False, **_extra} + for _index, _output in ((0, "output"), (1, "state")): # python index, Swift label + case( + f"gatedDeltaUpdate/{_name}/{_output}", + "Fast", + gated_delta_inputs(*_dims, **_extra), + f"mx.fast.gated_delta_update(queries, keys, values, gamma, beta{_py_kwargs})[{_index}]", + "MLXFast.gatedDeltaUpdate(queries: queries, keys: keys, values: values, " + f"gamma: gamma, beta: beta{_swift_args}).{_output}", + note=_note, + ) + +# `Memory.bufferSize` is deliberately not here: it reports the allocator length +# of each buffer, and the allocator hands out cached buffers up to ~2x the +# requested size, so the value depends on what the process freed earlier (the +# generator runs every case in one warm process). It is covered by property +# tests in Tests/MLXTests/MemoryTests.swift instead. + # -------------------------------------------------------------------- random case("gumbel", "Random", {}, "mx.random.gumbel([4, 3])", "MLXRandom.gumbel([4, 3])") diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 877861947..8e7993edc 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -357,6 +357,7 @@ "examples/example-metal-kernel.c", "examples/example-ops.c", "examples/example-safe-tensors.c", + "examples/example-stream-tls.c", "examples/example-thread-stream.c", examples/example.c, LICENSE, @@ -594,6 +595,16 @@ docs/src/usage/saving_and_loading.rst, docs/src/usage/unified_memory.rst, docs/src/usage/using_streams.rst, + examples/cmake_extension/axpby/axpby.cpp, + examples/cmake_extension/axpby/axpby.h, + examples/cmake_extension/axpby/axpby.metal, + examples/cmake_extension/bindings.cpp, + examples/cmake_extension/CMakeLists.txt, + examples/cmake_extension/pyproject.toml, + examples/cmake_extension/README.md, + examples/cmake_extension/requirements.txt, + examples/cmake_extension/setup.py, + examples/cmake_extension/test.py, examples/cmake_project/CMakeLists.txt, examples/cmake_project/example.cpp, examples/cmake_project/README.md, @@ -610,17 +621,13 @@ examples/export/README.md, examples/export/train_mlp.cpp, examples/export/train_mlp.py, - examples/extensions/axpby/axpby.cpp, - examples/extensions/axpby/axpby.h, - examples/extensions/axpby/axpby.metal, - examples/extensions/bindings.cpp, - examples/extensions/CMakeLists.txt, - examples/extensions/mlx_sample_extensions/__init__.py, - examples/extensions/pyproject.toml, - examples/extensions/README.md, - examples/extensions/requirements.txt, - examples/extensions/setup.py, - examples/extensions/test.py, + examples/metal_extension/activation.cpp, + examples/metal_extension/pyproject.toml, + examples/metal_extension/README.md, + examples/metal_extension/requirements.txt, + examples/metal_extension/setup.py, + examples/metal_extension/silu_and_mul.metal, + examples/metal_extension/test.py, examples/python/distributed_data_parallel.py, examples/python/distributed_tensor_parallel.py, examples/python/linear_regression.py, @@ -694,6 +701,7 @@ mlx/backend/cuda/binary/CMakeLists.txt, mlx/backend/cuda/binary/divide.cu, mlx/backend/cuda/binary/equal.cu, + mlx/backend/cuda/binary/floor_divide.cu, mlx/backend/cuda/binary/greater_equal.cu, mlx/backend/cuda/binary/greater.cu, mlx/backend/cuda/binary/less_equal.cu, @@ -928,8 +936,13 @@ mlx/backend/metal/kernels/fp_quantized.metal, mlx/backend/metal/kernels/fp4.h, mlx/backend/metal/kernels/fp8.h, + mlx/backend/metal/kernels/gated_delta_nax_ops.h, + mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h, + mlx/backend/metal/kernels/gated_delta_update_nax_vjp.metal, mlx/backend/metal/kernels/gated_delta_update_nax.h, mlx/backend/metal/kernels/gated_delta_update_nax.metal, + mlx/backend/metal/kernels/gated_delta_update_vjp.h, + mlx/backend/metal/kernels/gated_delta_update_vjp.metal, mlx/backend/metal/kernels/gated_delta_update.h, mlx/backend/metal/kernels/gated_delta_update.metal, mlx/backend/metal/kernels/gather_mm_offsets.metal, @@ -966,9 +979,11 @@ mlx/backend/metal/kernels/reduction/reduce_row.h, mlx/backend/metal/kernels/rms_norm.metal, mlx/backend/metal/kernels/rope.metal, + mlx/backend/metal/kernels/scaled_dot_product_attention_vjp.metal, mlx/backend/metal/kernels/scaled_dot_product_attention.metal, mlx/backend/metal/kernels/scan.h, mlx/backend/metal/kernels/scan.metal, + mlx/backend/metal/kernels/sdpa_blocked.h, mlx/backend/metal/kernels/sdpa_vector.h, mlx/backend/metal/kernels/searchsorted.h, mlx/backend/metal/kernels/searchsorted.metal, @@ -1019,6 +1034,8 @@ mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal, mlx/backend/metal/kernels/steel/gemm/loader.h, mlx/backend/metal/kernels/steel/gemm/mma.h, mlx/backend/metal/kernels/steel/gemm/nax.h, @@ -1218,6 +1235,7 @@ python/tests/mpi_test_distributed.py, python/tests/nccl_test_distributed.py, python/tests/ring_test_distributed.py, + python/tests/run.py, python/tests/test_array.py, python/tests/test_autograd.py, python/tests/test_bf16.py, @@ -1231,6 +1249,7 @@ python/tests/test_einsum.py, python/tests/test_eval.py, python/tests/test_export_import.py, + python/tests/test_extension.py, python/tests/test_fast_gated_delta.py, python/tests/test_fast_sdpa.py, python/tests/test_fast.py, @@ -1251,6 +1270,9 @@ python/tests/test_tree.py, python/tests/test_upsample.py, python/tests/test_vmap.py, + python/tests/test_zero_copy.py, + "python/tests/xp-tests/skips.txt", + "python/tests/xp-tests/xfail.txt", README.md, setup.py, tests/allocator_tests.cpp, @@ -1374,6 +1396,7 @@ metal/cexpf.h, metal/complex.h, metal/copy.h, + metal/cross_entropy.h, metal/defines.h, metal/dot.h, metal/erf.h, @@ -1385,6 +1408,11 @@ metal/fp_quantized.h, metal/fp4.h, metal/fp8.h, + metal/gated_delta_nax_ops.h, + metal/gated_delta_update_nax_vjp.h, + metal/gated_delta_update_nax.h, + metal/gated_delta_update_vjp.h, + metal/gated_delta_update.h, metal/gemv_masked.h, metal/gemv.h, metal/hadamard.h, @@ -1408,6 +1436,7 @@ metal/reduction/reduce_init.h, metal/reduction/reduce_row.h, metal/scan.h, + metal/sdpa_blocked.h, metal/sdpa_vector.h, metal/searchsorted.h, metal/softmax.h, @@ -1441,6 +1470,7 @@ metal/steel/gemm/kernels/steel_gemm_segmented.h, metal/steel/gemm/kernels/steel_gemm_splitk_nax.h, metal/steel/gemm/kernels/steel_gemm_splitk.h, + metal/steel/gemm/kernels/steel_gemm_thin_nax.h, metal/steel/gemm/loader.h, metal/steel/gemm/mma.h, metal/steel/gemm/nax.h, diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index 57d2b9115..1697339cb 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -7,7 +7,7 @@ BUILD_LIBRARY_FOR_DISTRIBUTION = YES HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann $(SRCROOT)/../Source/Cmlx/mlx/mlx/distributed/jaccl/lib -OTHER_CFLAGS = -isysroot $(SDKROOT) +OTHER_CFLAGS = -isysroot $(SDKROOT) -Wno-shorten-64-to-32 USER_HEADER_SEARCH_PATHS = $(SRCROOT)/../Source/Cmlx/mlx-c $(SRCROOT)/../Source/Cmlx/mlx