From f21a716ba830a339db0cf46fa02772ffbdd382c6 Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Mon, 5 Oct 2026 23:08:04 +0200 Subject: [PATCH 1/6] Vendor mlx core v0.32.3 and re-apply the SharpAI patches Replace Source/Cmlx/mlx with upstream ml-explore/mlx v0.32.3 (64ea011) plus the carried patches (SSD expert streaming, TurboKV, iOS I/O guards, quieter Metal logging), merged with git rather than copied. Conflicts were limited to mlx/io/load.{h,cpp}, where 0.32.3 reworked ParallelFileReader's thread pool; the upstream design is kept and only the iOS guards are carried on top. Regenerated Source/Cmlx/mlx-generated and the include dirs with tools/update-mlx.sh (moe_stream.metal is carried by hand). Build fixes for the new core: exclude the new CUDA sources from the Cmlx target, and make mlx_gather_qmm pass nullopt for the new global_scale argument (mlx #4458; mlx-c has not caught up). FP8Tests now expects NaN for the E4M3FN bytes 0x7F and 0xFF, which 0.32.3 decodes correctly (mlx #4376). Co-Authored-By: Claude Sonnet 5.5 --- Package.swift | 3 + Source/Cmlx/include-framework/Metal.hpp | 2 +- Source/Cmlx/include-framework/mlx-array.h | 28 +- Source/Cmlx/include-framework/mlx-fast.h | 10 + Source/Cmlx/include-framework/mlx-io-load.h | 97 +- Source/Cmlx/include-framework/mlx-memory.h | 15 +- Source/Cmlx/include-framework/mlx-ops.h | 4 + .../Cmlx/include-framework/mlx-primitives.h | 1 + Source/Cmlx/include-framework/mlx-stream.h | 4 + .../include-framework/mlx-types-complex.h | 17 +- .../include-framework/mlx-types-half_types.h | 7 - Source/Cmlx/include-framework/mlx-utils.h | 15 +- Source/Cmlx/include-framework/mlx-version.h | 2 +- Source/Cmlx/mlx-c/mlx/c/ops.cpp | 1 + Source/Cmlx/mlx-generated/binary_ops.cpp | 2 +- .../Cmlx/mlx-generated/compiled_preamble.cpp | 50 +- Source/Cmlx/mlx-generated/cross_entropy.cpp | 152 ++ .../mlx-generated/cuda/cuda_jit_sources.h | 980 +++++----- Source/Cmlx/mlx-generated/fp_quantized.cpp | 397 ++-- .../Cmlx/mlx-generated/fp_quantized_nax.cpp | 377 ++-- .../mlx-generated/gated_delta_update_nax.cpp | 1652 +++++++++++++++++ Source/Cmlx/mlx-generated/gemv.cpp | 71 + .../Cmlx/mlx-generated/metal/arg_reduce.metal | 24 +- Source/Cmlx/mlx-generated/metal/binary_ops.h | 2 +- .../Cmlx/mlx-generated/metal/cross_entropy.h | 133 ++ Source/Cmlx/mlx-generated/metal/defines.h | 1 + Source/Cmlx/mlx-generated/metal/fp8.h | 10 +- .../Cmlx/mlx-generated/metal/fp_quantized.h | 387 ++-- .../mlx-generated/metal/fp_quantized_nax.h | 367 ++-- .../mlx-generated/metal/gated_delta_update.h | 349 ++++ .../metal/gated_delta_update_nax.h | 527 ++++++ Source/Cmlx/mlx-generated/metal/quantized.h | 179 +- .../Cmlx/mlx-generated/metal/quantized_nax.h | 247 ++- .../metal/scaled_dot_product_attention.metal | 27 +- Source/Cmlx/mlx-generated/metal/sdpa_vector.h | 7 +- .../steel/attn/kernels/steel_attention.h | 124 +- .../steel/attn/kernels/steel_attention.metal | 45 +- .../steel/attn/kernels/steel_attention_nax.h | 92 +- .../steel/gemm/kernels/steel_gemm_gather.h | 283 ++- .../gemm/kernels/steel_gemm_gather_nax.h | 142 +- Source/Cmlx/mlx-generated/metal/unary_ops.h | 10 +- Source/Cmlx/mlx-generated/metal/utils.h | 70 + Source/Cmlx/mlx-generated/quantized.cpp | 179 +- Source/Cmlx/mlx-generated/quantized_nax.cpp | 247 ++- Source/Cmlx/mlx-generated/scan.cpp | 2 +- Source/Cmlx/mlx-generated/steel_attention.cpp | 124 +- .../mlx-generated/steel_attention_nax.cpp | 92 +- .../Cmlx/mlx-generated/steel_gemm_gather.cpp | 283 ++- .../mlx-generated/steel_gemm_gather_nax.cpp | 144 +- Source/Cmlx/mlx-generated/unary_ops.cpp | 20 +- Source/Cmlx/mlx-generated/utils.cpp | 71 + .../.github/actions/build-wheel/action.yml | 1 + .../Cmlx/mlx/.github/actions/setup/action.yml | 9 +- .../mlx/.github/actions/test-wheel/action.yml | 7 +- .../scripts/find-eligible-contributors.js | 12 + .../mlx/.github/workflows/build_and_test.yml | 2 +- Source/Cmlx/mlx/.github/workflows/release.yml | 2 +- .../.github/workflows/update_bypass_list.yml | 1 + Source/Cmlx/mlx/.gitignore | 1 + Source/Cmlx/mlx/ACKNOWLEDGMENTS.md | 1 + Source/Cmlx/mlx/CMakeLists.txt | 3 +- Source/Cmlx/mlx/CODE_OF_CONDUCT.md | 132 -- Source/Cmlx/mlx/CONTRIBUTING.md | 104 +- .../benchmarks/python/gated_delta_bench.py | 145 ++ .../benchmarks/python/gather_qqmm_bench.py | 108 ++ .../Cmlx/mlx/benchmarks/python/qqmm_bench.py | 62 + .../mlx/benchmarks/python/swiglu_bench.py | 77 + .../mlx/benchmarks/python/swiglu_qmm_bench.py | 72 + Source/Cmlx/mlx/docs/src/dev/extensions.rst | 2 +- .../mlx/docs/src/python/memory_management.rst | 1 + .../mlx/examples/extensions/CMakeLists.txt | 20 +- .../mlx/examples/extensions/pyproject.toml | 3 +- .../mlx/examples/extensions/requirements.txt | 1 - Source/Cmlx/mlx/mlx/CMakeLists.txt | 16 +- Source/Cmlx/mlx/mlx/array.cpp | 19 +- Source/Cmlx/mlx/mlx/array.h | 28 +- Source/Cmlx/mlx/mlx/backend/common/compiled.h | 4 +- Source/Cmlx/mlx/mlx/backend/common/utils.cpp | 17 + .../Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp | 4 +- Source/Cmlx/mlx/mlx/backend/cpu/binary.h | 6 +- Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h | 2 +- .../Cmlx/mlx/mlx/backend/cpu/distributed.cpp | 2 +- Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp | 9 +- Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp | 33 +- Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp | 8 + .../mlx/backend/cpu/simd/accelerate_simd.h | 20 +- .../Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h | 24 +- Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp | 41 +- Source/Cmlx/mlx/mlx/backend/cpu/unary_ops.h | 11 +- .../Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt | 12 + .../Cmlx/mlx/mlx/backend/cuda/allocator.cpp | 36 +- .../Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu | 20 +- Source/Cmlx/mlx/mlx/backend/cuda/cholesky.cu | 141 ++ Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h | 7 +- .../mlx/mlx/backend/cuda/cusolver_utils.cpp | 47 + .../mlx/mlx/backend/cuda/cusolver_utils.h | 18 + .../Cmlx/mlx/mlx/backend/cuda/delayload.cpp | 19 + Source/Cmlx/mlx/mlx/backend/cuda/device.cpp | 21 +- .../mlx/mlx/backend/cuda/device/qmm_naive.cuh | 4 + Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp | 2 + Source/Cmlx/mlx/mlx/backend/cuda/event.cu | 3 +- Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp | 42 +- .../Cmlx/mlx/mlx/backend/cuda/primitives.cpp | 15 +- .../mlx/mlx/backend/cuda/quantized/qmm/qmm.h | 1 + .../mlx/mlx/backend/cuda/quantized/qmm/qmv.cu | 24 +- .../mlx/mlx/backend/cuda/quantized/qqmm.cpp | 9 + .../mlx/backend/cuda/quantized/quantized.cpp | 16 +- Source/Cmlx/mlx/mlx/backend/cuda/scan.cu | 5 + Source/Cmlx/mlx/mlx/backend/cuda/sort.cu | 4 + Source/Cmlx/mlx/mlx/backend/cuda/wddm.cpp | 107 ++ Source/Cmlx/mlx/mlx/backend/cuda/wddm.h | 26 + Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp | 9 +- .../Cmlx/mlx/mlx/backend/metal/CMakeLists.txt | 4 + .../mlx/mlx/backend/metal/cross_entropy.cpp | 125 ++ Source/Cmlx/mlx/mlx/backend/metal/device.cpp | 9 +- Source/Cmlx/mlx/mlx/backend/metal/fence.cpp | 19 +- .../mlx/backend/metal/gated_delta_update.cpp | 167 ++ .../Cmlx/mlx/mlx/backend/metal/jit/includes.h | 4 + .../mlx/mlx/backend/metal/jit_kernels.cpp | 133 +- Source/Cmlx/mlx/mlx/backend/metal/kernels.h | 25 +- .../mlx/backend/metal/kernels/CMakeLists.txt | 5 + .../backend/metal/kernels/arg_reduce.metal | 24 +- .../mlx/backend/metal/kernels/binary_ops.h | 2 +- .../mlx/backend/metal/kernels/cross_entropy.h | 133 ++ .../backend/metal/kernels/cross_entropy.metal | 18 + .../mlx/mlx/backend/metal/kernels/defines.h | 1 + .../Cmlx/mlx/mlx/backend/metal/kernels/fp8.h | 10 +- .../mlx/backend/metal/kernels/fp_quantized.h | 387 ++-- .../backend/metal/kernels/fp_quantized.metal | 46 +- .../backend/metal/kernels/fp_quantized_nax.h | 367 ++-- .../metal/kernels/fp_quantized_nax.metal | 43 +- .../metal/kernels/gated_delta_update.h | 349 ++++ .../metal/kernels/gated_delta_update.metal | 42 + .../metal/kernels/gated_delta_update_nax.h | 527 ++++++ .../kernels/gated_delta_update_nax.metal | 31 + .../metal/kernels/gather_mm_offsets.metal | 23 + .../mlx/mlx/backend/metal/kernels/quantized.h | 179 +- .../mlx/backend/metal/kernels/quantized_nax.h | 247 ++- .../scaled_dot_product_attention.metal | 27 +- .../mlx/backend/metal/kernels/sdpa_vector.h | 7 +- .../steel/attn/kernels/steel_attention.h | 124 +- .../steel/attn/kernels/steel_attention.metal | 45 +- .../steel/attn/kernels/steel_attention_nax.h | 92 +- .../attn/kernels/steel_attention_nax.metal | 27 +- .../gemm/kernels/steel_gemm_fused_nax.metal | 2 + .../steel/gemm/kernels/steel_gemm_gather.h | 283 ++- .../gemm/kernels/steel_gemm_gather_nax.h | 142 +- .../gemm/kernels/steel_gemm_splitk_nax.metal | 2 + .../mlx/mlx/backend/metal/kernels/unary_ops.h | 10 +- .../mlx/mlx/backend/metal/kernels/utils.h | 70 + Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp | 103 +- Source/Cmlx/mlx/mlx/backend/metal/matmul.h | 7 + .../mlx/mlx/backend/metal/nojit_kernels.cpp | 27 + .../Cmlx/mlx/mlx/backend/metal/primitives.cpp | 23 - .../Cmlx/mlx/mlx/backend/metal/quantized.cpp | 273 ++- .../metal/scaled_dot_product_attention.cpp | 200 +- Source/Cmlx/mlx/mlx/backend/metal/scan.cpp | 5 + Source/Cmlx/mlx/mlx/backend/metal/sort.cpp | 4 + Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp | 9 +- .../mlx/mlx/backend/no_gpu/primitives.cpp | 11 + Source/Cmlx/mlx/mlx/compile.cpp | 9 +- .../distributed/jaccl/lib/jaccl/ring_impl.h | 29 +- Source/Cmlx/mlx/mlx/export.cpp | 1 + Source/Cmlx/mlx/mlx/fast.cpp | 127 ++ Source/Cmlx/mlx/mlx/fast.h | 10 + Source/Cmlx/mlx/mlx/fast_primitives.h | 32 + Source/Cmlx/mlx/mlx/fence.h | 15 +- Source/Cmlx/mlx/mlx/fft.cpp | 4 +- Source/Cmlx/mlx/mlx/io/gguf.cpp | 34 +- Source/Cmlx/mlx/mlx/io/gguf_quants.cpp | 13 +- Source/Cmlx/mlx/mlx/io/load.cpp | 78 +- Source/Cmlx/mlx/mlx/io/load.h | 97 +- Source/Cmlx/mlx/mlx/io/safetensors.cpp | 24 +- Source/Cmlx/mlx/mlx/linalg.cpp | 15 +- Source/Cmlx/mlx/mlx/memory.cpp | 30 + Source/Cmlx/mlx/mlx/memory.h | 15 +- Source/Cmlx/mlx/mlx/ops.cpp | 177 +- Source/Cmlx/mlx/mlx/ops.h | 4 + Source/Cmlx/mlx/mlx/primitives.cpp | 22 +- Source/Cmlx/mlx/mlx/primitives.h | 1 + Source/Cmlx/mlx/mlx/random.cpp | 2 +- Source/Cmlx/mlx/mlx/stream.cpp | 4 + Source/Cmlx/mlx/mlx/stream.h | 4 + Source/Cmlx/mlx/mlx/transforms.cpp | 34 +- Source/Cmlx/mlx/mlx/types/complex.h | 17 +- Source/Cmlx/mlx/mlx/types/half_types.h | 7 - Source/Cmlx/mlx/mlx/utils.cpp | 2 +- Source/Cmlx/mlx/mlx/utils.h | 15 +- Source/Cmlx/mlx/mlx/version.h | 2 +- .../mlx/nn/layers/positional_encoding.py | 55 +- .../Cmlx/mlx/python/mlx/nn/layers/upsample.py | 7 + Source/Cmlx/mlx/python/mlx/nn/losses.py | 11 +- .../mlx/python/mlx/optimizers/optimizers.py | 4 +- Source/Cmlx/mlx/python/src/array.cpp | 30 + Source/Cmlx/mlx/python/src/convert.cpp | 5 +- Source/Cmlx/mlx/python/src/convert.h | 18 - Source/Cmlx/mlx/python/src/fast.cpp | 30 +- Source/Cmlx/mlx/python/src/indexing.cpp | 14 +- Source/Cmlx/mlx/python/src/memory.cpp | 34 +- Source/Cmlx/mlx/python/src/ops.cpp | 45 +- Source/Cmlx/mlx/python/src/small_vector.h | 2 +- Source/Cmlx/mlx/python/src/stream.cpp | 32 +- Source/Cmlx/mlx/python/src/transforms.cpp | 12 +- .../mlx/python/tests/mlx_distributed_tests.py | 16 + Source/Cmlx/mlx/python/tests/mlx_tests.py | 21 + Source/Cmlx/mlx/python/tests/test_array.py | 72 + Source/Cmlx/mlx/python/tests/test_autograd.py | 11 + Source/Cmlx/mlx/python/tests/test_bf16.py | 15 + Source/Cmlx/mlx/python/tests/test_compile.py | 59 +- Source/Cmlx/mlx/python/tests/test_conv.py | 134 +- .../mlx/python/tests/test_conv_transpose.py | 32 +- Source/Cmlx/mlx/python/tests/test_device.py | 40 + Source/Cmlx/mlx/python/tests/test_fast.py | 3 + .../mlx/python/tests/test_fast_gated_delta.py | 286 +++ .../Cmlx/mlx/python/tests/test_fast_sdpa.py | 440 ++++- Source/Cmlx/mlx/python/tests/test_fft.py | 4 +- Source/Cmlx/mlx/python/tests/test_linalg.py | 29 + Source/Cmlx/mlx/python/tests/test_load.py | 18 + Source/Cmlx/mlx/python/tests/test_memory.py | 28 +- Source/Cmlx/mlx/python/tests/test_ops.py | 185 +- .../Cmlx/mlx/python/tests/test_optimizers.py | 23 + .../Cmlx/mlx/python/tests/test_quantized.py | 458 +++++ Source/Cmlx/mlx/python/tests/test_vmap.py | 8 +- .../Cmlx/mlx/python/tests/test_zero_copy.py | 62 + Source/Cmlx/mlx/setup.py | 2 + Source/Cmlx/mlx/tests/CMakeLists.txt | 11 +- Source/Cmlx/mlx/tests/allocator_tests.cpp | 20 + Source/Cmlx/mlx/tests/arg_reduce_tests.cpp | 23 + Source/Cmlx/mlx/tests/array_tests.cpp | 57 + Source/Cmlx/mlx/tests/autograd_tests.cpp | 32 +- Source/Cmlx/mlx/tests/cuda_tests.cpp | 111 ++ Source/Cmlx/mlx/tests/fft_tests.cpp | 11 +- Source/Cmlx/mlx/tests/gpu_tests.cpp | 18 + Source/Cmlx/mlx/tests/load_tests.cpp | 98 + Source/Cmlx/mlx/tests/ops_tests.cpp | 144 +- Source/Cmlx/mlx/tests/scheduler_tests.cpp | 55 + Tests/MLXTests/FP8Tests.swift | 12 +- UPSTREAM_MERGE_PLAN.md | 6 +- 238 files changed, 13496 insertions(+), 4257 deletions(-) create mode 100644 Source/Cmlx/mlx-generated/cross_entropy.cpp create mode 100644 Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp create mode 100644 Source/Cmlx/mlx-generated/metal/cross_entropy.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h delete mode 100644 Source/Cmlx/mlx/CODE_OF_CONDUCT.md create mode 100644 Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py create mode 100644 Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py create mode 100644 Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py create mode 100644 Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py create mode 100644 Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py create mode 100644 Source/Cmlx/mlx/mlx/backend/cuda/cholesky.cu create mode 100644 Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.cpp create mode 100644 Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.h create mode 100644 Source/Cmlx/mlx/mlx/backend/cuda/wddm.cpp create mode 100644 Source/Cmlx/mlx/mlx/backend/cuda/wddm.h create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/cross_entropy.cpp create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.h create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.metal create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal create mode 100644 Source/Cmlx/mlx/mlx/backend/metal/kernels/gather_mm_offsets.metal create mode 100644 Source/Cmlx/mlx/mlx/memory.cpp create mode 100644 Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py create mode 100644 Source/Cmlx/mlx/tests/cuda_tests.cpp diff --git a/Package.swift b/Package.swift index 7390b8f74..2fbab7f02 100644 --- a/Package.swift +++ b/Package.swift @@ -57,10 +57,12 @@ let noCudaCmlxExcludes = [ // Exclude CUDA backend files, but keep no_cuda.cpp for stubs // mlx/mlx/backend/cuda/no_cuda.cpp "mlx/mlx/backend/cuda/allocator.cpp", + "mlx/mlx/backend/cuda/cholesky.cu", "mlx/mlx/backend/cuda/compiled.cpp", "mlx/mlx/backend/cuda/conv.cpp", "mlx/mlx/backend/cuda/cublas_utils.cpp", "mlx/mlx/backend/cuda/cudnn_utils.cpp", + "mlx/mlx/backend/cuda/cusolver_utils.cpp", "mlx/mlx/backend/cuda/custom_kernel.cpp", "mlx/mlx/backend/cuda/delayload.cpp", "mlx/mlx/backend/cuda/device_info.cpp", @@ -77,6 +79,7 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp", "mlx/mlx/backend/cuda/slicing.cpp", "mlx/mlx/backend/cuda/utils.cpp", + "mlx/mlx/backend/cuda/wddm.cpp", "mlx/mlx/backend/cuda/worker.cpp", "mlx/mlx/backend/cuda/binary", "mlx/mlx/backend/cuda/conv", diff --git a/Source/Cmlx/include-framework/Metal.hpp b/Source/Cmlx/include-framework/Metal.hpp index a44ec5666..a30f8636c 100644 --- a/Source/Cmlx/include-framework/Metal.hpp +++ b/Source/Cmlx/include-framework/Metal.hpp @@ -2,7 +2,7 @@ // // Metal.hpp.in // -// Autogenerated from commit 8f784840de21cacd0c5c237ea860179e47dcec91. +// Autogenerated on October 05, 2026. // // Copyright 2020-2024 Apple Inc. // diff --git a/Source/Cmlx/include-framework/mlx-array.h b/Source/Cmlx/include-framework/mlx-array.h index 9c6819c2e..4af752e76 100644 --- a/Source/Cmlx/include-framework/mlx-array.h +++ b/Source/Cmlx/include-framework/mlx-array.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -81,13 +82,22 @@ class MLX_API array { array& operator=(array&& other) && = delete; /** Default copy and move constructors otherwise. */ - array& operator=(array&& other) & = default; array(const array& other) = default; array(array&& other) = default; + /** Assignment goes through release() to break sibling reference cycle. */ array& operator=(const array& other) & { - if (this->id() != other.id()) { - this->array_desc_ = other.array_desc_; + if (array_desc_ != other.array_desc_) { + // Hold the old descriptor while assigning, and then release it. Otherwise + // when |this| and |other| are siblings they would be both destroyed. + release(std::exchange(array_desc_, other.array_desc_)); + } + return *this; + } + + array& operator=(array&& other) & noexcept { + if (array_desc_ != other.array_desc_) { + release(std::exchange(array_desc_, std::move(other.array_desc_))); } return *this; } @@ -467,10 +477,12 @@ class MLX_API array { void copy_shared_buffer(const array& other); void overwrite_descriptor(const array& other) { - array_desc_ = other.array_desc_; + release(std::exchange(array_desc_, other.array_desc_)); } - ~array(); + ~array() { + release(std::move(array_desc_)); + } private: // Initialize the arrays data @@ -533,6 +545,12 @@ class MLX_API array { // the primitive which knows how to compute the array's data from its inputs // and the list of array's inputs for the primitive. std::shared_ptr array_desc_; + + /** + * Lets go of an ArrayDesc, breaking the sibling reference cycle of a + * multi-output primitive. + */ + static void release(std::shared_ptr desc); }; template diff --git a/Source/Cmlx/include-framework/mlx-fast.h b/Source/Cmlx/include-framework/mlx-fast.h index 9e8786d53..ad99de725 100644 --- a/Source/Cmlx/include-framework/mlx-fast.h +++ b/Source/Cmlx/include-framework/mlx-fast.h @@ -61,6 +61,16 @@ MLX_API array scaled_dot_product_attention( bool force_fused = false, StreamOrDevice s = {}); +MLX_API std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state = std::nullopt, + const std::optional& mask = std::nullopt, + StreamOrDevice s = {}); + using TemplateArg = std::variant; using ScalarArg = std::variant; diff --git a/Source/Cmlx/include-framework/mlx-io-load.h b/Source/Cmlx/include-framework/mlx-io-load.h index c5e4fc7b4..1c133548d 100644 --- a/Source/Cmlx/include-framework/mlx-io-load.h +++ b/Source/Cmlx/include-framework/mlx-io-load.h @@ -4,11 +4,13 @@ #pragma once #include +#include #include -#include // SIZE_MAX +#include // SIZE_MAX +#include #include -#ifdef _MSC_VER +#ifdef _WIN32 #include #else #include @@ -67,6 +69,8 @@ class Writer { std::ios_base::seekdir way = std::ios_base::beg) = 0; virtual void write(const char* data, size_t n) = 0; virtual std::string label() const = 0; + virtual void open() {} + virtual ~Writer() = default; }; @@ -89,20 +93,30 @@ class ParallelFileReader : public Reader { } size_t tell() override { +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } // Warning: do not use this function from multiple threads as // it advances the file descriptor void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } // Warning: do not use this function from multiple threads as @@ -116,42 +130,56 @@ class ParallelFileReader : public Reader { } private: - // On iOS, always use sequential pread() — no std::async thread spawning. - // std::async wakeup bursts during model loading exceed iOS's per-process - // wakeup rate limit → EXC_RESOURCE (RESOURCE_TYPE_WAKEUPS) crash. - // On macOS, 32MB parallel batches maximize NVMe throughput. + // Reads larger than this are split in batches and read in parallel. + // Fork: on iOS, always use sequential pread() (no thread spawning). Wakeup + // bursts during model loading exceed iOS's per-process wakeup rate limit + // (EXC_RESOURCE, RESOURCE_TYPE_WAKEUPS). #ifdef MLX_IOS_SEQUENTIAL_IO static constexpr size_t batch_size_ = SIZE_MAX; #else - static constexpr size_t batch_size_ = 1 << 25; // 32MB + static constexpr size_t batch_size_ = 1 << 25; #endif - static ThreadPool& thread_pool(); + + // The pool that reads the batches, held from the first batched read until + // the reader is destroyed. It cannot be io::thread_pool(), the tasks that + // run there wait for these batches and would deadlock in the same pool. + ThreadPool& thread_pool(); + int fd_; std::string label_; + std::once_flag pool_once_; + std::shared_ptr pool_; }; class FileWriter : public Writer { public: explicit FileWriter() {} explicit FileWriter(std::string file_path) - : fd_(open( - file_path.c_str(), - O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, - 0644)), - label_(std::move(file_path)) {} + : file_path_(std::move(file_path)) {} FileWriter(const FileWriter&) = delete; FileWriter& operator=(const FileWriter&) = delete; - FileWriter(FileWriter&& other) { - std::swap(fd_, other.fd_); + FileWriter(FileWriter&& other) + : fd_(std::exchange(other.fd_, -1)), + file_path_(std::move(other.file_path_)) { + other.file_path_.clear(); } ~FileWriter() override { - if (fd_ != 0) { + if (fd_ >= 0) { close(fd_); } } + // Kept separate from construction so lazy inputs can be evaluated first, + // they may still read from the file. + void open() override { + if (fd_ < 0 && !file_path_.empty()) { + fd_ = ::open( + file_path_.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, 0644); + } + } + bool is_open() const override { return fd_ >= 0; } @@ -161,21 +189,34 @@ class FileWriter : public Writer { } size_t tell() override { + check_open(); +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + check_open(); + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } void write(const char* data, size_t n) override { + check_open(); while (n != 0) { auto m = ::write(fd_, data, std::min(n, static_cast(INT32_MAX))); if (m <= 0) { @@ -189,12 +230,18 @@ class FileWriter : public Writer { } std::string label() const override { - return "file " + label_; + return "file " + file_path_; } private: - int fd_{0}; - std::string label_; + void check_open() const { + if (!is_open()) { + throw std::runtime_error("[write] File " + file_path_ + " is not open."); + } + } + + int fd_{-1}; + std::string file_path_; }; } // namespace io diff --git a/Source/Cmlx/include-framework/mlx-memory.h b/Source/Cmlx/include-framework/mlx-memory.h index bfa2f22d7..448aa8e68 100644 --- a/Source/Cmlx/include-framework/mlx-memory.h +++ b/Source/Cmlx/include-framework/mlx-memory.h @@ -1,11 +1,13 @@ #ifdef __cplusplus -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #pragma once #include +#include #include +#include namespace mlx::core { @@ -16,6 +18,14 @@ namespace mlx::core { * */ MLX_API size_t get_active_memory(); +/* 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. + * */ +MLX_API size_t get_array_buffer_size(const std::vector& arrays); + /* Get the peak amount of used memory in bytes. * * The maximum memory used recorded from the beginning of the program @@ -43,6 +53,9 @@ MLX_API size_t get_cache_memory(); * When Metal is available the memory limit defaults to 1.5 times the maximum * recommended working set size reported by the device. * + * On native Windows with CUDA the memory limit is also bounded by the current + * video memory budget reported by the operating system. + * * Returns the previous memory limit. * */ MLX_API size_t set_memory_limit(size_t limit); diff --git a/Source/Cmlx/include-framework/mlx-ops.h b/Source/Cmlx/include-framework/mlx-ops.h index 96e3ecfe2..d3ff5a0d5 100644 --- a/Source/Cmlx/include-framework/mlx-ops.h +++ b/Source/Cmlx/include-framework/mlx-ops.h @@ -382,6 +382,9 @@ inline array transpose( return transpose(a, std::vector(axes), s); } +/** Permutes last two dimensions of an array. */ +MLX_API array matrix_transpose(const array& a, StreamOrDevice s = {}); + /** Swap two axes of an array. */ MLX_API array swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {}); @@ -1627,6 +1630,7 @@ MLX_API array gather_qmm( std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "affine", + const std::optional& global_scale = std::nullopt, bool sorted_indices = false, StreamOrDevice s = {}); diff --git a/Source/Cmlx/include-framework/mlx-primitives.h b/Source/Cmlx/include-framework/mlx-primitives.h index 33d2bb094..692f821ab 100644 --- a/Source/Cmlx/include-framework/mlx-primitives.h +++ b/Source/Cmlx/include-framework/mlx-primitives.h @@ -1889,6 +1889,7 @@ class Scan : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS(); + DEFINE_INPUT_OUTPUT_SHAPE() const char* name() const override { switch (reduce_type_) { diff --git a/Source/Cmlx/include-framework/mlx-stream.h b/Source/Cmlx/include-framework/mlx-stream.h index 9b9bae946..2d219bc3e 100644 --- a/Source/Cmlx/include-framework/mlx-stream.h +++ b/Source/Cmlx/include-framework/mlx-stream.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -30,6 +31,9 @@ struct MLX_API ThreadLocalStream : public Stream { /** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); +/** Get the default stream of current thread if it exists. */ +MLX_API std::optional peek_default_stream(Device d); + /** Make the stream the default for its device on current thread. */ MLX_API void set_default_stream(Stream s); diff --git a/Source/Cmlx/include-framework/mlx-types-complex.h b/Source/Cmlx/include-framework/mlx-types-complex.h index c5c1f2849..f9dda2b71 100644 --- a/Source/Cmlx/include-framework/mlx-types-complex.h +++ b/Source/Cmlx/include-framework/mlx-types-complex.h @@ -107,9 +107,24 @@ inline complex64_t operator-(const complex64_t& v) { complex_binop_helper(_op_, _operator_, float16_t) \ complex_binop_helper(_op_, _operator_, bfloat16_t) \ complex_binop_helper(_op_, _operator_, float) -// clang-format on complex_binop(+, operator+) +template +constexpr bool is_complex = false; + +template +constexpr bool is_complex().real())>> = true; +// clang-format on + +template +inline bool isnan(T v) { + if constexpr (is_complex) { + return std::isnan(std::real(v)) || std::isnan(std::imag(v)); + } else { + return std::isnan(v); + } +} + } // namespace mlx::core #endif diff --git a/Source/Cmlx/include-framework/mlx-types-half_types.h b/Source/Cmlx/include-framework/mlx-types-half_types.h index b6869f528..49140ca13 100644 --- a/Source/Cmlx/include-framework/mlx-types-half_types.h +++ b/Source/Cmlx/include-framework/mlx-types-half_types.h @@ -57,13 +57,6 @@ fp16_bf16_binop_helper(/, operator/) #endif -template -constexpr bool is_complex = false; - -template -constexpr bool is_complex().real())>> = - true; - template constexpr bool is_signed_v = std::is_signed_v || std::is_same_v || std::is_same_v; diff --git a/Source/Cmlx/include-framework/mlx-utils.h b/Source/Cmlx/include-framework/mlx-utils.h index a0d8478a4..53665bb87 100644 --- a/Source/Cmlx/include-framework/mlx-utils.h +++ b/Source/Cmlx/include-framework/mlx-utils.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -29,7 +30,9 @@ MLX_API Stream to_stream(StreamOrDevice s, Device default_); struct StreamContext { public: - StreamContext(StreamOrDevice s) : _stream(default_stream(default_device())) { + StreamContext(StreamOrDevice s) + : _stream(default_stream(default_device())), + _owner(std::this_thread::get_id()) { if (std::holds_alternative(s)) { throw std::runtime_error( "[StreamContext] Invalid argument, please specify a stream or device."); @@ -39,13 +42,21 @@ struct StreamContext { set_default_stream(_s); } - ~StreamContext() { + // noexcept(false) so the thread check below can throw. A destructor that + // throws while another exception unwinds this thread calls std::terminate. + ~StreamContext() noexcept(false) { + if (std::this_thread::get_id() != _owner) { + throw std::runtime_error( + "[StreamContext] Destroyed in a different thread than it was " + "constructed in."); + } set_default_device(_stream.device); set_default_stream(_stream); } private: Stream _stream; + std::thread::id _owner; }; struct MLX_API PrintOptions { diff --git a/Source/Cmlx/include-framework/mlx-version.h b/Source/Cmlx/include-framework/mlx-version.h index 3b3eb3f2a..f672cd22f 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 2 +#define MLX_VERSION_PATCH 3 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) diff --git a/Source/Cmlx/mlx-c/mlx/c/ops.cpp b/Source/Cmlx/mlx-c/mlx/c/ops.cpp index bf565d1e3..9e1b06bd4 100644 --- a/Source/Cmlx/mlx-c/mlx/c/ops.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/ops.cpp @@ -1708,6 +1708,7 @@ 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) sorted_indices, mlx_stream_get_(s))); } catch (std::exception& e) { diff --git a/Source/Cmlx/mlx-generated/binary_ops.cpp b/Source/Cmlx/mlx-generated/binary_ops.cpp index 2c8157725..9c2d6f475 100644 --- a/Source/Cmlx/mlx-generated/binary_ops.cpp +++ b/Source/Cmlx/mlx-generated/binary_ops.cpp @@ -273,7 +273,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/compiled_preamble.cpp b/Source/Cmlx/mlx-generated/compiled_preamble.cpp index 8b89f66a8..c23168e74 100644 --- a/Source/Cmlx/mlx-generated/compiled_preamble.cpp +++ b/Source/Cmlx/mlx-generated/compiled_preamble.cpp @@ -76,11 +76,6 @@ inline float operator+(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator-(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator*(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator*(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator/(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) / static_cast(rhs); } inline float operator/(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) / static_cast(rhs); } -template -constexpr bool is_complex = false; -template -constexpr bool is_complex().real())>> = - true; template constexpr bool is_signed_v = std::is_signed_v || std::is_same_v || std::is_same_v; @@ -147,6 +142,18 @@ inline complex64_t operator-(const complex64_t& v) { return -static_cast>(v); } inline complex64_t operator+(const std::complex& x, const complex64_t& y) { return x + static_cast>(y); } inline complex64_t operator+(const complex64_t& x, const std::complex& y) { return static_cast>(x) + y; } inline complex64_t operator+(const complex64_t& x, const complex64_t& y) { return static_cast>(x) + static_cast>(y); } inline complex64_t operator+(bool x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, bool y) { return x + static_cast(y); } inline complex64_t operator+(uint32_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, uint32_t y) { return x + static_cast(y); } inline complex64_t operator+(uint64_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, uint64_t y) { return x + static_cast(y); } inline complex64_t operator+(int32_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, int32_t y) { return x + static_cast(y); } inline complex64_t operator+(int64_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, int64_t y) { return x + static_cast(y); } inline complex64_t operator+(float16_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, float16_t y) { return x + static_cast(y); } inline complex64_t operator+(bfloat16_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, bfloat16_t y) { return x + static_cast(y); } inline complex64_t operator+(float x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, float y) { return x + static_cast(y); } +template +constexpr bool is_complex = false; +template +constexpr bool is_complex().real())>> = true; +template +inline bool isnan(T v) { + if constexpr (is_complex) { + return std::isnan(std::real(v)) || std::isnan(std::imag(v)); + } else { + return std::isnan(v); + } +} } namespace mlx::core::simd { template @@ -269,7 +276,7 @@ auto imag(Simd in) -> Simd { } template Simd isnan(Simd in) { - return std::isnan(in.value); + return mlx::core::isnan(in.value); } template auto operator +(Simd a, Simd b) ->Simd { return a.value + b.value; } template auto operator +(T1 a, Simd b)->Simd { return a + b.value; } template auto operator +(Simd a, T2 b)->Simd { return a.value + b; } template auto operator -(Simd a, Simd b) ->Simd { return a.value - b.value; } template auto operator -(T1 a, Simd b)->Simd { return a - b.value; } template auto operator -(Simd a, T2 b)->Simd { return a.value - b; } @@ -287,11 +294,25 @@ Simd clz(Simd x_) { return __builtin_clz(x_.value); } template +Simd 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); + } + } + return a / b; +} +template Simd remainder(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; T r; if constexpr (std::is_integral_v) { + if (b == 0) { + return a; + } r = a % b; } else { r = std::remainder(a, b); @@ -308,7 +329,7 @@ Simd maximum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -319,7 +340,7 @@ Simd minimum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -640,19 +661,20 @@ struct ToFP8 { struct FromFP8 { template Simd operator()(Simd x) { - auto v = Simd(x & 127) << 7; + auto v = Simd(x & 127); + auto sign_bit = Simd((x >> 7) & 1) << 15; + auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; Simd out; if constexpr (simd::max_size >= N) { - auto converted = *(Simd*)(&v); + auto converted = *(Simd*)(&u); out = converted * 256.0; } else { for (int i = 0; i < N; ++i) { - auto converted = *(float16_t*)(&v[i]); + auto converted = *(float16_t*)(&u[i]); out[i] = converted * 256.0; } } - auto sign = Simd(x & 128); - return select(sign, -out, out); + return out; } float operator()(uint8_t x) { return (*this)(Simd(x)).value; @@ -663,7 +685,7 @@ namespace mlx::core::detail { 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 operator/(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 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/cross_entropy.cpp b/Source/Cmlx/mlx-generated/cross_entropy.cpp new file mode 100644 index 000000000..53d046999 --- /dev/null +++ b/Source/Cmlx/mlx-generated/cross_entropy.cpp @@ -0,0 +1,152 @@ +namespace mlx::core::metal { + +const char* cross_entropy() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/cross_entropy.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/cross_entropy.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/cross_entropy.h" +// Copyright © 2026 Apple Inc. + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy( + const device T* in, + const device int32_t* targets, + device float* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + in += gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(in[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = + (offset + i < axis_size) ? AccT(in[offset + i]) : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + uint n_simdgroups = ceildiv(lsize, uint(SIMD_SIZE)); + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = (simd_lane_id < n_simdgroups) ? local_max[simd_lane_id] + : Limits::finite_min; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = + (simd_lane_id < n_simdgroups) ? local_normalizer[simd_lane_id] : AccT(0); + normalizer = simd_sum(normalizer); + + if (lid == 0) { + // Subtract the two logits first; they are the same magnitude, so less is + // lost than in logsumexp - x_t. + AccT gap = maxval - x_t; + out[gid] = isinf(maxval) ? float(gap) : float(log(normalizer) + gap); + } +} + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy_vjp( + const device T* in, + const device int32_t* targets, + const device float* loss, + const device float* cotan, + device T* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + size_t row_offset = gid * size_t(axis_size); + in += row_offset; + out += row_offset; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + AccT g = AccT(cotan[gid]); + AccT l = AccT(loss[gid]); + + // out aliases in when the input is donated, so every thread must read the + // target column before any thread writes. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } else { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + if (col < axis_size) { + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } + } + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h index 3ed14c37c..b7089f3a2 100644 --- a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -2043,140 +2043,6 @@ 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, @@ -2319,6 +2185,140 @@ 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, @@ -4466,143 +4466,155 @@ constexpr char jit_source_qmm_naive[] = { 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, 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, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, 0x65, 0x20, 0x67, 0x6c, + 0x6f, 0x62, 0x61, 0x6c, 0x20, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x69, + 0x73, 0x20, 0x70, 0x65, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, + 0x65, 0x64, 0x20, 0x65, 0x78, 0x70, 0x65, 0x72, 0x74, 0x3b, 0x20, 0x61, + 0x20, 0x6e, 0x6f, 0x6e, 0x2d, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x65, + 0x64, 0x20, 0x63, 0x61, 0x6c, 0x6c, 0x20, 0x68, 0x61, 0x73, 0x20, 0x6f, + 0x6e, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x67, 0x6c, + 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x26, + 0x26, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, 0x6c, 0x6f, + 0x62, 0x61, 0x6c, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x2b, 0x3d, + 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 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, 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, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x5b, + 0x26, 0x5d, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, + 0x69, 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 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, 0x20, 0x20, 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, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 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, 0x20, 0x20, - 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x5b, 0x26, - 0x5d, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, - 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, - 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, - 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 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, 0x20, 0x20, 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, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 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, 0x20, 0x20, 0x20, - 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x75, 0x6d, 0x6d, 0x79, - 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3b, 0x20, 0x6e, 0x6f, 0x2d, - 0x62, 0x69, 0x61, 0x73, 0x20, 0x70, 0x61, 0x74, 0x68, 0x73, 0x20, 0x6e, - 0x65, 0x76, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, - 0x6f, 0x72, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x67, 0x5a, 0x2e, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, - 0x20, 0x67, 0x53, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, - 0x20, 0x7d, 0x28, 0x29, 0x3b, 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, 0x69, 0x6e, 0x74, 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, 0x63, - 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 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, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, - 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, - 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, - 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, - 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, 0x20, 0x2d, 0x20, 0x73, - 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, 0x2a, - 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x67, 0x41, 0x29, - 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, 0x69, - 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x3c, - 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x48, 0x61, 0x73, 0x4b, - 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x53, 0x4d, 0x38, - 0x30, 0x3e, 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, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, - 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, - 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, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x75, 0x6d, 0x6d, + 0x79, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3b, 0x20, 0x6e, 0x6f, + 0x2d, 0x62, 0x69, 0x61, 0x73, 0x20, 0x70, 0x61, 0x74, 0x68, 0x73, 0x20, + 0x6e, 0x65, 0x76, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x6f, 0x72, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x67, 0x5a, 0x2e, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x67, 0x53, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x28, 0x29, 0x3b, 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, 0x69, 0x6e, 0x74, 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, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 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, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, 0x20, 0x42, 0x4c, + 0x4b, 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, + 0x73, 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, 0x20, 0x2d, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, + 0x2a, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x67, 0x41, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, + 0x69, 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, + 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x48, 0x61, 0x73, + 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x53, 0x4d, + 0x38, 0x30, 0x3e, 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, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, + 0x65, 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, }; constexpr char jit_source_qmm_sm80[] = { @@ -5982,6 +5994,224 @@ constexpr char jit_source_qmm_sm90[] = { 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, +}; + constexpr char jit_source_scatter[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, @@ -6347,224 +6577,6 @@ 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/fp_quantized.cpp b/Source/Cmlx/mlx-generated/fp_quantized.cpp index 1c56cdb56..c9fade97e 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized.cpp @@ -98,11 +98,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -299,13 +299,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { } template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -317,7 +316,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -327,9 +327,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -343,6 +348,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -350,7 +358,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -365,14 +374,19 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales( scales_ + bi * src_ld / group_size + - (bj * pack_factor) / group_size) {} + (bj * pack_factor) / group_size) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -384,21 +398,15 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -951,12 +959,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -994,7 +1004,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -1013,12 +1024,15 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -1028,7 +1042,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -1038,9 +1052,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -1050,7 +1073,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -1061,6 +1084,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -1076,12 +1108,14 @@ template < typename T, int group_size, int bits, + bool has_global_scale = false, int BM = 32, int BK = 32, int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -1119,7 +1153,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; auto wl = (const device uint8_t*)w; @@ -1134,16 +1169,23 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1152,9 +1194,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1172,7 +1224,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1181,9 +1237,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1199,8 +1265,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } @@ -1241,11 +1307,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device uint8_t*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -1288,6 +1355,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -1607,12 +1678,14 @@ template < 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> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -1652,8 +1725,22 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1709,8 +1796,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1739,10 +1826,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1798,10 +1886,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1857,10 +1946,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1896,12 +1986,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> [[kernel]] void fp_gather_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1931,10 +2023,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1951,8 +2044,22 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1992,9 +2099,10 @@ 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, x, y, Xs, @@ -2020,6 +2128,7 @@ template < [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -2050,10 +2159,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -2070,8 +2180,21 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -2083,16 +2206,19 @@ template < int BK, int WM, int WN, - bool transpose> + bool transpose, + bool has_global_scale = false> [[kernel]] void fp_gather_qmm_rhs( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2123,7 +2249,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -2136,13 +2263,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2158,112 +2290,61 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp index 6b8473d85..78b045258 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp @@ -98,11 +98,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -163,7 +163,7 @@ inline float mx_scale_round_up(float x) { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/fp_quantized_nax.h" -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -223,13 +223,12 @@ struct Dequantize { }; template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -241,7 +240,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -269,6 +269,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -276,7 +279,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -289,7 +293,11 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id) {} + scales(scales_ + bi * src_ld / group_size + group_id) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -298,7 +306,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -312,14 +321,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -328,7 +330,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -354,6 +357,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -363,6 +367,7 @@ template < METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup Wtype* Ws, @@ -392,7 +397,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -408,7 +414,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -498,6 +504,7 @@ template < typename T, const int group_size, const int bits, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -507,6 +514,7 @@ template < METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Ws, @@ -536,7 +544,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -554,7 +563,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -649,11 +658,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device S*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -696,6 +706,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -710,10 +724,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -752,8 +768,19 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -811,8 +838,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -825,10 +852,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -857,10 +886,11 @@ template < threadgroup Wtype Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -877,8 +907,19 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -890,10 +931,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_n_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -924,10 +967,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -944,8 +988,8 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -958,16 +1002,19 @@ template < int WM, int WN, bool transpose, + bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_rhs_nax( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -984,7 +1031,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -996,13 +1044,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1028,10 +1081,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1039,151 +1092,125 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; - } - } + NAXTile Dtile; + Dtile.clear(); - xn += BK; - loader_w.next(); + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); } - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + xn += BK; + loader_w.next(); + } + + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp new file mode 100644 index 000000000..56afceba5 --- /dev/null +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp @@ -0,0 +1,1652 @@ +namespace mlx::core::metal { + +const char* gated_delta_update_nax() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/gated_delta_update_nax.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_update_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax.h" + +#include + +#include +#include + + +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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + // 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + // 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 +[[kernel]] void gated_delta_fused_nax( + 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)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + 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); + + 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)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile 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} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + 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); + } + }; + + 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)) + : 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; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + tmp_tile.clear(); + QKt_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) + + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/gemv.cpp b/Source/Cmlx/mlx-generated/gemv.cpp index 36ed559bb..0a4a5c78a 100644 --- a/Source/Cmlx/mlx-generated/gemv.cpp +++ b/Source/Cmlx/mlx-generated/gemv.cpp @@ -739,6 +739,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: @@ -1133,6 +1134,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -1251,6 +1275,53 @@ inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/utils.h" /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal index ba0268dc9..dd643afe3 100644 --- a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal +++ b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal @@ -1,11 +1,17 @@ // Copyright © 2023 Apple Inc. +#include #include #include "utils.h" using namespace metal; +template , bool> = true> +inline bool isnan(T) { + return false; +} + template struct IndexValPair { uint32_t index; @@ -17,8 +23,11 @@ struct ArgMin { static constexpr constant U init = Limits::max; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val > current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val > current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -29,7 +38,7 @@ struct ArgMin { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -43,8 +52,11 @@ struct ArgMax { static constexpr constant U init = Limits::min; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val < current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val < current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -55,7 +67,7 @@ struct ArgMax { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/Source/Cmlx/mlx-generated/metal/binary_ops.h b/Source/Cmlx/mlx-generated/metal/binary_ops.h index 37650c7d9..968087f8e 100644 --- a/Source/Cmlx/mlx-generated/metal/binary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/binary_ops.h @@ -261,7 +261,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/metal/cross_entropy.h b/Source/Cmlx/mlx-generated/metal/cross_entropy.h new file mode 100644 index 000000000..bef0bce5f --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/cross_entropy.h @@ -0,0 +1,133 @@ +// Copyright © 2026 Apple Inc. + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy( + const device T* in, + const device int32_t* targets, + device float* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + in += gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(in[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = + (offset + i < axis_size) ? AccT(in[offset + i]) : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + uint n_simdgroups = ceildiv(lsize, uint(SIMD_SIZE)); + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = (simd_lane_id < n_simdgroups) ? local_max[simd_lane_id] + : Limits::finite_min; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = + (simd_lane_id < n_simdgroups) ? local_normalizer[simd_lane_id] : AccT(0); + normalizer = simd_sum(normalizer); + + if (lid == 0) { + // Subtract the two logits first; they are the same magnitude, so less is + // lost than in logsumexp - x_t. + AccT gap = maxval - x_t; + out[gid] = isinf(maxval) ? float(gap) : float(log(normalizer) + gap); + } +} + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy_vjp( + const device T* in, + const device int32_t* targets, + const device float* loss, + const device float* cotan, + device T* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + size_t row_offset = gid * size_t(axis_size); + in += row_offset; + out += row_offset; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + AccT g = AccT(cotan[gid]); + AccT l = AccT(loss[gid]); + + // out aliases in when the input is donated, so every thread must read the + // target column before any thread writes. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } else { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + if (col < axis_size) { + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } + } + } +} diff --git a/Source/Cmlx/mlx-generated/metal/defines.h b/Source/Cmlx/mlx-generated/metal/defines.h index c369adb7e..0e6590616 100644 --- a/Source/Cmlx/mlx-generated/metal/defines.h +++ b/Source/Cmlx/mlx-generated/metal/defines.h @@ -14,6 +14,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: diff --git a/Source/Cmlx/mlx-generated/metal/fp8.h b/Source/Cmlx/mlx-generated/metal/fp8.h index 42c5ec128..00efdeb79 100644 --- a/Source/Cmlx/mlx-generated/metal/fp8.h +++ b/Source/Cmlx/mlx-generated/metal/fp8.h @@ -30,11 +30,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized.h b/Source/Cmlx/mlx-generated/metal/fp_quantized.h index 733361958..880ad589f 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized.h @@ -136,13 +136,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { } template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -154,7 +153,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -164,9 +164,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -180,6 +185,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -187,7 +195,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -202,14 +211,19 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales( scales_ + bi * src_ld / group_size + - (bj * pack_factor) / group_size) {} + (bj * pack_factor) / group_size) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -221,21 +235,15 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -788,12 +796,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -831,7 +841,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -850,12 +861,15 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -865,7 +879,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -875,9 +889,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -887,7 +910,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -898,6 +921,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -913,12 +945,14 @@ template < typename T, int group_size, int bits, + bool has_global_scale = false, int BM = 32, int BK = 32, int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -956,7 +990,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; auto wl = (const device uint8_t*)w; @@ -971,16 +1006,23 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -989,9 +1031,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1009,7 +1061,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1018,9 +1074,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1036,8 +1102,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } @@ -1078,11 +1144,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device uint8_t*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -1125,6 +1192,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -1444,12 +1515,14 @@ template < 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> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -1489,8 +1562,22 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1546,8 +1633,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1576,10 +1663,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1635,10 +1723,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1694,10 +1783,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1733,12 +1823,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> [[kernel]] void fp_gather_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1768,10 +1860,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1788,8 +1881,22 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1829,9 +1936,10 @@ 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, x, y, Xs, @@ -1857,6 +1965,7 @@ template < [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1887,10 +1996,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1907,8 +2017,21 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1920,16 +2043,19 @@ template < int BK, int WM, int WN, - bool transpose> + bool transpose, + bool has_global_scale = false> [[kernel]] void fp_gather_qmm_rhs( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1960,7 +2086,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -1973,13 +2100,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1995,112 +2127,61 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h index 35c61c2fa..0e4f7ec83 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -60,13 +60,12 @@ struct Dequantize { }; template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -78,7 +77,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -106,6 +106,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -113,7 +116,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -126,7 +130,11 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id) {} + scales(scales_ + bi * src_ld / group_size + group_id) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -135,7 +143,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -149,14 +158,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -165,7 +167,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -191,6 +194,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -200,6 +204,7 @@ template < METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup Wtype* Ws, @@ -229,7 +234,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -245,7 +251,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -335,6 +341,7 @@ template < typename T, const int group_size, const int bits, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -344,6 +351,7 @@ template < METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Ws, @@ -373,7 +381,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -391,7 +400,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -486,11 +495,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device S*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -533,6 +543,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -547,10 +561,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -589,8 +605,19 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -648,8 +675,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -662,10 +689,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -694,10 +723,11 @@ template < threadgroup Wtype Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -714,8 +744,19 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -727,10 +768,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_n_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -761,10 +804,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -781,8 +825,8 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -795,16 +839,19 @@ template < int WM, int WN, bool transpose, + bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_rhs_nax( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -821,7 +868,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -833,13 +881,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -865,10 +918,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -876,149 +929,123 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; - } - } + NAXTile Dtile; + Dtile.clear(); - xn += BK; - loader_w.next(); + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); } - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + xn += BK; + loader_w.next(); + } + + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h new file mode 100644 index 000000000..af29ff161 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h @@ -0,0 +1,349 @@ +#pragma once + +#include +#include "utils.h" + +#define AT(TILE, IDX) TILE.thread_elements()[IDX] +#define SUB(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) - AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) - AT(TILE2, 1); \ + } +#define ADD(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) + AT(TILE2, 1); \ + } +#define FMA(TILE0, S, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = S * AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = S * AT(TILE1, 1) + AT(TILE2, 1); \ + } + +#define SCALE(TILE0, S) \ + { \ + AT(TILE0, 0) *= S; \ + AT(TILE0, 1) *= S; \ + } +#define SCALE2(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= S0; \ + AT(TILE0, 1) *= S1; \ + } +#define SCALE_TRI(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn > fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 > fm ? 0.f : S1; \ + } +#define SCALE_TRIEQ(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn >= fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ + } + +// lambdas are not supported in metal 14 so porting to macros. + +// non transposed +#define LOAD_M(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fm < valid_rows) ? ((SRC)[fm * (LD) + fn]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fm < valid_rows) ? ((SRC)[fm * (LD) + fn + 1]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fm * (LD) + fn]); \ + AT(M, 1) = static_cast((SRC)[fm * (LD) + fn + 1]); \ + } + +// transposed load: sequence is the column -> mask fn / fn+1 +#define LOAD_MT(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fn < valid_rows) ? ((SRC)[fn * (LD) + fm]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fn + 1 < valid_rows) ? ((SRC)[(fn + 1) * (LD) + fm]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fn * (LD) + fm]); \ + AT(M, 1) = static_cast((SRC)[(fn + 1) * (LD) + fm]); \ + } + +#define PROCESS_CHUNK_SG(B, S_tile, VALID) \ + { \ + const short valid_rows = (VALID); \ + \ + 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)) \ + : 0.0f; \ + \ + float gamma_val = simd_prefix_inclusive_sum(g_val); \ + \ + if (thread_index_in_simdgroup < C) { \ + gamma[thread_index_in_simdgroup] = gamma_val; \ + } \ + simdgroup_barrier(mem_flags::mem_threadgroup); \ + \ + float gamma_fm = metal::fast::exp(gamma[fm]); \ + float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); \ + float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); \ + float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); \ + float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); \ + float gamma_C = metal::fast::exp(gamma[C - 1]); \ + \ + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; \ + \ + KKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + LOAD_MT(KT_tile, k_ + kk, Dk * Hk, B) \ + simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); \ + } \ + \ + KKtK_tile = KKt_tile; \ + SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) \ + \ + simdgroup_float8x8 Tinv, P; \ + AT(P, 0) = AT(KKtK_tile, 0); \ + AT(P, 1) = AT(KKtK_tile, 1); \ + SUB(Tinv, I_tile, KKtK_tile) \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int step = 1; (1 << step) < C; step++) { \ + simdgroup_multiply(P, P, P); \ + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); \ + } \ + \ + WS_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + SCALE(K_tile, beta_fm) \ + simdgroup_multiply(W_tile, Tinv, K_tile); \ + SCALE(W_tile, gamma_fm) \ + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); \ + } \ + \ + SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) \ + \ + LOAD_M(V_tile, v_ + dv_idx, Dv * Hv, B) \ + SCALE(V_tile, beta_fm) \ + simdgroup_multiply(U_tile, Tinv, V_tile); \ + SUB(delta_tile, U_tile, WS_tile) \ + \ + tmp_tile = make_filled_simdgroup_matrix(0.f); \ + QKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(Q_tile, q_ + kk, Hk * Dk, B) \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); \ + SCALE(Q_tile, gamma_fm) \ + simdgroup_multiply_accumulate( \ + tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); \ + } \ + \ + SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) \ + \ + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); \ + \ + if (fm < valid_rows) { \ + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); \ + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); \ + } \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) \ + simdgroup_multiply(KD_tile, K_tile, delta_tile); \ + FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) \ + } \ + } + +template +[[kernel]] void gated_delta_fused_chunk( + 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)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + 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); + + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + + ((thread_index_in_simdgroup / 2) % 4); // row coordinate of the held tile + const short fn = (qid & 2) * 2 + + (thread_index_in_simdgroup % 2) * 2; // column coordinate of the held tile + + auto dv_idx = thread_position_in_grid.y * 8; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + +#define OUTPUT(T) \ + if (true) { \ + simdgroup_store(T, y); \ + return; \ + } + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + simdgroup_float8x8 S_tile[Dk / 8]; + + // simdgroup matrices + simdgroup_float8x8 V_tile, K_tile, KT_tile, Q_tile; + simdgroup_float8x8 W_tile, U_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + // tiles for WY form computation + simdgroup_float8x8 KKtK_tile, KKt_tile; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + simdgroup_float8x8 I_tile = make_filled_simdgroup_matrix(0.f); + AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; + AT(I_tile, 1) = (fm == fn + 1) ? 1.0f : 0.0f; + + // load initial state into registers + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); + } + + int t = 0; + MLX_MTL_PRAGMA_UNROLL + for (; t + C <= T; t += C) { + PROCESS_CHUNK_SG(false, S_tile, C); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + PROCESS_CHUNK_SG(true, S_tile, short(T - t)); + } + + MLX_MTL_PRAGMA_UNROLL + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); + } +} + +/* + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + */ +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* 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)]], + 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; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } +} \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h new file mode 100644 index 000000000..15b98fd21 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h @@ -0,0 +1,527 @@ +#pragma once + +#include + +#include +#include + +#include "steel/gemm/nax.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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + // 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + // 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 +[[kernel]] void gated_delta_fused_nax( + 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)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + 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); + + 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)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile 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} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + 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); + } + }; + + 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)) + : 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; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + tmp_tile.clear(); + QKt_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) + + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 42855e14d..be1ebff59 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -596,6 +596,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -651,14 +657,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1897,6 +1896,7 @@ template < 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> @@ -2276,6 +2276,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -2437,11 +2438,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2485,13 +2487,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2508,113 +2515,61 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/metal/quantized_nax.h b/Source/Cmlx/mlx-generated/metal/quantized_nax.h index ed32eb59a..a34774992 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/quantized_nax.h @@ -653,14 +653,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -793,14 +786,7 @@ struct QuantizedBlockLoader< return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1215,7 +1201,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1334,7 +1321,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_gather_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1477,11 +1465,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1510,13 +1499,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1543,12 +1537,11 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1556,145 +1549,119 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); + NAXTile Dtile; + Dtile.clear(); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); + } - volatile int compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } - } - xn += BK; - loader_w.next(); - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); + (void)compiler_barrier; + } + } - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + xn += BK; + loader_w.next(); + } - volatile int compiler_barrier; + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); + threadgroup_barrier(mem_flags::mem_threadgroup); - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + } - threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); - } + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } 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 8fdd08762..6797bdf3e 100644 --- a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal @@ -28,30 +28,35 @@ using namespace metal; qk_dim, \ value_dim) -#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ - instantiate_kernel( \ - "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ - sdpa_vector_2pass_1_gqa, \ - type, \ - qk_dim, \ - value_dim, \ - 8, \ +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, g, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #g "_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + g, \ hpt) #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ + instantiate_sdpa_vector(type, 96, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ - instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ - instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ + instantiate_sdpa_vector(type, 512, 512) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 8, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 12, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 16, 2) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ instantiate_sdpa_vector_aggregation(type, 192) \ - instantiate_sdpa_vector_aggregation(type, 256) + instantiate_sdpa_vector_aggregation(type, 256) \ + instantiate_sdpa_vector_aggregation(type, 512) instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h index 2f466a567..4594f0389 100644 --- a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h +++ b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h @@ -512,6 +512,7 @@ template const int num_kv_heads = tpg.x; const int num_q_heads = num_kv_heads * G; const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx; const int chunk = (N + blocks - 1) / blocks; const int kstart = block_idx * chunk; @@ -520,9 +521,9 @@ template const int s0 = kstart + cchunk * sub; const int s1 = min(kend, s0 + sub); - const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + - simd_lid * qk_per_thread; - const device T* vp = values + kv_head_idx * v_head_stride + + const device T* kp = keys + kv_batch_head_idx * k_head_stride + + s0 * k_seq_stride + simd_lid * qk_per_thread; + const device T* vp = values + kv_batch_head_idx * v_head_stride + s0 * v_seq_stride + simd_lid * v_per_thread; U q[HPT][qk_per_thread]; 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 1926aa3cd..bc833474e 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 @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "../../../steel/attn/attn.h" @@ -63,6 +63,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -114,13 +115,15 @@ template < constexpr short LDQ_tgp = BD + padQ; constexpr short LDK_tgp = BK + padK; - constexpr short LDV_tgp = BD + padV; + constexpr short LDV_tgp = BV + padV; constexpr short tgp_mem_0 = (BK + padK) * (BD); - constexpr short tgp_mem_1 = BK * (BD + padV); + constexpr short tgp_mem_1 = BK * (BV + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = (WN == 2) && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -150,7 +153,7 @@ template < using VBlockLoader = BlockLoaderT< /* typename T = */ T, /* short BROWS = */ BK, - /* short BCOLS = */ BD, + /* short BCOLS = */ BV, /* short kDstStrRow = */ LDV_tgp, /* short kDstStrCol = */ 1, /* short reduction_dim = */ 0, @@ -168,8 +171,15 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; + + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert( + BD % WN == 0 && BV % WN == 0, + "The head dims must split evenly across WN"); - constexpr int kNWarps = WM * WN; + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -180,14 +190,19 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + constexpr int TVh = BV / (kFragSize * WN); + constexpr int BVh = BV / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -195,22 +210,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BVh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -262,17 +300,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -373,7 +445,7 @@ template < // Load V blocks if (!align_K && kb == (params->NK_aligned)) { - loader_v.load_safe(short2(BD, params->kL_rem)); + loader_v.load_safe(short2(BV, params->kL_rem)); } else { loader_v.load_unsafe(); } @@ -425,7 +497,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TVh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -461,10 +533,10 @@ template < threadgroup_barrier(mem_flags::mem_none); // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BVh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal index 9efa26fe0..7d0fe1453 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal @@ -1,31 +1,36 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. // clang-format off #include "../../../utils.h" #include "../../../steel/attn/kernels/steel_attention.h" -#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ - instantiate_kernel( \ - "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ - "_wm" #wm "_wn" #wn "_mask" #mname, \ - attention, dtype, bq, bk, bd, wm, wn, mtype, float) +#define instantiate_attn(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention, dtype, bq, bk, bd, bv, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) +#define instantiate_attn_helper(iname, itype, bq, bk, bd, bv, wm, wn) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, iname, itype) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, bool_, bool) -#define instantiate_attn_mask_helper(iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, bool_, bool) +#define instantiate_attn_shapes_helper(iname, itype) \ + instantiate_attn_helper(iname, itype, 32, 16, 192, 192, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 16, 128, 128, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 96, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 64, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 80, 80, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 72, 72, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 64, 64, 4, 1) -instantiate_attn_mask_helper(float16, half); -instantiate_attn_mask_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float16, half); +instantiate_attn_shapes_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float32, float); -instantiate_attn_mask_helper(float32, float); +instantiate_attn_helper(float16, half, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(float16, half, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(float32, float, 32, 16, 256, 256, 4, 2); // clang-format on 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 5010f1dee..bf2f026ba 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 @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "../../../steel/attn/nax.h" #include "../../../steel/attn/params.h" @@ -76,6 +76,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -136,11 +137,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -421,8 +424,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } @@ -489,15 +492,10 @@ template < // Head-dim split attention kernel /////////////////////////////////////////////////////////////////////////////// -// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup -// accumulator working set of attention_nax (TD output fragments plus the S -// fragments) is what gates tensor-unit throughput, so this kernel splits the -// head dim across the WN = 2 simdgroups of the second warp dimension: each -// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, -// halving its accumulator set. The pair exchanges its partial Q@K.T sums -// through threadgroup memory, then both simdgroups run softmax redundantly on -// the full S tile (the row statistics are cheap) and each accumulates P@V for -// its own half of Dv. +// Split wide heads across WN simdgroups to reduce the accumulator working set. +// Each group owns one D slice. The groups exchange their partial Q@K.T scores +// through threadgroup memory and reduce them in the same order. Each group then +// runs softmax and accumulates its P@V slice. // clang-format off template < @@ -556,9 +554,9 @@ template < constexpr short kU = 16; // The WM simdgroups along the first warp dimension split the Q sequence; - // the WN simdgroups along the second split the head dim. The exchange - // below reduces exactly one peer, so WN is fixed at 2. - static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + // the WN simdgroups along the second split the head dim. + static_assert( + WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4"); constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, @@ -579,20 +577,20 @@ template < constexpr int BDh = BD / WN; static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); - static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + static_assert(TK % 2 == 0, "S fragments are computed pair by pair"); const short row_group = simd_group_id / WN; - const short d_half = simd_group_id % WN; + const short d_group = simd_group_id % WN; using otile_t = NAXTile; otile_t Otile; Otile.clear(); const short tm = kU * TQ * row_group; - Q += tm * int(params->Q_strides[2]) + d_half * BDh; - K += d_half * BDh; - V += d_half * BDh; - O += tm * int(params->O_strides[2]) + d_half * BDh; + Q += tm * int(params->Q_strides[2]) + d_group * BDh; + K += d_group * BDh; + V += d_group * BDh; + O += tm * int(params->O_strides[2]) + d_group * BDh; constexpr short kRowsPT = otile_t::kRowsPerThread; @@ -632,12 +630,11 @@ template < using stile_t = NAXTile; constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; - // One slot per (row group, half): a fragment pair in per-lane-linear - // layout. Both halves share the fragment-to-lane mapping, so the - // exchange needs no coordinate math. - threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + // One slot per (row group, D group) in per-lane-linear layout. All groups + // share the fragment-to-lane mapping. + threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32]; - // Keep the simdgroup's Q half resident in registers for the whole KV + // Keep the simdgroup's Q slice resident in registers for the whole KV // loop: TDh fragments of T are cheap next to the accumulators. NAXTile Qtiles[TDh]; STEEL_PRAGMA_UNROLL @@ -662,7 +659,7 @@ template < stile_t Stile; Stile.clear(); - // S = Q @ K.T, this half of D only, exchanged pair by pair. + // S = Q @ K.T for this D slice. STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { STEEL_PRAGMA_UNROLL @@ -686,26 +683,37 @@ template < Ktile.frag_at(1, 0), metal::true_type{}); } + } - // Exchange the partial pair and reduce. - threadgroup AccumType* slot = s_xchg[row_group][d_half]; - thread auto& s0 = Stile.frag_at(0, ik); - thread auto& s1 = Stile.frag_at(0, ik + 1); - const short base = short(simd_lane_id) * (2 * kEPF); + // Exchange all partial scores and reduce them in a fixed order. + threadgroup AccumType* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - slot[base + i] = s0[i]; - slot[base + kEPF + i] = s1[i]; + slot[base + ik * kEPF + i] = s[i]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - s0[i] += peer[base + i]; - s1[i] += peer[base + kEPF + i]; + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short peer_group = 1; peer_group < WN; peer_group++) { + const threadgroup AccumType* peer = s_xchg[row_group][peer_group]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += peer[base + ik * kEPF + i]; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); // Scale S STEEL_PRAGMA_UNROLL @@ -826,7 +834,7 @@ template < } } - // Do softmax (redundantly per half; the row statistics are cheap) + // Do softmax in each D group; the row statistics are cheap. metal::vec new_max; metal::vec factor; STEEL_PRAGMA_UNROLL @@ -854,7 +862,7 @@ template < simdgroup_barrier(mem_flags::mem_none); - // O = P @ V, this half of Dv only. + // O = P @ V for this D slice. STEEL_PRAGMA_UNROLL for (short id = 0; id < TDh; id += 2) { STEEL_PRAGMA_UNROLL diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h index 4c055e69d..54e45d1e9 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h @@ -20,9 +20,10 @@ template < [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { @@ -44,8 +45,19 @@ template < using loader_b_t = typename gemm_kernel::loader_b_t; using mma_t = typename gemm_kernel::mma_t; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } @@ -53,188 +65,127 @@ template < threadgroup T As[gemm_kernel::tgp_mem_size_a]; threadgroup T Bs[gemm_kernel::tgp_mem_size_b]; - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; const size_t c_row_long = size_t(c_row); const size_t c_col_long = size_t(c_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row)); const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col)); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[c_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (rhs_indices[c_row + n] != index) { - offset_next = n; - index_next = rhs_indices[c_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); + // Prepare threadgroup loading operations + thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); + thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id); + + // Prepare iterations + const int gemm_k_iterations = params->gemm_k_iterations_aligned; - // Prepare threadgroup loading operations - thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); - thread loader_b_t loader_b( - B + index * params->batch_stride_b, - params->ldb, + // Do unaligned K iterations first + if (!align_K) { + const int k_last = params->gemm_k_iterations_aligned * BK; + const int k_remain = params->K - k_last; + const size_t k_jump_a = + transpose_a ? params->lda * size_t(k_last) : size_t(k_last); + const size_t k_jump_b = + transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); + + // Move loader source ahead to end + loader_a.src += k_jump_a; + loader_b.src += k_jump_b; + + // Load tile + const short2 tile_dims_A = + transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); + const short2 tile_dims_B = + transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); + + loader_a.load_safe(tile_dims_A); + loader_b.load_safe(tile_dims_B); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Do matmul + mma_op.mma(As, Bs); + + // Reset source back to start + loader_a.src -= k_jump_a; + loader_b.src -= k_jump_b; + } + + const short lbk = 0; + + // Tile aligned don't check + if ((tgp_bm == BM) && (align_N || tgp_bn == BN)) { + gemm_kernel::gemm_loop( + As, Bs, - simd_group_id, - simd_lane_id); - - // Prepare iterations - const int gemm_k_iterations = params->gemm_k_iterations_aligned; - - // Do unaligned K iterations first - if (!align_K) { - const int k_last = params->gemm_k_iterations_aligned * BK; - const int k_remain = params->K - k_last; - const size_t k_jump_a = - transpose_a ? params->lda * size_t(k_last) : size_t(k_last); - const size_t k_jump_b = - transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); - - // Move loader source ahead to end - loader_a.src += k_jump_a; - loader_b.src += k_jump_b; - - // Load tile - const short2 tile_dims_A = - transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); - const short2 tile_dims_B = - transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); - - loader_a.load_safe(tile_dims_A); - loader_b.load_safe(tile_dims_B); + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - threadgroup_barrier(mem_flags::mem_threadgroup); + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Do matmul - mma_op.mma(As, Bs); + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Reset source back to start - loader_a.src -= k_jump_a; - loader_b.src -= k_jump_b; - } + // Nothing aligned so check both rows and cols + else { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Matrix level aligned never check - if (align_M && align_N) { - for (int k = 0; k < gemm_k_iterations; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Load elements into threadgroup - loader_a.load_unsafe(); - loader_b.load_unsafe(); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Multiply and accumulate threadgroup elements - mma_op.mma(As, Bs); - - // Prepare for next iteration - loader_a.next(); - loader_b.next(); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } else { - const short lbk = 0; - - // Tile aligned don't check - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } - - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - - // Nothing aligned so check both rows and cols - else { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - } + // Store results to device memory + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + mma_op.store_result(C, params->ldd); + } else { + mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h index 7fec27060..afd645ff9 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h @@ -20,9 +20,11 @@ template < gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], + uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { constexpr short SM = BM / WM; @@ -31,106 +33,66 @@ gather_mm_rhs_nax( constexpr short TM = SM / 16; constexpr short TN = SN / 16; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; - const size_t c_row_long = size_t(c_row); - const size_t c_col_long = size_t(c_col); + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + const short sgp_sm = short(clamp(tgp_bm - tm, 0, int(SM))); + const short sgp_sn = + align_N ? SN : short(clamp(params->N - (c_col + tn), 0, int(SN))); + const size_t c_row_long = size_t(c_row + tm); + const size_t c_col_long = size_t(c_col + tn); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - rhs_indices += c_row; - - const short tm = SM * (simd_group_id / WN); - const short tn = SN * (simd_group_id % WN); - - const int sgp_sm_int = - align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); - const short sgp_sm = short(sgp_sm_int); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); - const int sgp_sn_int = - align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); - const short sgp_sn = short(sgp_sn_int); - const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); + dispatch_bool(align_K, [&](auto kAlignedK) { + dispatch_bool(sgp_sm == SM, [&](auto kAlignedM) { + dispatch_bool(sgp_sn == SN, [&](auto kAlignedN) { + NAXTile Ctile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + params->K, + params->gemm_k_iterations_aligned, + sgp_sm, + sgp_sn); - A += transpose_a ? tm : (tm * params->lda); - B += transpose_b ? (tn * params->ldb) : tn; - C += tm * params->ldd + tn; - rhs_indices += tm; - - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[0]; - short offset_next = 0; - int n = 0; - while (n < sgp_sm) { - n++; - offset = offset_next; - index = index_next; - offset_next = sgp_sm; - for (; n < sgp_sm; n++) { - if (rhs_indices[n] != index) { - offset_next = n; - index_next = rhs_indices[n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - NAXTile Ctile; - - dispatch_bool(align_K, [&](auto kAlignedK) { - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K - T, - SM, - SN, - SK, - BK, - transpose_a, - transpose_b, - kAlignedM.value, - kAlignedN.value, - kAlignedK.value, - AccumType>; - Ctile = do_gemm( - A, - B + index * params->batch_stride_b, - params->lda, - params->ldb, - params->K, - params->gemm_k_iterations_aligned, - sgp_sm, - sgp_sn); - - if constexpr (kAlignedN.value) { - if (offset_next - offset == SM) { - Ctile.store(C, int(params->ldd)); - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(SN, offset_next)); - } - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(sgp_sn, offset_next)); - } - }); + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.store(C, int(params->ldd)); + } else { + Ctile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } }); }); - } + }); } diff --git a/Source/Cmlx/mlx-generated/metal/unary_ops.h b/Source/Cmlx/mlx-generated/metal/unary_ops.h index a2dc80426..5909350a7 100644 --- a/Source/Cmlx/mlx-generated/metal/unary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/unary_ops.h @@ -35,7 +35,7 @@ struct Abs { return x; }; complex64_t operator()(complex64_t x) thread { - return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; + return {hypot(x.real, x.imag), 0}; }; }; @@ -308,7 +308,7 @@ struct Round { struct Sigmoid { template T operator()(T x) thread { - auto y = 1 / (1 + metal::exp(metal::abs(x))); + auto y = 1 / (1 + metal::precise::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -325,8 +325,8 @@ struct Sign { if (x == complex64_t(0)) { return x; } - return x / - (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag); + auto r = hypot(x.real, x.imag); + return {x.real / r, x.imag / r}; }; }; @@ -378,7 +378,7 @@ struct Sqrt { auto b_abs = metal::precise::sqrt((r - x.real) / 2.0); auto b = metal::copysign(b_abs, x.imag); return {a, b}; - } + }; }; struct Rsqrt { diff --git a/Source/Cmlx/mlx-generated/metal/utils.h b/Source/Cmlx/mlx-generated/metal/utils.h index cbf5798a1..7c82bf3cd 100644 --- a/Source/Cmlx/mlx-generated/metal/utils.h +++ b/Source/Cmlx/mlx-generated/metal/utils.h @@ -353,6 +353,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -470,3 +493,50 @@ template <> inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } + +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 583bba8ac..0ec148ade 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -609,6 +609,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -664,14 +670,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1910,6 +1909,7 @@ template < 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> @@ -2289,6 +2289,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -2450,11 +2451,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2498,13 +2500,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2521,113 +2528,61 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/quantized_nax.cpp b/Source/Cmlx/mlx-generated/quantized_nax.cpp index 123a4a37c..05277f64d 100644 --- a/Source/Cmlx/mlx-generated/quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/quantized_nax.cpp @@ -666,14 +666,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -806,14 +799,7 @@ struct QuantizedBlockLoader< return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1228,7 +1214,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1347,7 +1334,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_gather_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1490,11 +1478,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1523,13 +1512,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1556,12 +1550,11 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1569,147 +1562,121 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); + NAXTile Dtile; + Dtile.clear(); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); + } - volatile int compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } - } - xn += BK; - loader_w.next(); - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); + (void)compiler_barrier; + } + } - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + xn += BK; + loader_w.next(); + } - volatile int compiler_barrier; + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); + threadgroup_barrier(mem_flags::mem_threadgroup); - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + } - threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); - } + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/scan.cpp b/Source/Cmlx/mlx-generated/scan.cpp index 1fd1d74dc..9a291a0b8 100644 --- a/Source/Cmlx/mlx-generated/scan.cpp +++ b/Source/Cmlx/mlx-generated/scan.cpp @@ -273,7 +273,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/steel_attention.cpp b/Source/Cmlx/mlx-generated/steel_attention.cpp index 75c2edbdd..acd115969 100644 --- a/Source/Cmlx/mlx-generated/steel_attention.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention.cpp @@ -1780,7 +1780,7 @@ struct GEMMKernel { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h" -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. using namespace mlx::steel; @@ -1844,6 +1844,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -1895,13 +1896,15 @@ template < constexpr short LDQ_tgp = BD + padQ; constexpr short LDK_tgp = BK + padK; - constexpr short LDV_tgp = BD + padV; + constexpr short LDV_tgp = BV + padV; constexpr short tgp_mem_0 = (BK + padK) * (BD); - constexpr short tgp_mem_1 = BK * (BD + padV); + constexpr short tgp_mem_1 = BK * (BV + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = (WN == 2) && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -1931,7 +1934,7 @@ template < using VBlockLoader = BlockLoaderT< /* typename T = */ T, /* short BROWS = */ BK, - /* short BCOLS = */ BD, + /* short BCOLS = */ BV, /* short kDstStrRow = */ LDV_tgp, /* short kDstStrCol = */ 1, /* short reduction_dim = */ 0, @@ -1949,8 +1952,15 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; + + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert( + BD % WN == 0 && BV % WN == 0, + "The head dims must split evenly across WN"); - constexpr int kNWarps = WM * WN; + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -1961,14 +1971,19 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + constexpr int TVh = BV / (kFragSize * WN); + constexpr int BVh = BV / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -1976,22 +1991,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BVh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -2043,17 +2081,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -2154,7 +2226,7 @@ template < // Load V blocks if (!align_K && kb == (params->NK_aligned)) { - loader_v.load_safe(short2(BD, params->kL_rem)); + loader_v.load_safe(short2(BV, params->kL_rem)); } else { loader_v.load_unsafe(); } @@ -2206,7 +2278,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TVh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -2242,10 +2314,10 @@ template < threadgroup_barrier(mem_flags::mem_none); // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BVh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; diff --git a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp index df0831701..3d4595d3c 100644 --- a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp @@ -1290,7 +1290,7 @@ struct BlockSwizzle { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h" -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. using namespace mlx::steel; @@ -1364,6 +1364,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -1424,11 +1425,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -1709,8 +1712,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } @@ -1777,15 +1780,10 @@ template < // Head-dim split attention kernel /////////////////////////////////////////////////////////////////////////////// -// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup -// accumulator working set of attention_nax (TD output fragments plus the S -// fragments) is what gates tensor-unit throughput, so this kernel splits the -// head dim across the WN = 2 simdgroups of the second warp dimension: each -// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, -// halving its accumulator set. The pair exchanges its partial Q@K.T sums -// through threadgroup memory, then both simdgroups run softmax redundantly on -// the full S tile (the row statistics are cheap) and each accumulates P@V for -// its own half of Dv. +// Split wide heads across WN simdgroups to reduce the accumulator working set. +// Each group owns one D slice. The groups exchange their partial Q@K.T scores +// through threadgroup memory and reduce them in the same order. Each group then +// runs softmax and accumulates its P@V slice. // clang-format off template < @@ -1844,9 +1842,9 @@ template < constexpr short kU = 16; // The WM simdgroups along the first warp dimension split the Q sequence; - // the WN simdgroups along the second split the head dim. The exchange - // below reduces exactly one peer, so WN is fixed at 2. - static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + // the WN simdgroups along the second split the head dim. + static_assert( + WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4"); constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, @@ -1867,20 +1865,20 @@ template < constexpr int BDh = BD / WN; static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); - static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + static_assert(TK % 2 == 0, "S fragments are computed pair by pair"); const short row_group = simd_group_id / WN; - const short d_half = simd_group_id % WN; + const short d_group = simd_group_id % WN; using otile_t = NAXTile; otile_t Otile; Otile.clear(); const short tm = kU * TQ * row_group; - Q += tm * int(params->Q_strides[2]) + d_half * BDh; - K += d_half * BDh; - V += d_half * BDh; - O += tm * int(params->O_strides[2]) + d_half * BDh; + Q += tm * int(params->Q_strides[2]) + d_group * BDh; + K += d_group * BDh; + V += d_group * BDh; + O += tm * int(params->O_strides[2]) + d_group * BDh; constexpr short kRowsPT = otile_t::kRowsPerThread; @@ -1920,12 +1918,11 @@ template < using stile_t = NAXTile; constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; - // One slot per (row group, half): a fragment pair in per-lane-linear - // layout. Both halves share the fragment-to-lane mapping, so the - // exchange needs no coordinate math. - threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + // One slot per (row group, D group) in per-lane-linear layout. All groups + // share the fragment-to-lane mapping. + threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32]; - // Keep the simdgroup's Q half resident in registers for the whole KV + // Keep the simdgroup's Q slice resident in registers for the whole KV // loop: TDh fragments of T are cheap next to the accumulators. NAXTile Qtiles[TDh]; STEEL_PRAGMA_UNROLL @@ -1950,7 +1947,7 @@ template < stile_t Stile; Stile.clear(); - // S = Q @ K.T, this half of D only, exchanged pair by pair. + // S = Q @ K.T for this D slice. STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { STEEL_PRAGMA_UNROLL @@ -1974,26 +1971,37 @@ template < Ktile.frag_at(1, 0), metal::true_type{}); } + } - // Exchange the partial pair and reduce. - threadgroup AccumType* slot = s_xchg[row_group][d_half]; - thread auto& s0 = Stile.frag_at(0, ik); - thread auto& s1 = Stile.frag_at(0, ik + 1); - const short base = short(simd_lane_id) * (2 * kEPF); + // Exchange all partial scores and reduce them in a fixed order. + threadgroup AccumType* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - slot[base + i] = s0[i]; - slot[base + kEPF + i] = s1[i]; + slot[base + ik * kEPF + i] = s[i]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - s0[i] += peer[base + i]; - s1[i] += peer[base + kEPF + i]; + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short peer_group = 1; peer_group < WN; peer_group++) { + const threadgroup AccumType* peer = s_xchg[row_group][peer_group]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += peer[base + ik * kEPF + i]; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); // Scale S STEEL_PRAGMA_UNROLL @@ -2114,7 +2122,7 @@ template < } } - // Do softmax (redundantly per half; the row statistics are cheap) + // Do softmax in each D group; the row statistics are cheap. metal::vec new_max; metal::vec factor; STEEL_PRAGMA_UNROLL @@ -2142,7 +2150,7 @@ template < simdgroup_barrier(mem_flags::mem_none); - // O = P @ V, this half of Dv only. + // O = P @ V for this D slice. STEEL_PRAGMA_UNROLL for (short id = 0; id < TDh; id += 2) { STEEL_PRAGMA_UNROLL diff --git a/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp b/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp index 3219ac624..bd6979f6a 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp @@ -33,9 +33,10 @@ template < [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { @@ -57,8 +58,19 @@ template < using loader_b_t = typename gemm_kernel::loader_b_t; using mma_t = typename gemm_kernel::mma_t; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } @@ -66,188 +78,127 @@ template < threadgroup T As[gemm_kernel::tgp_mem_size_a]; threadgroup T Bs[gemm_kernel::tgp_mem_size_b]; - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; const size_t c_row_long = size_t(c_row); const size_t c_col_long = size_t(c_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row)); const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col)); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[c_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (rhs_indices[c_row + n] != index) { - offset_next = n; - index_next = rhs_indices[c_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); + // Prepare threadgroup loading operations + thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); + thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id); + + // Prepare iterations + const int gemm_k_iterations = params->gemm_k_iterations_aligned; - // Prepare threadgroup loading operations - thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); - thread loader_b_t loader_b( - B + index * params->batch_stride_b, - params->ldb, + // Do unaligned K iterations first + if (!align_K) { + const int k_last = params->gemm_k_iterations_aligned * BK; + const int k_remain = params->K - k_last; + const size_t k_jump_a = + transpose_a ? params->lda * size_t(k_last) : size_t(k_last); + const size_t k_jump_b = + transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); + + // Move loader source ahead to end + loader_a.src += k_jump_a; + loader_b.src += k_jump_b; + + // Load tile + const short2 tile_dims_A = + transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); + const short2 tile_dims_B = + transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); + + loader_a.load_safe(tile_dims_A); + loader_b.load_safe(tile_dims_B); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Do matmul + mma_op.mma(As, Bs); + + // Reset source back to start + loader_a.src -= k_jump_a; + loader_b.src -= k_jump_b; + } + + const short lbk = 0; + + // Tile aligned don't check + if ((tgp_bm == BM) && (align_N || tgp_bn == BN)) { + gemm_kernel::gemm_loop( + As, Bs, - simd_group_id, - simd_lane_id); - - // Prepare iterations - const int gemm_k_iterations = params->gemm_k_iterations_aligned; - - // Do unaligned K iterations first - if (!align_K) { - const int k_last = params->gemm_k_iterations_aligned * BK; - const int k_remain = params->K - k_last; - const size_t k_jump_a = - transpose_a ? params->lda * size_t(k_last) : size_t(k_last); - const size_t k_jump_b = - transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); - - // Move loader source ahead to end - loader_a.src += k_jump_a; - loader_b.src += k_jump_b; - - // Load tile - const short2 tile_dims_A = - transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); - const short2 tile_dims_B = - transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); - - loader_a.load_safe(tile_dims_A); - loader_b.load_safe(tile_dims_B); + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - threadgroup_barrier(mem_flags::mem_threadgroup); + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Do matmul - mma_op.mma(As, Bs); + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Reset source back to start - loader_a.src -= k_jump_a; - loader_b.src -= k_jump_b; - } + // Nothing aligned so check both rows and cols + else { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Matrix level aligned never check - if (align_M && align_N) { - for (int k = 0; k < gemm_k_iterations; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Load elements into threadgroup - loader_a.load_unsafe(); - loader_b.load_unsafe(); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Multiply and accumulate threadgroup elements - mma_op.mma(As, Bs); - - // Prepare for next iteration - loader_a.next(); - loader_b.next(); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } else { - const short lbk = 0; - - // Tile aligned don't check - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } - - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - - // Nothing aligned so check both rows and cols - else { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - } + // Store results to device memory + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + mma_op.store_result(C, params->ldd); + } else { + mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp index 7312ba458..8ff7cdb90 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp @@ -33,9 +33,11 @@ template < gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], + uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { constexpr short SM = BM / WM; @@ -44,108 +46,68 @@ gather_mm_rhs_nax( constexpr short TM = SM / 16; constexpr short TN = SN / 16; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; - const size_t c_row_long = size_t(c_row); - const size_t c_col_long = size_t(c_col); + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + const short sgp_sm = short(clamp(tgp_bm - tm, 0, int(SM))); + const short sgp_sn = + align_N ? SN : short(clamp(params->N - (c_col + tn), 0, int(SN))); + const size_t c_row_long = size_t(c_row + tm); + const size_t c_col_long = size_t(c_col + tn); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - rhs_indices += c_row; - - const short tm = SM * (simd_group_id / WN); - const short tn = SN * (simd_group_id % WN); - - const int sgp_sm_int = - align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); - const short sgp_sm = short(sgp_sm_int); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); - const int sgp_sn_int = - align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); - const short sgp_sn = short(sgp_sn_int); - const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); - - A += transpose_a ? tm : (tm * params->lda); - B += transpose_b ? (tn * params->ldb) : tn; - C += tm * params->ldd + tn; - rhs_indices += tm; - - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[0]; - short offset_next = 0; - int n = 0; - while (n < sgp_sm) { - n++; - offset = offset_next; - index = index_next; - offset_next = sgp_sm; - for (; n < sgp_sm; n++) { - if (rhs_indices[n] != index) { - offset_next = n; - index_next = rhs_indices[n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - NAXTile Ctile; - - dispatch_bool(align_K, [&](auto kAlignedK) { - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K - T, - SM, - SN, - SK, - BK, - transpose_a, - transpose_b, - kAlignedM.value, - kAlignedN.value, - kAlignedK.value, - AccumType>; - Ctile = do_gemm( - A, - B + index * params->batch_stride_b, - params->lda, - params->ldb, - params->K, - params->gemm_k_iterations_aligned, - sgp_sm, - sgp_sn); - - if constexpr (kAlignedN.value) { - if (offset_next - offset == SM) { - Ctile.store(C, int(params->ldd)); - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(SN, offset_next)); - } - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(sgp_sn, offset_next)); - } - }); + dispatch_bool(align_K, [&](auto kAlignedK) { + dispatch_bool(sgp_sm == SM, [&](auto kAlignedM) { + dispatch_bool(sgp_sn == SN, [&](auto kAlignedN) { + NAXTile Ctile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + params->K, + params->gemm_k_iterations_aligned, + sgp_sm, + sgp_sn); + + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.store(C, int(params->ldd)); + } else { + Ctile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } }); }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/unary_ops.cpp b/Source/Cmlx/mlx-generated/unary_ops.cpp index 3a680b7f2..4ab9ac5f4 100644 --- a/Source/Cmlx/mlx-generated/unary_ops.cpp +++ b/Source/Cmlx/mlx-generated/unary_ops.cpp @@ -349,11 +349,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -446,7 +446,7 @@ struct Abs { return x; }; complex64_t operator()(complex64_t x) thread { - return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; + return {hypot(x.real, x.imag), 0}; }; }; @@ -719,7 +719,7 @@ struct Round { struct Sigmoid { template T operator()(T x) thread { - auto y = 1 / (1 + metal::exp(metal::abs(x))); + auto y = 1 / (1 + metal::precise::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -736,8 +736,8 @@ struct Sign { if (x == complex64_t(0)) { return x; } - return x / - (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag); + auto r = hypot(x.real, x.imag); + return {x.real / r, x.imag / r}; }; }; @@ -789,7 +789,7 @@ struct Sqrt { auto b_abs = metal::precise::sqrt((r - x.real) / 2.0); auto b = metal::copysign(b_abs, x.imag); return {a, b}; - } + }; }; struct Rsqrt { diff --git a/Source/Cmlx/mlx-generated/utils.cpp b/Source/Cmlx/mlx-generated/utils.cpp index 18d9644d0..c8cf785e0 100644 --- a/Source/Cmlx/mlx-generated/utils.cpp +++ b/Source/Cmlx/mlx-generated/utils.cpp @@ -739,6 +739,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: @@ -1133,6 +1134,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -1251,6 +1275,53 @@ inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} + /////////////////////////////////////////////////////////////////////////////// )preamble"; } diff --git a/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml index 2e15efd93..c8f86420a 100644 --- a/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml +++ b/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml @@ -89,6 +89,7 @@ runs: --exclude libcuda* \ --exclude libcudnn* \ --exclude libcufft* \ + --exclude libcusolver* \ --exclude libnccl* \ --exclude libnvrtc* fi diff --git a/Source/Cmlx/mlx/.github/actions/setup/action.yml b/Source/Cmlx/mlx/.github/actions/setup/action.yml index 510086390..6ba85028b 100644 --- a/Source/Cmlx/mlx/.github/actions/setup/action.yml +++ b/Source/Cmlx/mlx/.github/actions/setup/action.yml @@ -88,8 +88,7 @@ runs: max-size: |- ${{ case(inputs.ccache-key == 'release', case(startsWith(inputs.toolkit, 'cuda'), - case(runner.os == 'Linux', '300MB', - '160MB'), + '300MB', '60MB'), case(startsWith(inputs.toolkit, 'cuda'), case(runner.os == 'Linux' && runner.arch == 'x64', '600MB', @@ -170,9 +169,9 @@ runs: } PACKAGES: | { - "cuda-12.6": ["cudart_12.6", "nvcc_12.6", "cublas_12.6", "cublas_dev_12.6", "cufft_12.6", "cufft_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", "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", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"], + "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" diff --git a/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml index a29e3c1dc..e06087650 100644 --- a/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml +++ b/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml @@ -46,9 +46,4 @@ runs: echo "No matching backend wheel to install" exit 1 fi - if ${{ runner.os == 'macOS' }} ; then - # FIXME: run.py does not quit in macOS CI. - python -m unittest discover -v python/tests - else - uv run python/tests/run.py -v --failfast - fi + uv run python/tests/run.py -v --failfast diff --git a/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js b/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js index f5a05396c..296c8f2be 100644 --- a/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js +++ b/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js @@ -8,6 +8,7 @@ * - 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 @@ -95,6 +96,8 @@ async function fetchCollaborators(owner, repo, token) { * open: number, * closed: number, * merged: number, + * total: number, + * lowPriority: number, * mostRecentCreatedAt: Date * }> */ @@ -110,11 +113,18 @@ function aggregateUserStats(pullRequests) { 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') { @@ -138,6 +148,7 @@ function aggregateUserStats(pullRequests) { * - 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 */ @@ -152,6 +163,7 @@ function filterEligibleUsers(stats, collaboratorLogins) { 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); } diff --git a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml b/Source/Cmlx/mlx/.github/workflows/build_and_test.yml index d590d10b1..a4b5ddf8f 100644 --- a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml +++ b/Source/Cmlx/mlx/.github/workflows/build_and_test.yml @@ -85,7 +85,7 @@ jobs: matrix: macos-target: ['14.0', '15.0', '26.2'] toolkit: ['cpu', 'metal', 'jit'] - runs-on: 'macos-26' + runs-on: 'macos-26-xlarge' needs: check_lint steps: - uses: actions/checkout@v7 diff --git a/Source/Cmlx/mlx/.github/workflows/release.yml b/Source/Cmlx/mlx/.github/workflows/release.yml index b53f003ff..0445da036 100644 --- a/Source/Cmlx/mlx/.github/workflows/release.yml +++ b/Source/Cmlx/mlx/.github/workflows/release.yml @@ -141,7 +141,7 @@ jobs: strategy: matrix: python-version: *pyver - runs-on: 'macos-26' + runs-on: 'macos-26-xlarge' env: *build-env steps: - uses: actions/checkout@v7 diff --git a/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml b/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml index b9b71eee8..ec4025893 100644 --- a/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml +++ b/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml @@ -15,6 +15,7 @@ permissions: jobs: update_bypass_list: name: 'Update bypass list' + if: github.repository == 'ml-explore/mlx' runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v7 diff --git a/Source/Cmlx/mlx/.gitignore b/Source/Cmlx/mlx/.gitignore index 161d8d67e..f5d859ff3 100644 --- a/Source/Cmlx/mlx/.gitignore +++ b/Source/Cmlx/mlx/.gitignore @@ -79,3 +79,4 @@ uv.lock .cache/ # vim *.swp +wip \ No newline at end of file diff --git a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md index 7005396e7..a97754922 100644 --- a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md +++ b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md @@ -22,6 +22,7 @@ MLX was developed with contributions from the following individuals: - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. - katlun-lgtm: Added `reflect` and `symmetric` padding modes. - Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting. +- Oleksandr Zakharchuk: Added the cuSOLVER `cholesky` for the CUDA backend. diff --git a/Source/Cmlx/mlx/CMakeLists.txt b/Source/Cmlx/mlx/CMakeLists.txt index feb7ce5eb..01b329b2f 100644 --- a/Source/Cmlx/mlx/CMakeLists.txt +++ b/Source/Cmlx/mlx/CMakeLists.txt @@ -43,6 +43,7 @@ option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON) option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON) option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF) option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON) +option(MLX_EXPORT_ALL_SYMBOLS "Export all symbols in the library" OFF) option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF) option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF) option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF) @@ -395,7 +396,7 @@ if(MLX_BUILD_PYTHON_BINDINGS) FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v2.15.0 + GIT_TAG v3.0.1 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) diff --git a/Source/Cmlx/mlx/CODE_OF_CONDUCT.md b/Source/Cmlx/mlx/CODE_OF_CONDUCT.md deleted file mode 100644 index e8d213c7b..000000000 --- a/Source/Cmlx/mlx/CODE_OF_CONDUCT.md +++ /dev/null @@ -1,132 +0,0 @@ -# Contributor Covenant Code of Conduct - -## Our Pledge - -We as members, contributors, and leaders pledge to make participation in our -community a harassment-free experience for everyone, regardless of age, body -size, visible or invisible disability, ethnicity, sex characteristics, gender -identity and expression, level of experience, education, socio-economic status, -nationality, personal appearance, race, caste, color, religion, or sexual -identity and orientation. - -We pledge to act and interact in ways that contribute to an open, welcoming, -diverse, inclusive, and healthy community. - -## Our Standards - -Examples of behavior that contributes to a positive environment for our -community include: - -* Demonstrating empathy and kindness toward other people -* Being respectful of differing opinions, viewpoints, and experiences -* Giving and gracefully accepting constructive feedback -* Accepting responsibility and apologizing to those affected by our mistakes, - and learning from the experience -* Focusing on what is best not just for us as individuals, but for the overall - community - -Examples of unacceptable behavior include: - -* The use of sexualized language or imagery, and sexual attention or advances of - any kind -* Trolling, insulting or derogatory comments, and personal or political attacks -* Public or private harassment -* Publishing others' private information, such as a physical or email address, - without their explicit permission -* Other conduct which could reasonably be considered inappropriate in a - professional setting - -## Enforcement Responsibilities - -Community leaders are responsible for clarifying and enforcing our standards of -acceptable behavior and will take appropriate and fair corrective action in -response to any behavior that they deem inappropriate, threatening, offensive, -or harmful. - -Community leaders have the right and responsibility to remove, edit, or reject -comments, commits, code, wiki edits, issues, and other contributions that are -not aligned to this Code of Conduct, and will communicate reasons for moderation -decisions when appropriate. - -## Scope - -This Code of Conduct applies within all community spaces, and also applies when -an individual is officially representing the community in public spaces. -Examples of representing our community include using an official e-mail address, -posting via an official social media account, or acting as an appointed -representative at an online or offline event. - -## Enforcement - -Instances of abusive, harassing, or otherwise unacceptable behavior may be -reported to the community leaders responsible for enforcement at -[opensource-conduct@group.apple.com](mailto:opensource-conduct@group.apple.com). -All complaints will be reviewed and investigated promptly and fairly. - -All community leaders are obligated to respect the privacy and security of the -reporter of any incident. - -## Enforcement Guidelines - -Community leaders will follow these Community Impact Guidelines in determining -the consequences for any action they deem in violation of this Code of Conduct: - -### 1. Correction - -**Community Impact**: Use of inappropriate language or other behavior deemed -unprofessional or unwelcome in the community. - -**Consequence**: A private, written warning from community leaders, providing -clarity around the nature of the violation and an explanation of why the -behavior was inappropriate. A public apology may be requested. - -### 2. Warning - -**Community Impact**: A violation through a single incident or series of -actions. - -**Consequence**: A warning with consequences for continued behavior. No -interaction with the people involved, including unsolicited interaction with -those enforcing the Code of Conduct, for a specified period of time. This -includes avoiding interactions in community spaces as well as external channels -like social media. Violating these terms may lead to a temporary or permanent -ban. - -### 3. Temporary Ban - -**Community Impact**: A serious violation of community standards, including -sustained inappropriate behavior. - -**Consequence**: A temporary ban from any sort of interaction or public -communication with the community for a specified period of time. No public or -private interaction with the people involved, including unsolicited interaction -with those enforcing the Code of Conduct, is allowed during this period. -Violating these terms may lead to a permanent ban. - -### 4. Permanent Ban - -**Community Impact**: Demonstrating a pattern of violation of community -standards, including sustained inappropriate behavior, harassment of an -individual, or aggression toward or disparagement of classes of individuals. - -**Consequence**: A permanent ban from any sort of public interaction within the -community. - -## Attribution - -This Code of Conduct is adapted from the [Contributor Covenant][homepage], -version 2.1, available at -[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]. - -Community Impact Guidelines were inspired by -[Mozilla's code of conduct enforcement ladder][Mozilla CoC]. - -For answers to common questions about this code of conduct, see the FAQ at -[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at -[https://www.contributor-covenant.org/translations][translations]. - -[homepage]: https://www.contributor-covenant.org -[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html -[Mozilla CoC]: https://github.com/mozilla/diversity -[FAQ]: https://www.contributor-covenant.org/faq -[translations]: https://www.contributor-covenant.org/translations diff --git a/Source/Cmlx/mlx/CONTRIBUTING.md b/Source/Cmlx/mlx/CONTRIBUTING.md index eaccfec88..d3a4f444c 100644 --- a/Source/Cmlx/mlx/CONTRIBUTING.md +++ b/Source/Cmlx/mlx/CONTRIBUTING.md @@ -1,39 +1,97 @@ -# Contributing to MLX +# Welcome to the MLX community! -We want to make contributing to this project as easy and transparent as -possible. +We welcome contributors of all backgrounds and experience levels. A community with diverse perspectives builds software with lasting impact. -## AI Usage Policy +[🤝 How you can help](#how-you-can-help) · [⚙️ Setting up your environment](#setting-up-your-environment) · [📝 Submitting issues and pull requests](#submitting-issues-and-pull-requests) · [🔀 Pull requests](#pull-requests) · [❤️ Code of conduct](#code-of-conduct) -AI-generated code is allowed. What is not allowed is submitting code you do not -understand. You are 100% responsible for every line, however it was produced, -and must explicitly disclose the manner in which AI was employed. +## How you can help -It is strictly prohibited to use AI to write your posts for you (bug reports, -feature requests, pull request descriptions, Github discussions, responding to -humans, ...). +- Report bugs with clear, reproducible steps +- Improve documentation to make the project more accessible +- Triage issues by providing feedback, testing, and validation +- Propose and implement enhancements of all sizes +- Share your experience by speaking about MLX at conferences or meetups -## Pull Requests +## Setting up your environment -- Make sure new code is covered by tests. Add new tests if not, and confirm - the new tests fail in the main branch. -- If performance may be impacted, run benchmarks for both the main branch and - the pull request. +See [Build and Install](https://ml-explore.github.io/mlx/build/html/install.html) for setup instructions. + +## Submitting issues and pull requests + +We maintain high standards to ensure quality across the project. + +Your contributions are evaluated on: + +- Technical quality and correctness +- Adherence to existing conventions and architectural patterns +- Demonstrated understanding of the implementation and its implications +- Clarity of communication +- Maintainability + +You may use productivity tools, including AI-assisted coding, to help you work more efficiently. +However, you remain fully accountable for all submitted code, issues, pull requests, and comments. +AI-generated content must meet the same standards as human-written contributions. Review, validate, and, as needed, +refine or rearrange AI-assisted work so the final contribution reflects your human creativity, understanding, and +control. Use your voice and expression when producing written materials. Misuse of AI tools in your contributions +and conversations may be considered a violation of our [Code of Conduct](#code-of-conduct). + +### Issues + +Before submitting an issue: + +- Search [existing issues](../../issues) to avoid duplicates. +- Write a clear, concise title and a detailed description. Ensure all information is accurate and fully understood by you. +- Include precise steps to reproduce the issue, expected behavior, and actual behavior. Verify these steps yourself. +- Specify your environment, MLX version, and any relevant configurations. +- Explain the severity and impact of the problem. +- For security issues, please see the [MLX security advisories](https://github.com/ml-explore/mlx?tab=security-ov-file). Do **not** open a public GitHub issue. + +#### Proposing features + +We welcome feature requests. Describe the proposed changes needed and why they matter. Please do not start with a pull request. Use the issue template that best matches your needs. + +### Pull requests + +⚠️ **Important:** We recommend an issue for larger changes to the codebase. + +Pull requests represent proposed solutions or enhancements. When submitting a pull request, verify it meets these expectations: + +- They should match an open issue but may not always. +- Address the stated problem or feature request completely and effectively. You must fully understand and validate your proposed solution. +- Solutions must be thorough, handle edge cases, and integrate cleanly. +- Include comprehensive tests that validate your changes and prevent regressions. +- If performance may be impacted, run benchmarks for both the main branch and the pull request. - When providing benchmarking results, include scripts and reproduction steps. -- Format the code with `uvx pre-commit run --all` before submitting a pull - request. You can also install git hooks to run it automatically: + +#### Pull request style guide and format + +Follow our coding style and formatting to maintain a consistent, readable codebase. + +- Write comprehensive and clear documentation. Concisely explain the 'why' behind complex decisions. +- Write clear, concise, and descriptive commit messages. +- Format the code with `pre-commit run --all` before submitting a pull request. You can also install git hooks to run it automatically: ```shell pip install pre-commit pre-commit install ``` -## Issues +#### Testing + +All contributions require thorough testing. + +**Local Testing:** Ensure all existing tests pass before submitting: + +```bash +python3 python/tests/run.py +``` + +**Automated Tests:** New features and bug fixes require corresponding automated tests that validate the intended behavior and prevent regressions. + +#### Running CI -We use GitHub issues to track public bugs. Please ensure your description is -clear and has sufficient instructions to be able to reproduce the issue. +CI runs automatically on every pull request and needs to be approved by a maintainer before it will run. -## License +## Code of conduct -By contributing to MLX, you agree that your contributions will be licensed -under the LICENSE file in the root directory of this source tree. +We are committed to fostering a community where different experiences and perspectives come together to create and collaborate. Good collaboration depends on honest feedback and respect for the time and effort every contributor brings to this project. [Please review our Code of Conduct](https://github.com/ml-explore/.github/blob/main/CODE_OF_CONDUCT.md); all community members are expected to adhere to it. diff --git a/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py b/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py new file mode 100644 index 000000000..277897699 --- /dev/null +++ b/Source/Cmlx/mlx/benchmarks/python/gated_delta_bench.py @@ -0,0 +1,145 @@ +import argparse +import csv +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) + mx.synchronize() + + s = time.perf_counter_ns() + for _ in range(N_iter_bench): + f(*args) + mx.synchronize() + e = time.perf_counter_ns() + return (e - s) * 1e-9 # total seconds for N_iter_bench * N_iter_func calls + + +def do_kernel_bench(f, *args): + ys = [] + for _ in range(N_iter_func): + out, hf = f(*args) + ys.append(out) + ys.append(hf) + mx.eval(ys) + return ys + + +def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): + mx.random.seed(42) + 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.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + + 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) + / denom + * 1e3 + ) + + speedups = [] + 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) + / denom + * 1e3 + ) + speedups.append(ms_seq / ms_c if ms_c > 0 else float("nan")) + except Exception as ex: + print(f" chunk {C} failed: {ex}") + speedups.append(float("nan")) + + return shape_str, f"{ms_seq:.3f}", speedups, ms_seq + + +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] + else: + Bs = [1, 8, 16] + Ts = [8, 512, 1024, 2048] + Hks = [16] + Hvs = [32] + Dks = [128] + Dvs = [128] + + chunk_sizes = [0, 8, 16] + non_zero_Cs = [C for C in chunk_sizes if C != 0] + + headers = ["B", "T", "Hk", "Hv", "Dk", "Dv", "time_seq (ms)"] + [ + f"C={C} (speedup)" for C in 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))) + + 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( + B, T, Hk, Hv, Dk, Dv, chunk_sizes + ) + row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}", base_time_s] + for speed in speedups: + row.append(f"{(base_time / speed):<8.2f} ({speed:<5.2f}x)") + + print(fmt.format(*row), end="") + print(f"{RESET}") + + rows.append(row) + + if to_csv: + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(headers) + writer.writerows(rows) + print(f"\nResults also written to {csv_path}") + + +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") + args = parser.parse_args() + + run_benchmark(args.full, to_csv=args.csv, csv_path=args.csv_out) diff --git a/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py new file mode 100644 index 000000000..83257a4cd --- /dev/null +++ b/Source/Cmlx/mlx/benchmarks/python/gather_qqmm_bench.py @@ -0,0 +1,108 @@ +# Copyright © 2026 Apple Inc. + +import argparse + +import mlx.core as mx +from time_utils import time_fn + +SHAPES = { + "qwen3.6-35b-a3b-up": { + "experts": 256, + "top_k": 8, + "input": 2048, + "output": 512, + }, + "qwen3.6-35b-a3b-down": { + "experts": 256, + "top_k": 8, + "input": 512, + "output": 2048, + }, +} + + +def make_weights(shape): + weight = mx.random.uniform( + low=-0.5, + high=0.5, + shape=(shape["experts"], shape["output"], shape["input"]), + ).astype(mx.bfloat16) + quantized, scales = mx.quantize( + weight, + mode="nvfp4", + global_scale=mx.array(1.0, dtype=mx.float32), + ) + global_scale = mx.ones((shape["experts"],), dtype=mx.float32) + mx.eval(quantized, scales, global_scale) + return quantized, scales, global_scale + + +def gather_qqmm(x, weight, scales, global_scale_x, global_scale_w, rhs, sorted_): + return mx.gather_qqmm( + x, + weight, + scales, + rhs_indices=rhs, + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + sorted_indices=sorted_, + ) + + +def benchmark(shape, workload, tokens): + weight, scales, global_scale_w = make_weights(shape) + global_scale_x = mx.array(1.0, dtype=mx.float32) + + if workload == "prompt": + routes = tokens * shape["top_k"] + x = mx.random.uniform( + low=-0.5, high=0.5, shape=(routes, 1, shape["input"]) + ).astype(mx.bfloat16) + rhs = mx.repeat( + mx.arange(shape["experts"], dtype=mx.uint32), + (routes + shape["experts"] - 1) // shape["experts"], + )[:routes] + sorted_ = True + elif workload == "matrix": + x = mx.random.uniform( + low=-0.5, + high=0.5, + shape=(shape["top_k"], tokens, shape["input"]), + ).astype(mx.bfloat16) + rhs = (mx.arange(shape["top_k"], dtype=mx.uint32) * 17) % shape["experts"] + sorted_ = False + else: + x = mx.random.uniform( + low=-0.5, high=0.5, shape=(1, 1, 1, shape["input"]) + ).astype(mx.bfloat16) + rhs = ( + (mx.arange(shape["top_k"], dtype=mx.uint32) * 17) % shape["experts"] + ).reshape(1, shape["top_k"]) + sorted_ = False + + mx.eval(x, rhs) + time_fn( + gather_qqmm, + x, + weight, + scales, + global_scale_x, + global_scale_w, + rhs, + sorted_, + msg=( + f"{workload} routes={rhs.size} " f"N={shape['output']} K={shape['input']}" + ), + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--shape", choices=SHAPES, default="qwen3.6-35b-a3b-up") + parser.add_argument( + "--workload", choices=("prompt", "matrix", "decode"), default="prompt" + ) + parser.add_argument("--tokens", type=int, default=2048) + args = parser.parse_args() + benchmark(SHAPES[args.shape], args.workload, args.tokens) diff --git a/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py new file mode 100644 index 000000000..3939d2a8a --- /dev/null +++ b/Source/Cmlx/mlx/benchmarks/python/qqmm_bench.py @@ -0,0 +1,62 @@ +# Copyright © 2026 Apple Inc. + +import argparse + +import mlx.core as mx +from time_utils import time_fn + +SHAPES = { + "qwen3.6-35b-a3b-up": (512, 2048), + "qwen3.6-35b-a3b-down": (2048, 512), +} + + +def qqmm(x, weight, scales, global_scale_x, global_scale_w): + return mx.qqmm( + x, + weight, + scales, + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + + +def benchmark(shape, tokens): + output_dims, input_dims = shape + mx.random.seed(0) + global_scale_x = mx.array(1.0, dtype=mx.float32) + global_scale_w = mx.array(1.0, dtype=mx.float32) + x = mx.random.uniform(low=-0.5, high=0.5, shape=(tokens, input_dims)).astype( + mx.bfloat16 + ) + weight = mx.random.uniform( + low=-0.5, high=0.5, shape=(output_dims, input_dims) + ).astype(mx.bfloat16) + weight, scales = mx.quantize( + weight, + mode="nvfp4", + global_scale=global_scale_w, + ) + mx.eval(x, weight, scales) + + time_fn( + qqmm, + x, + weight, + scales, + global_scale_x, + global_scale_w, + msg=( + f"mode=nvfp4 dtype=bfloat16 tokens={tokens} " + f"N={output_dims} K={input_dims}" + ), + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--shape", choices=SHAPES, default="qwen3.6-35b-a3b-up") + parser.add_argument("--tokens", type=int, default=2048) + args = parser.parse_args() + benchmark(SHAPES[args.shape], args.tokens) diff --git a/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py b/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py new file mode 100644 index 000000000..805bc7c1d --- /dev/null +++ b/Source/Cmlx/mlx/benchmarks/python/swiglu_bench.py @@ -0,0 +1,77 @@ +# Copyright © 2025 Apple Inc. + +import mlx.core as mx +import mlx.nn as nn +from time_utils import time_fn + +SEQ_LENS = [512, 713, 1024, 2048, 4123, 8192, 8192 * 2] +MAX_UNSORTED_N = 1024 + +# https://huggingface.co/zai-org/GLM-5.3-Flash-BF16/blob/main/config.json +# https://huggingface.co/zai-org/GLM-5.3/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-35B-A3B/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-122B-A10B/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-397B-A17B/blob/main/config.json +CONFIGS = { + "glm-5.3-flash": (4096, 2048, 288, 8), + "glm-5.3": (6144, 2048, 256, 8), + "qwen3.5-35b-a3b": (2048, 512, 256, 8), + "qwen3.5-122b-a10b": (3072, 1024, 256, 8), + "qwen3.5-397b-a17b": (4096, 1024, 512, 10), +} + + +def gather_sort(x, indices): + N, M = indices.shape + indices = indices.flatten() + order = mx.argsort(indices) + inv_order = mx.argsort(order) + return x.flatten(0, -3)[order // M], indices[order], inv_order + + +def scatter_unsort(x, inv_order, shape=None): + x = x[inv_order] + if shape is not None: + x = mx.unflatten(x, 0, shape) + return x + + +def time_gather_mm(name, D, M, E, I): + 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) + mx.eval(w1, w2, w3) + + def gather_mm(x, w1, w2, w3, indices, sort): + idx = indices + inv_order = None + if sort: + x, idx, inv_order = gather_sort(x, indices) + gate = mx.gather_mm( + x, w1.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort + ) + up = mx.gather_mm(x, w2.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort) + x = mx.gather_mm( + nn.silu(gate) * up, + w3.swapaxes(-1, -2), + rhs_indices=idx, + sorted_indices=sort, + ) + if sort: + x = scatter_unsort(x, inv_order, indices.shape) + return x + + 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) + mx.eval(x, indices, sorted_indices) + + label = f"{name} N={N}" + time_fn(gather_mm, x, w1, w2, w3, indices, True, msg=f"{label} swiglu") + + +if __name__ == "__main__": + for name, config in CONFIGS.items(): + time_gather_mm(name, *config) diff --git a/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py new file mode 100644 index 000000000..2e80ec2fa --- /dev/null +++ b/Source/Cmlx/mlx/benchmarks/python/swiglu_qmm_bench.py @@ -0,0 +1,72 @@ +# Copyright © 2026 Apple Inc. + +import mlx.core as mx +import mlx.nn as nn +from time_utils import time_fn + +SEQ_LENS = [512, 713, 1024, 2048, 4123, 8192, 8192 * 2] +MODES = ["affine", "nvfp4", "mxfp8"] + +# https://huggingface.co/zai-org/GLM-5.3-Flash-BF16/blob/main/config.json +# https://huggingface.co/zai-org/GLM-5.3/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-35B-A3B/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-122B-A10B/blob/main/config.json +# https://huggingface.co/Qwen/Qwen3.5-397B-A17B/blob/main/config.json +CONFIGS = { + "glm-5.3-flash": (4096, 2048, 288, 8), + "glm-5.3": (6144, 2048, 256, 8), + "qwen3.5-35b-a3b": (2048, 512, 256, 8), + "qwen3.5-122b-a10b": (3072, 1024, 256, 8), + "qwen3.5-397b-a17b": (4096, 1024, 512, 10), +} + + +def gather_sort(x, indices): + N, M = indices.shape + indices = indices.flatten() + order = mx.argsort(indices) + inv_order = mx.argsort(order) + return x.flatten(0, -3)[order // M], indices[order], inv_order + + +def scatter_unsort(x, inv_order, shape=None): + x = x[inv_order] + if shape is not None: + x = mx.unflatten(x, 0, shape) + return x + + +def time_gather_qmm(name, D, M, E, I, 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) + w1, w2, w3 = (mx.quantize(w, mode=mode) for w in (w1, w2, w3)) + mx.eval(w1, w2, w3) + + def gather_qmm(x, w1, w2, w3, indices, sort): + idx = indices + inv_order = None + if sort: + x, idx, inv_order = gather_sort(x, indices) + kwargs = dict(transpose=True, mode=mode, rhs_indices=idx, sorted_indices=sort) + gate = mx.gather_qmm(x, *w1, **kwargs) + up = mx.gather_qmm(x, *w2, **kwargs) + x = mx.gather_qmm(nn.silu(gate) * up, *w3, **kwargs) + if sort: + x = scatter_unsort(x, inv_order, indices.shape) + return x + + 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) + mx.eval(x, indices) + + label = f"{name} {mode} N={N}" + time_fn(gather_qmm, x, w1, w2, w3, indices, True, msg=f"{label} swiglu") + + +if __name__ == "__main__": + for mode in MODES: + for name, config in CONFIGS.items(): + time_gather_qmm(name, *config, mode) diff --git a/Source/Cmlx/mlx/docs/src/dev/extensions.rst b/Source/Cmlx/mlx/docs/src/dev/extensions.rst index 89c8991e5..5f5fab642 100644 --- a/Source/Cmlx/mlx/docs/src/dev/extensions.rst +++ b/Source/Cmlx/mlx/docs/src/dev/extensions.rst @@ -711,7 +711,7 @@ build utilities defined in :mod:`mlx.extension`: package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, extras_require={"dev":[]}, zip_safe=False, - python_requires=">=3.8", + python_requires=">=3.10", ) .. note:: diff --git a/Source/Cmlx/mlx/docs/src/python/memory_management.rst b/Source/Cmlx/mlx/docs/src/python/memory_management.rst index f708efbfd..b0564bd0f 100644 --- a/Source/Cmlx/mlx/docs/src/python/memory_management.rst +++ b/Source/Cmlx/mlx/docs/src/python/memory_management.rst @@ -7,6 +7,7 @@ Memory Management :toctree: _autosummary get_active_memory + get_array_buffer_size get_peak_memory reset_peak_memory get_cache_memory diff --git a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt b/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt index eedc3a7da..014f71240 100644 --- a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt +++ b/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt @@ -10,15 +10,23 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON) option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON) # ----------------------------- Dependencies ----------------------------- + +include(FetchContent) + +# Avoid warning about DOWNLOAD_EXTRACT_TIMESTAMP in CMake 3.24: +cmake_policy(SET CMP0135 NEW) + find_package( - Python 3.8 + Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) -execute_process( - COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir - OUTPUT_STRIP_TRAILING_WHITESPACE - OUTPUT_VARIABLE nanobind_ROOT) -find_package(nanobind CONFIG REQUIRED) +FetchContent_Declare( + nanobind + GIT_REPOSITORY https://github.com/wjakob/nanobind.git + GIT_TAG v3.0.1 + GIT_SHALLOW TRUE + EXCLUDE_FROM_ALL) +FetchContent_MakeAvailable(nanobind) execute_process( COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir diff --git a/Source/Cmlx/mlx/examples/extensions/pyproject.toml b/Source/Cmlx/mlx/examples/extensions/pyproject.toml index 560a58bc2..2386c9ec4 100644 --- a/Source/Cmlx/mlx/examples/extensions/pyproject.toml +++ b/Source/Cmlx/mlx/examples/extensions/pyproject.toml @@ -2,7 +2,6 @@ requires = [ "setuptools>=42", "cmake>=3.25", - "mlx>=0.18.0", - "nanobind==2.15.0", + "mlx>=0.31.2", ] build-backend = "setuptools.build_meta" diff --git a/Source/Cmlx/mlx/examples/extensions/requirements.txt b/Source/Cmlx/mlx/examples/extensions/requirements.txt index 917d125ee..41dce3ca7 100644 --- a/Source/Cmlx/mlx/examples/extensions/requirements.txt +++ b/Source/Cmlx/mlx/examples/extensions/requirements.txt @@ -1,4 +1,3 @@ setuptools>=42 cmake>=3.25 mlx>=0.31.2 -nanobind==2.15.0 diff --git a/Source/Cmlx/mlx/mlx/CMakeLists.txt b/Source/Cmlx/mlx/mlx/CMakeLists.txt index 0df4f4234..f6d1a4aa2 100644 --- a/Source/Cmlx/mlx/mlx/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/CMakeLists.txt @@ -18,6 +18,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/memory.cpp ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h) # Define MLX_VERSION only in the version.cpp file. @@ -27,11 +28,13 @@ target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR}) target_link_libraries(mlx PRIVATE $) # Do not export symbols by default. -set_target_properties( - mlx mlx_version - PROPERTIES VISIBILITY_INLINES_HIDDEN ON - CXX_VISIBILITY_PRESET hidden - CUDA_VISIBILITY_PRESET hidden) +if(NOT MLX_EXPORT_ALL_SYMBOLS) + set_target_properties( + mlx mlx_version + PROPERTIES VISIBILITY_INLINES_HIDDEN ON + CXX_VISIBILITY_PRESET hidden + CUDA_VISIBILITY_PRESET hidden) +endif() # Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries. set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT) @@ -52,6 +55,9 @@ endif() if(WIN32) # windows.h defines min/max macros that conflict with std::min/std::max. target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN) + # TODO: Remove after removing std::atomic_store/load. + target_compile_definitions( + mlx PUBLIC _SILENCE_CXX20_OLD_SHARED_PTR_ATOMIC_SUPPORT_DEPRECATION_WARNING) endif() if(MSVC) # Unicode support in fmt does not compile in .cu files. diff --git a/Source/Cmlx/mlx/mlx/array.cpp b/Source/Cmlx/mlx/mlx/array.cpp index 07d3acd61..e6dbbe785 100644 --- a/Source/Cmlx/mlx/mlx/array.cpp +++ b/Source/Cmlx/mlx/mlx/array.cpp @@ -208,34 +208,33 @@ void array::copy_shared_buffer(const array& other) { copy_shared_buffer(other, other.strides(), other.flags(), other.data_size()); } -array::~array() { - if (array_desc_ == nullptr) { +void array::release(std::shared_ptr desc) { + if (desc == nullptr) { return; } // Detached/detaching - if (array_desc_->primitive == nullptr) { + if (desc->primitive == nullptr) { return; } // Break circular reference for non-detached arrays with siblings - if (auto n = siblings().size(); n > 0) { + if (auto n = desc->siblings.size(); n > 0) { bool do_detach = true; // If all siblings have siblings.size() references except - // the one we are currently destroying (which has siblings.size() + 1) + // the one we are currently releasing (which has siblings.size() + 1) // then there are no more external references - do_detach &= (array_desc_.use_count() == (n + 1)); - for (auto& s : siblings()) { + do_detach &= (desc.use_count() == (n + 1)); + for (auto& s : desc->siblings) { do_detach &= (s.array_desc_.use_count() == n); if (!do_detach) { break; } } if (do_detach) { - for (auto& s : siblings()) { + for (auto& s : desc->siblings) { for (auto& ss : s.siblings()) { - // Set to null here to avoid descending into array destructor - // for siblings + // Set to null here to avoid descending into reset() for siblings ss.array_desc_ = nullptr; } s.array_desc_->siblings.clear(); diff --git a/Source/Cmlx/mlx/mlx/array.h b/Source/Cmlx/mlx/mlx/array.h index 3f45e9cb9..cc29290d2 100644 --- a/Source/Cmlx/mlx/mlx/array.h +++ b/Source/Cmlx/mlx/mlx/array.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "mlx/allocator.h" @@ -80,13 +81,22 @@ class MLX_API array { array& operator=(array&& other) && = delete; /** Default copy and move constructors otherwise. */ - array& operator=(array&& other) & = default; array(const array& other) = default; array(array&& other) = default; + /** Assignment goes through release() to break sibling reference cycle. */ array& operator=(const array& other) & { - if (this->id() != other.id()) { - this->array_desc_ = other.array_desc_; + if (array_desc_ != other.array_desc_) { + // Hold the old descriptor while assigning, and then release it. Otherwise + // when |this| and |other| are siblings they would be both destroyed. + release(std::exchange(array_desc_, other.array_desc_)); + } + return *this; + } + + array& operator=(array&& other) & noexcept { + if (array_desc_ != other.array_desc_) { + release(std::exchange(array_desc_, std::move(other.array_desc_))); } return *this; } @@ -466,10 +476,12 @@ class MLX_API array { void copy_shared_buffer(const array& other); void overwrite_descriptor(const array& other) { - array_desc_ = other.array_desc_; + release(std::exchange(array_desc_, other.array_desc_)); } - ~array(); + ~array() { + release(std::move(array_desc_)); + } private: // Initialize the arrays data @@ -532,6 +544,12 @@ class MLX_API array { // the primitive which knows how to compute the array's data from its inputs // and the list of array's inputs for the primitive. std::shared_ptr array_desc_; + + /** + * Lets go of an ArrayDesc, breaking the sibling reference cycle of a + * multi-output primitive. + */ + static void release(std::shared_ptr desc); }; template diff --git a/Source/Cmlx/mlx/mlx/backend/common/compiled.h b/Source/Cmlx/mlx/mlx/backend/common/compiled.h index 8c6466da0..94ba38b8f 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/compiled.h +++ b/Source/Cmlx/mlx/mlx/backend/common/compiled.h @@ -37,9 +37,9 @@ void print_float_constant(std::ostream& os, const array& x) { auto old_precision = os.precision(); if constexpr (std::is_same_v) { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } else { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } os << value << std::setprecision(old_precision); } diff --git a/Source/Cmlx/mlx/mlx/backend/common/utils.cpp b/Source/Cmlx/mlx/mlx/backend/common/utils.cpp index ae169e35e..66764173e 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/utils.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include +#include #include "mlx/backend/common/utils.h" @@ -174,6 +175,22 @@ Dims get_2d_grid_dims_common( const Shape& shape, const Strides& strides, size_t divisor) { + // Shortcut when size is a multiple of divisor, for example: + // 36: {2 / 2, 4 / 2, 9 / 9} = {1, 2, 1} + Shape divided_shape = shape; + size_t remaining_divisor = divisor; + for (int i = 0; i < shape.size(); ++i) { + if (strides[i] == 0) { + continue; + } + auto factor = std::gcd(static_cast(shape[i]), remaining_divisor); + divided_shape[i] /= factor; + remaining_divisor /= factor; + } + if (remaining_divisor == 1) { + return get_2d_grid_dims_common(divided_shape, strides); + } + // Compute the 2d grid dimensions such that the total size of the grid is // divided by divisor. size_t grid_x = 1; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp index d93f74cf0..57c765bc6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp @@ -41,7 +41,7 @@ void arg_reduce_dispatch( switch (rtype) { case ArgReduce::ArgMin: { auto op = [](auto ind_x, auto x, auto ind_y, auto y) { - if (x < (*y)) { + if ((!mlx::core::isnan(*y) && mlx::core::isnan(x)) || x < (*y)) { (*y) = x; (*ind_y) = ind_x; } @@ -51,7 +51,7 @@ void arg_reduce_dispatch( } case ArgReduce::ArgMax: { auto op = [](auto ind_x, auto x, auto ind_y, auto y) { - if (x > (*y)) { + if ((!mlx::core::isnan(*y) && mlx::core::isnan(x)) || x > (*y)) { (*y) = x; (*ind_y) = ind_x; } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h index 5333aff9a..fb7db76f9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h @@ -16,7 +16,7 @@ namespace mlx::core { template struct VectorScalar { template - void operator()(const T* a, const T* b, U* dst, int size) { + void operator()(const T* a, const T* b, U* dst, size_t size) { T scalar = *b; constexpr int N = simd::max_size; while (size >= N) { @@ -36,7 +36,7 @@ struct VectorScalar { template struct ScalarVector { template - void operator()(const T* a, const T* b, U* dst, int size) { + void operator()(const T* a, const T* b, U* dst, size_t size) { T scalar = *a; constexpr int N = simd::max_size; while (size >= N) { @@ -56,7 +56,7 @@ struct ScalarVector { template struct VectorVector { template - void operator()(const T* a, const T* b, U* dst, int size) { + void operator()(const T* a, const T* b, U* dst, size_t size) { constexpr int N = simd::max_size; while (size >= N) { simd::store(dst, Op{}(simd::load(a), simd::load(b))); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h b/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h index d50751ce3..0bfdc0501 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary_ops.h @@ -25,7 +25,7 @@ using namespace mlx::core::simd; DEFAULT_BINARY_OP(Add, operator+) DEFAULT_BINARY_OP(ArcTan2, atan2) -DEFAULT_BINARY_OP(Divide, operator/) +DEFAULT_BINARY_OP(Divide, 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/distributed.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp index 6b788cde2..88ec269fe 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp @@ -98,7 +98,7 @@ void Recv::eval_cpu( void ReduceScatter::eval_cpu( const std::vector& inputs, std::vector& outputs) { - assert(inputs.size() == 0); + assert(inputs.size() == 1); assert(outputs.size() == 1); auto [in, copied] = ensure_row_contiguous(inputs[0], stream()); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp index bb2c06ea0..0b4654b22 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp @@ -22,8 +22,13 @@ CommandEncoder& get_command_encoder(Stream s) { } std::unordered_map& get_command_encoders() { - static thread_local std::unordered_map encoders; - return encoders; + if (is_main_thread()) { + static auto* encoders = new std::unordered_map(); + return *encoders; + } else { + static thread_local std::unordered_map encoders; + return encoders; + } } std::unordered_map& get_global_command_encoders() { diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp index 64b3fab33..04965690b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp @@ -68,6 +68,18 @@ const complex64_t Limits::max = const complex64_t Limits::min = -std::numeric_limits::infinity(); +template +struct ReductionAccumulator { + static constexpr int N = std::min(simd::max_size, simd::max_size); + // Widen to float32 only if the input is float16 (with N=1) or bfloat16 as it + // improves performance. + static constexpr bool widen_to_float = + (std::is_same_v || + (N == 1 && std::is_same_v)); + + using type = std::conditional_t; +}; + template void strided_reduce( const T* x, @@ -261,6 +273,23 @@ void reduction_op( } } +template +void float_reduction( + const array& x, + array& out, + const std::vector& axes, + U init) { + using AccT = ReductionAccumulator::type; + if constexpr (std::is_same_v) { + reduction_op(x, out, axes, init); + } else { + array temp(out.shape(), TypeToDtype(), nullptr, {}); + temp.set_data(allocator::malloc(temp.nbytes())); + reduction_op(x, temp, axes, static_cast(init)); + std::copy_n(temp.data(), out.size(), out.data()); + } +} + struct AndReduce { template bool operator()(bool x, T y) { @@ -420,13 +449,13 @@ void reduce_dispatch_sum_prod( if constexpr (std::is_integral_v && sizeof(InT) <= 4) { reduction_op(in, out, axes, 0); } else { - reduction_op(in, out, axes, 0); + float_reduction(in, out, axes, 0); } } else { if constexpr (std::is_integral_v && sizeof(InT) <= 4) { reduction_op(in, out, axes, 1); } else { - reduction_op(in, out, axes, 1); + float_reduction(in, out, axes, 1); } } } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp index ab2ca1436..6632e44e7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp @@ -269,6 +269,11 @@ void scan_dispatch( void Scan::eval_cpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); + if (out.size() == 0) { + out.set_data(allocator::malloc(0)); + return; + } + auto& encoder = cpu::get_command_encoder(stream()); // Ensure contiguity @@ -278,6 +283,9 @@ void Scan::eval_cpu(const std::vector& inputs, array& out) { encoder.add_temporary(in); } out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } encoder.set_input_array(in); encoder.set_output_array(out); 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 e11cc7f5f..5d6128328 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h @@ -233,6 +233,16 @@ Simd minimum(Simd a, Simd b) { return out; } +// Integer division by zero gives 0. See the scalar divide in base_simd.h. +template +Simd divide(Simd a, Simd b) { + if constexpr (std::is_integral_v) { + return select(b == Simd(0), Simd(0), Simd(a / b)); + } else { + return a / b; + } +} + template Simd remainder(Simd a, Simd b) { Simd r; @@ -268,15 +278,17 @@ Simd pow(Simd base, Simd exp) { return asd::pow(base.value, exp.value); } else { Simd res = 1; - // Raising an integer to a negative power is undefined - if (any(exp < 0)) { - return 0; - } while (any(exp > 0)) { res = select((exp & 1) != 0, res * base, res); base = select(exp > 0, base * base, base); exp = exp >> 1; } + if constexpr (std::is_signed_v) { + // Raising an integer to a negative power is undefined, so return 0 for + // those lanes. Select per lane: a single negative exponent must not zero + // the whole vector. exp stays negative through the shifts above. + res = select(exp < 0, Simd(0), res); + } return res; } } 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 1ae883f7e..626dd93df 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h @@ -13,6 +13,7 @@ #include // For _BitScanReverse #endif +#include "mlx/types/complex.h" #include "mlx/types/half_types.h" namespace mlx::core::simd { @@ -159,7 +160,7 @@ auto imag(Simd in) -> Simd { } template Simd isnan(Simd in) { - return std::isnan(in.value); + return mlx::core::isnan(in.value); } #define DEFAULT_BINARY(OP) \ @@ -203,12 +204,29 @@ Simd clz(Simd x_) { #endif } +// Integer division by zero traps on x86 and returns 0 on arm64. Define the +// quotient as 0 on all platforms, which keeps a == (a / b) * b + (a % b). +template +Simd 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); + } + } + return a / b; +} + template Simd remainder(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; T r; if constexpr (std::is_integral_v) { + if (b == 0) { + return a; + } r = a % b; } else { r = std::remainder(a, b); @@ -226,7 +244,7 @@ Simd maximum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -238,7 +256,7 @@ Simd minimum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp index 767e7ef1d..92378981b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp @@ -2,7 +2,6 @@ #include #include -#include #include #include "mlx/backend/common/utils.h" @@ -19,9 +18,9 @@ namespace { template bool nan_aware_less(T a, T b) { if constexpr (is_floating_point_v || std::is_same_v) { - if (std::isnan(a)) + if (mlx::core::isnan(a)) return false; - if (std::isnan(b)) + if (mlx::core::isnan(b)) return true; } return a < b; @@ -118,6 +117,10 @@ struct StridedIterator { template void sort(array& out, int axis) { + if (out.size() == 0) { + return; + } + // Get axis, shape and stride info axis = axis < 0 ? axis + out.ndim() : axis; size_t in_size = out.size(); @@ -149,6 +152,10 @@ void sort(array& out, int axis) { template void argsort(const array& in, array& out, int axis) { + if (in.size() == 0) { + return; + } + // Get axis, shape and stride info axis = axis < 0 ? axis + in.ndim() : axis; size_t n_rows = in.size() / in.shape(axis); @@ -198,10 +205,10 @@ void argsort(const array& in, array& out, int axis) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_point_v) { - if (std::isnan(v1)) + if constexpr (is_floating_point_v || std::is_same_v) { + if (mlx::core::isnan(v1)) return false; - if (std::isnan(v2)) + if (mlx::core::isnan(v2)) return true; } @@ -212,6 +219,10 @@ void argsort(const array& in, array& out, int axis) { template void partition(array& out, int axis, int kth) { + if (out.size() == 0) { + return; + } + // Get axis, shape and stride info axis = axis < 0 ? axis + out.ndim() : axis; size_t in_size = out.size(); @@ -246,6 +257,10 @@ void partition(array& out, int axis, int kth) { template void argpartition(const array& in, array& out, int axis, int kth) { + if (in.size() == 0) { + return; + } + // Get axis, shape and stride info axis = axis < 0 ? axis + in.ndim() : axis; size_t n_rows = in.size() / in.shape(axis); @@ -299,10 +314,10 @@ void argpartition(const array& in, array& out, int axis, int kth) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if constexpr (is_floating_point_v) { - if (std::isnan(v1)) + if constexpr (is_floating_point_v || std::is_same_v) { + if (mlx::core::isnan(v1)) return false; - if (std::isnan(v2)) + if (mlx::core::isnan(v2)) return true; } @@ -357,6 +372,10 @@ void ArgSort::eval_cpu(const std::vector& inputs, array& out) { // Allocate output out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + auto& encoder = cpu::get_command_encoder(stream()); encoder.set_input_array(in); encoder.set_input_array(out); @@ -384,6 +403,10 @@ void Sort::eval_cpu(const std::vector& inputs, array& out) { : CopyType::General; copy_cpu(in, out, ctype, stream()); + if (out.size() == 0) { + return; + } + auto& encoder = cpu::get_command_encoder(stream()); encoder.set_output_array(out); encoder.dispatch([out = array::unsafe_weak_copy(out), axis]() mutable { diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/unary_ops.h b/Source/Cmlx/mlx/mlx/backend/cpu/unary_ops.h index f441e88bd..e2dd7413f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/unary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/unary_ops.h @@ -154,19 +154,20 @@ struct ToFP8 { struct FromFP8 { template Simd operator()(Simd x) { - auto v = Simd(x & 127) << 7; + auto v = Simd(x & 127); + auto sign_bit = Simd((x >> 7) & 1) << 15; + auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; Simd out; if constexpr (simd::max_size >= N) { - auto converted = *(Simd*)(&v); + auto converted = *(Simd*)(&u); out = converted * 256.0; } else { for (int i = 0; i < N; ++i) { - auto converted = *(float16_t*)(&v[i]); + auto converted = *(float16_t*)(&u[i]); out[i] = converted * 256.0; } } - auto sign = Simd(x & 128); - return select(sign, -out, out); + return out; } float operator()(uint8_t x) { return (*this)(Simd(x)).value; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt index 9c8d3174c..50be6f404 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt @@ -9,6 +9,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/arange.cu ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/binary_two.cu + ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cu ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp ${CMAKE_CURRENT_SOURCE_DIR}/copy.cu ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_contiguous.cu @@ -21,6 +22,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cu ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cusolver_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp @@ -77,6 +79,10 @@ target_sources( add_library(mlx_dirs OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/dirs.cpp) target_include_directories(mlx_dirs PRIVATE "${PROJECT_SOURCE_DIR}") target_link_libraries(mlx PRIVATE $) +# TODO: Remove after removing std::atomic_store/load. +target_compile_definitions( + mlx_dirs + PUBLIC _SILENCE_CXX20_OLD_SHARED_PTR_ATOMIC_SUPPORT_DEPRECATION_WARNING) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary) @@ -181,6 +187,8 @@ set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES # Configure Windows CUDA DLL loading. if(WIN32) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/wddm.cpp) + target_link_libraries(mlx PRIVATE dxgi) set(MLX_CUDA_BIN_DIR "" CACHE STRING "Directory containing CUDA DLLs for Windows delay-loading") @@ -241,6 +249,7 @@ else() "$ORIGIN/../../nvidia/cuda_nvrtc/lib" "$ORIGIN/../../nvidia/cudnn/lib" "$ORIGIN/../../nvidia/cufft/lib" + "$ORIGIN/../../nvidia/cusolver/lib" "$ORIGIN/../../nvidia/nccl/lib") endif() endif() @@ -286,6 +295,9 @@ target_link_libraries(mlx PRIVATE CUDA::cublasLt) # Use cuFFT. target_link_libraries(mlx PRIVATE CUDA::cufft) +# Use cuSOLVER. +target_link_libraries(mlx PRIVATE CUDA::cusolver) + # Use NVRTC and driver APIs. target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp index b04d72da0..5f4b9f2d5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp @@ -1,13 +1,18 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/utils.h" #include "mlx/backend/gpu/device_info.h" +#include "mlx/device.h" #include "mlx/memory.h" #include "mlx/scheduler.h" #include "mlx/utils.h" +#ifdef _WIN32 +#include "mlx/backend/cuda/wddm.h" +#endif + #include #include @@ -64,9 +69,9 @@ inline void* unified_malloc(size_t size) { inline void unified_free(void* data) { if (supports_managed_memory()) { - CHECK_CUDA_ERROR(cudaFree(data)); + cudaFree(data); } else { - CHECK_CUDA_ERROR(cudaFreeHost(data)); + cudaFreeHost(data); } } @@ -186,7 +191,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (!buf) { // If we have a lot of memory pressure try to reclaim memory from the cache. int64_t mem_to_free = - get_active_memory() + get_cache_memory() + size - memory_limit_; + get_active_memory() + get_cache_memory() + size - get_memory_limit(); if (mem_to_free > 0) { buffer_cache_.release_cached_buffers(mem_to_free); } @@ -319,9 +324,9 @@ void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) { if (!stream) { stream = free_streams_[buf.device]; } - CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream)); + cudaFreeAsync(buf.data, stream); } else { - CHECK_CUDA_ERROR(cudaFree(buf.data)); + cudaFree(buf.data); } } } @@ -340,7 +345,13 @@ void CudaAllocator::reset_peak_memory() { } size_t CudaAllocator::get_memory_limit() { +#ifdef _WIN32 + int device = default_device().index; + return std::min( + memory_limit_, get_wddm_memory_limit(device, mem_pools_[device])); +#else return memory_limit_; +#endif } size_t CudaAllocator::set_memory_limit(size_t limit) { @@ -360,8 +371,17 @@ size_t CudaAllocator::set_cache_limit(size_t limit) { } void CudaAllocator::clear_cache() { - std::lock_guard lk(mutex_); - buffer_cache_.clear(); + { + std::lock_guard lk(mutex_); + buffer_cache_.clear(); + } + for (size_t i = 0; i < mem_pools_.size(); ++i) { + if (mem_pools_[i]) { + cu::device(static_cast(i)).make_current(); + CHECK_CUDA_ERROR(cudaStreamSynchronize(free_streams_[i])); + CHECK_CUDA_ERROR(cudaMemPoolTrimTo(mem_pools_[i], 0)); + } + } } CudaAllocator& allocator() { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu index 9609b8597..39abb274a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu @@ -35,8 +35,11 @@ struct ArgMin { __device__ IndexValPair operator()( const IndexValPair& best, const IndexValPair& current) { - if (best.val > current.val || - (best.val == current.val && best.index > current.index)) { + if ((cuda::std::isnan(current.val) && + (!cuda::std::isnan(best.val) || best.index > current.index)) || + (!cuda::std::isnan(best.val) && + (best.val > current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -50,7 +53,8 @@ struct ArgMin { uint32_t offset) { #pragma unroll for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if ((!cuda::std::isnan(best.val) && cuda::std::isnan(vals[i])) || + vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -68,8 +72,11 @@ struct ArgMax { __device__ IndexValPair operator()( const IndexValPair& best, const IndexValPair& current) { - if (best.val < current.val || - (best.val == current.val && best.index > current.index)) { + if ((cuda::std::isnan(current.val) && + (!cuda::std::isnan(best.val) || best.index > current.index)) || + (!cuda::std::isnan(best.val) && + (best.val < current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -83,7 +90,8 @@ struct ArgMax { uint32_t offset) { #pragma unroll for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if ((!cuda::std::isnan(best.val) && cuda::std::isnan(vals[i])) || + vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cholesky.cu b/Source/Cmlx/mlx/mlx/backend/cuda/cholesky.cu new file mode 100644 index 000000000..47c993fb0 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cholesky.cu @@ -0,0 +1,141 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/cusolver_utils.h" +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/cuda/utils.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/primitives.h" + +#include +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +// potrf writes only one triangle; zero the other one to match the CPU op. +template +__global__ void zero_triangle(T* out, IdxT size, int32_t n, bool zero_below) { + IdxT index = cg::this_grid().thread_rank(); + if (index >= size) { + return; + } + int32_t row = (index / n) % n; + int32_t col = index % n; + if (zero_below ? (col < row) : (col > row)) { + out[index] = T(0); + } +} + +// potrfBatched wants a device array of per matrix pointers. +__global__ void +fill_matrix_pointers(void** ptrs, float* base, int64_t n2, int64_t batch) { + int64_t i = cg::this_grid().thread_rank(); + if (i < batch) { + ptrs[i] = base + i * n2; + } +} + +} // namespace cu + +void Cholesky::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("Cholesky::eval_gpu"); + auto& s = stream(); + auto& a = inputs[0]; + + // Copy the input to the output; the factorization runs in place. + copy_gpu( + a, + out, + a.flags().row_contiguous ? CopyType::Vector : CopyType::General, + s); + + if (a.size() == 0) { + return; + } + + auto& encoder = cu::get_command_encoder(s); + + auto handle = get_cusolver_handle(encoder.device()); + CHECK_CUSOLVER_ERROR(cusolverDnSetStream(handle, encoder.stream())); + + // cuSOLVER is column major. The input is symmetric and a column major + // lower triangle is a row major upper one, so pass the opposite of + // upper_, same as the CPU op. + cublasFillMode_t uplo = + upper_ ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER; + + // Only float32 reaches eval_gpu: float64 is rejected on GPU streams at + // array construction. + auto type = CUDA_R_32F; + auto* out_ptr = gpu_ptr(out); + + int64_t n = a.shape(-1); + int64_t batch = a.size() / (n * n); + + // Matching the CPU op, info is never read back: a non positive definite + // input gives an undefined factor rather than an error. + auto* info = static_cast( + allocate_workspace(encoder, batch * sizeof(int))); + + // TODO: For some consumer cards serialized loop might be faster. + if (batch > 1) { + array ptrs( + cu::malloc_async(batch * sizeof(void*), encoder), + {static_cast(batch)}, + uint64); + encoder.add_temporary(ptrs); + encoder.set_output_array(ptrs); + encoder.add_kernel_node( + &cu::fill_matrix_pointers, + cuda::ceil_div(batch, 256), + 256, + gpu_ptr(ptrs), + out_ptr, + n * n, + batch); + + encoder.set_input_array(ptrs); + encoder.set_output_array(out); + auto capture = encoder.capture_context(); + CHECK_CUSOLVER_ERROR(cusolverDnSpotrfBatched( + handle, + uplo, + n, + gpu_ptr(ptrs), + /* lda */ n, + info, + batch)); + } else { + int lwork = 0; + CHECK_CUSOLVER_ERROR(cusolverDnSpotrf_bufferSize( + handle, uplo, n, out_ptr, n, &lwork)); + auto* workspace = static_cast( + allocate_workspace(encoder, lwork * sizeof(float))); + + encoder.set_input_array(out); + encoder.set_output_array(out); + auto capture = encoder.capture_context(); + CHECK_CUSOLVER_ERROR(cusolverDnSpotrf( + handle, uplo, n, out_ptr, n, workspace, lwork, info)); + } + + encoder.set_output_array(out); + dispatch_bool(out.size() > INT32_MAX, [&](auto large) { + using IdxT = std::conditional_t; + auto [num_blocks, block_dims] = get_launch_args(out, large()); + encoder.add_kernel_node( + cu::zero_triangle, + num_blocks, + block_dims, + out_ptr, + static_cast(out.size()), + static_cast(n), + upper_); + }); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h index f8a234ee6..083d19e6f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h @@ -26,11 +26,10 @@ class CudaHandle { } ~CudaHandle() { - // Skip if there was an error to avoid throwing in the destructors - if (cudaPeekAtLastError() != cudaSuccess) { - return; + // Skip error check to avoid throwing in destructors. + if (handle_ != nullptr) { + Destroy(handle_); } - reset(); } CudaHandle(const CudaHandle&) = delete; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.cpp new file mode 100644 index 000000000..b4366d8d7 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.cpp @@ -0,0 +1,47 @@ +// Copyright © 2026 Apple Inc. +#include "mlx/backend/cuda/cusolver_utils.h" +#include "mlx/backend/gpu/device_info.h" + +#include + +namespace mlx::core { + +namespace { + +auto& cusolver_handles_cache() { + struct CusolverHandle { + ~CusolverHandle() { + if (handle) { + // Not checked: runs at thread exit where a throw would terminate. + cusolverDnDestroy(handle); + } + } + cusolverDnHandle_t handle{nullptr}; + }; + static thread_local std::vector cache(gpu::device_count()); + return cache; +} + +} // namespace + +cusolverDnHandle_t get_cusolver_handle(cu::Device& device) { + auto& storage = cusolver_handles_cache().at(device.cuda_device()); + if (!storage.handle) { + device.make_current(); + CHECK_CUSOLVER_ERROR(cusolverDnCreate(&storage.handle)); + } + return storage.handle; +} + +void init_cusolver_handles_cache() { + cusolver_handles_cache(); +} + +void check_cusolver_error(const char* name, cusolverStatus_t err) { + if (err != CUSOLVER_STATUS_SUCCESS) { + throw std::runtime_error( + fmt::format("{} failed with code: {}.", name, static_cast(err))); + } +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.h new file mode 100644 index 000000000..ad457a982 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cusolver_utils.h @@ -0,0 +1,18 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "mlx/backend/cuda/device.h" + +#include + +namespace mlx::core { + +void check_cusolver_error(const char* name, cusolverStatus_t err); + +#define CHECK_CUSOLVER_ERROR(cmd) check_cusolver_error(#cmd, (cmd)) + +void init_cusolver_handles_cache(); + +cusolverDnHandle_t get_cusolver_handle(cu::Device& device); + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp index 7092d4497..44f6d163c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp @@ -61,6 +61,22 @@ fs::path load_cudnn() { return cudnn_dir; } +fs::path load_cusolver() { + fs::path cusolver_dir; + if (const char* dir = cuda_bin_dir()) { + cusolver_dir = dir; + } else { + cusolver_dir = relative_to_current_binary("../nvidia/cusolver/bin"); + ::AddDllDirectory(cublas_dir().c_str()); + ::AddDllDirectory( + relative_to_current_binary("../nvidia/cusparse/bin").c_str()); + ::AddDllDirectory( + relative_to_current_binary("../nvidia/nvjitlink/bin").c_str()); + } + ::AddDllDirectory(cusolver_dir.c_str()); + return cusolver_dir; +} + // Called by system when failed to locate a lazy-loaded DLL. FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) { HMODULE mod = NULL; @@ -71,6 +87,9 @@ FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) { mod = ::LoadLibraryW((cudnn_dir / dll).c_str()); } else if (dll.starts_with("cublas")) { mod = ::LoadLibraryW((cublas_dir() / dll).c_str()); + } else if (dll.starts_with("cusolver")) { + static auto cusolver_dir = load_cusolver(); + mod = ::LoadLibraryW((cusolver_dir / dll).c_str()); } else if (dll.starts_with("nvrtc")) { static auto nvrtc_dir = load_nvrtc(); mod = ::LoadLibraryW((nvrtc_dir / dll).c_str()); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp index 472b3d99f..6e40b3403 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp @@ -213,7 +213,11 @@ CommandEncoder::CommandEncoder(Device& d) } CommandEncoder::~CommandEncoder() { - synchronize(); + try { + synchronize(); + } catch (...) { + // Synchronizing can fail when the CUDA runtime is shutting down. + } worker_->stop(); } @@ -612,13 +616,20 @@ CommandEncoder& get_command_encoder(Stream s) { } std::unordered_map& get_command_encoders() { - static thread_local std::unordered_map encoders; - return encoders; + if (is_main_thread()) { + static auto* encoders = new std::unordered_map(); + return *encoders; + } else { + static thread_local std::unordered_map encoders; + return 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/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh index e2ad855b9..5eea030c5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh @@ -269,6 +269,10 @@ void qmm_naive_kernel( // For gather, use index lookup for input batch slicing. uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord; + // The global scale is per gathered expert; a non-gathered call has one. + if (global_scale && rhs_indices) { + global_scale += b_batch; + } // Get batch slice. Tensor mA = mA_mkl(_,_,a_batch); // (M,K) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp index f8a71d71b..2862f3822 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp @@ -4,6 +4,7 @@ #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/cublas_utils.h" #include "mlx/backend/cuda/cudnn_utils.h" +#include "mlx/backend/cuda/cusolver_utils.h" #include "mlx/backend/cuda/event.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" @@ -20,6 +21,7 @@ void init() { mlx::core::cu::CudaEvent::init_pool(); init_cublas_handles_cache(); init_cudnn_handles_cache(); + init_cusolver_handles_cache(); init_cudnn_conv_cache(); init_cudnn_sdpa_cache(); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu index d3b6f97f5..bcea00544 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu @@ -232,8 +232,7 @@ AtomicEvent::AtomicEvent(Device& d) { cuda_free = cudaFree; coherent_ = false; } - buf_ = std::shared_ptr( - buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); }); + buf_ = std::shared_ptr(buf, [cuda_free](void* buf) { cuda_free(buf); }); if (coherent_) { *ptr() = 0; } else { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp index 3a3acdba0..e775d73ec 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp @@ -5,26 +5,43 @@ #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/event.h" +#include + namespace mlx::core { struct FenceImpl { uint32_t count; - Event event; + std::vector gpu_events; + Event cpu_event; - FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} + FenceImpl(uint32_t count, Stream s) : count(count), cpu_event(s) { + // Ensure that we use AtomicEvent, it is the only event that can order a CPU + // stream against the GPU. + cpu_event.cast().ensure_created(s, 2); + } }; Fence::Fence(Stream s) { fence_ = std::make_shared(0, s); - // Ensure that we use AtomicEvent. - cast().event.cast().ensure_created(s, 2); } -void Fence::wait(Stream s, const array&) { - cast().event.wait(); +void Fence::wait(Stream s, const array&, uint32_t value) { + auto& f = cast(); + if (value == 0) { + return; + } + if (!f.gpu_events.empty() && s.device == Device::gpu) { + f.gpu_events.at(value - 1).wait(s); + } else { + // AtomicEvent can not reliably notify a GPU stream, so a dependency that + // involves the CPU keeps the synchronous wait. + auto& event = f.cpu_event; + event.set_value(value); + event.wait(); + } } -void Fence::update(Stream s, const array& a, bool cross_device) { +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 @@ -37,8 +54,15 @@ void Fence::update(Stream s, const array& a, bool cross_device) { } } f.count++; - f.event.set_value(f.count); - f.event.signal(s); + if (s.device == Device::gpu) { + // Keep each recording so a consumer can wait for an earlier update. + auto& event = f.gpu_events.emplace_back(s); + event.set_value(1); + event.signal(s); + } + f.cpu_event.set_value(f.count); + f.cpu_event.signal(s); + return f.count; } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp index 94f260767..c02357b08 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp @@ -24,14 +24,27 @@ namespace mlx::core { throw std::runtime_error(#func " has no CUDA implementation."); \ } +bool fast::GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + return true; +} + NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) NO_GPU_MULTI(SVD) NO_GPU(Inverse) -NO_GPU(Cholesky) NO_GPU_MULTI(Eig) NO_GPU_MULTI(Eigh) +namespace fast { +NO_GPU_MULTI(GatedDeltaUpdate) +} + namespace distributed { NO_GPU_MULTI(Send) NO_GPU_MULTI(Recv) 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 64cfecfd5..6991baad1 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h @@ -135,6 +135,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu index 4ec6b95a7..63c847844 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu @@ -222,6 +222,7 @@ __global__ void gather_qmv_kernel( const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, const uint32_t* lhs_indices, const uint32_t* rhs_indices, @@ -244,8 +245,12 @@ __global__ void gather_qmv_kernel( x += block.group_index().y * k + m * k * x_idx; out += block.group_index().y * n + m * n * l; + // The global scale is per gathered expert. + if (global_scale) { + global_scale += w_idx; + } qmv_kernel_impl( - x, w, scales, biases, nullptr, out, row, w_idx, n, k); + x, w, scales, biases, global_scale, out, row, w_idx, n, k); } template < @@ -305,6 +310,7 @@ void gather_qmv( const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, const uint32_t* lhs_indices, const uint32_t* rhs_indices, @@ -321,7 +327,16 @@ void gather_qmv( uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m), uint32_t(l)}; dim3 block_dims{WARP_SIZE, rows_per_block}; void* args[] = { - &x, &w, &scales, &biases, &out, &lhs_indices, &rhs_indices, &n, &k}; + &x, + &w, + &scales, + &biases, + &global_scale, + &out, + &lhs_indices, + &rhs_indices, + &n, + &k}; dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) { auto* kernel = &gather_qmv_kernel< @@ -464,6 +479,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -490,6 +506,9 @@ void gather_qmv( if (biases) { encoder.set_input_array(*biases); } + if (global_scale) { + encoder.set_input_array(*global_scale); + } encoder.set_input_array(lhs_indices); encoder.set_input_array(rhs_indices); encoder.set_output_array(out); @@ -499,6 +518,7 @@ void gather_qmv( gpu_ptr(w), gpu_ptr(scales), biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, gpu_ptr(out), gpu_ptr(lhs_indices), gpu_ptr(rhs_indices), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp index 196bd9c05..5a9fba33d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp @@ -1,5 +1,6 @@ // Copyright © 2025 Apple Inc. +#include "mlx/backend/common/broadcasting.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/backend/cuda/quantized/qqmm_impl.h" @@ -244,6 +245,14 @@ void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { ensure_contiguous(w_pre, encoder, s), ensure_contiguous(inputs[base_size - 1], encoder, s)); + // The gather kernels take one global scale per expert. + if (global_scale_w) { + int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2); + array gs_e(Shape{E}, float32, nullptr, {}); + broadcast(*global_scale_w, gs_e); + global_scale_w = ensure_row_contiguous(gs_e, encoder, s); + } + // Quantize activation. array x = quantize_dequantize_input( x_pre, global_scale_x, bits_, group_size_, encoder, s); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp index 4d25f3c3e..5c3e25580 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/device.h" @@ -159,9 +159,13 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { array x = ensure_row_contiguous(inputs[0], encoder, s); const array& w = inputs[1]; const array& scales = inputs[2]; + // Affine gets biases at index 3, nvfp4 an optional global scale. std::optional biases; - if (inputs.size() == 6) { + std::optional global_scale; + if (mode_ == QuantizationMode::Affine) { biases = inputs[3]; + } else if (inputs.size() == 6) { + global_scale = ensure_row_contiguous(inputs[3], encoder, s); } array lhs_indices = ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s); @@ -186,7 +190,10 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { mode_, encoder.device()); }; - bool can_use_qmm_sm80 = supports(supports_qmm_sm80); + // qmm_sm80 does not apply global scales yet; route such calls to the + // naive kernel until it does. + bool can_use_qmm_sm80 = + !global_scale.has_value() && supports(supports_qmm_sm80); bool can_use_qmm_naive = supports(supports_qmm_naive); bool can_use_qmv = supports(supports_qmv); @@ -212,7 +219,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, - std::nullopt, + global_scale, lhs_indices, rhs_indices, out, @@ -229,6 +236,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + global_scale, lhs_indices, rhs_indices, out, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu index a7dd2eea2..c2ba87527 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu @@ -458,6 +458,11 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); + if (out.size() == 0) { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + return; + } + if (in.flags().contiguous && in.strides()[axis_] != 0) { if (in.is_donatable() && in.itemsize() == out.itemsize()) { out.copy_shared_buffer(in); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu index d9c77387c..5d73681a6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu @@ -1084,6 +1084,10 @@ void gpu_sort( int axis, bool argsort) { auto& encoder = cu::get_command_encoder(s); + if (out.size() == 0) { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + return; + } gpu_merge_sort(s, in, out, axis, argsort); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/wddm.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/wddm.cpp new file mode 100644 index 000000000..71027f77c --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/wddm.cpp @@ -0,0 +1,107 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/wddm.h" +#include "mlx/backend/cuda/utils.h" +#include "mlx/backend/gpu/device_info.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace mlx::core::cu { + +namespace { + +struct DeviceMemoryBudget { + Microsoft::WRL::ComPtr adapter; + UINT node{0}; + std::chrono::steady_clock::time_point next_query; + size_t memory_limit{std::numeric_limits::max()}; +}; + +class WddmMemoryBudget { + public: + explicit WddmMemoryBudget(size_t device_count) : budgets_(device_count) { + Microsoft::WRL::ComPtr factory; + if (FAILED(CreateDXGIFactory1(IID_PPV_ARGS(&factory)))) { + return; + } + + for (int device = 0; device < device_count; ++device) { + cudaDeviceProp properties{}; + CHECK_CUDA_ERROR(cudaGetDeviceProperties(&properties, device)); + if (properties.integrated || properties.tccDriver || + properties.luidDeviceNodeMask == 0) { + continue; + } + + LUID luid{}; + std::memcpy(&luid, properties.luid, sizeof(luid)); + if (luid.HighPart == 0 && luid.LowPart == 0) { + continue; + } + + auto& budget = budgets_[device]; + if (FAILED(factory->EnumAdapterByLuid( + luid, IID_PPV_ARGS(&budget.adapter)))) { + continue; + } + + auto mask = properties.luidDeviceNodeMask; + while ((mask & 1) == 0) { + ++budget.node; + mask >>= 1; + } + } + } + + size_t get_memory_limit(int device, cudaMemPool_t pool) { + size_t max_size = std::numeric_limits::max(); + auto& budget = budgets_.at(device); + if (!budget.adapter || !pool) { + return max_size; + } + + auto now = std::chrono::steady_clock::now(); + std::lock_guard lock(mutex_); + if (now < budget.next_query) { + return budget.memory_limit; + } + budget.next_query = now + std::chrono::milliseconds(20); + + DXGI_QUERY_VIDEO_MEMORY_INFO info{}; + uint64_t pool_reserved = 0; + if (FAILED(budget.adapter->QueryVideoMemoryInfo( + budget.node, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &info))) { + return max_size; + } + if (info.Budget == 0) { + return max_size; + } + + CHECK_CUDA_ERROR(cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &pool_reserved)); + budget.memory_limit = + memory_limit_from_budget(info.Budget, info.CurrentUsage, pool_reserved); + return budget.memory_limit; + } + + private: + std::vector budgets_; + std::mutex mutex_; +}; + +} // namespace + +size_t get_wddm_memory_limit(int device, cudaMemPool_t pool) { + static WddmMemoryBudget budget(gpu::device_count()); + return budget.get_memory_limit(device, pool); +} + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/wddm.h b/Source/Cmlx/mlx/mlx/backend/cuda/wddm.h new file mode 100644 index 000000000..c76224b7b --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/wddm.h @@ -0,0 +1,26 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include + +namespace mlx::core::cu { + +inline size_t memory_limit_from_budget( + uint64_t budget, + uint64_t usage, + uint64_t pool_reserved) { + auto non_pool_usage = usage > pool_reserved ? usage - pool_reserved : 0; + auto margin = budget / 20; + if ((non_pool_usage >= budget) || (margin >= budget - non_pool_usage)) { + return 0; + } + uint64_t size_limit = std::numeric_limits::max(); + return std::min(budget - non_pool_usage - margin, size_limit); +} + +size_t get_wddm_memory_limit(int device, cudaMemPool_t pool); + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp index de38a06e5..6266d6099 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp @@ -57,14 +57,17 @@ void Worker::commit(cudaStream_t stream) { } void Worker::thread_fn() { - while (!stop_) { - uint64_t current_batch = 0; + uint64_t current_batch = 0; + while (true) { Tasks tasks; { std::unique_lock lk(mtx_); - cond_.wait(lk, [this, ¤t_batch] { + cond_.wait(lk, [this, current_batch] { return this->signaled_batch_ > current_batch || this->stop_; }); + if (stop_) { + return; + } current_batch = signaled_batch_; auto end = worker_tasks_.upper_bound(current_batch); for (auto it = worker_tasks_.begin(); it != end; ++it) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt index ea4a995ad..bc5c2fcf7 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt @@ -44,6 +44,7 @@ if(MLX_METAL_JIT) make_jit_source(binary_two) make_jit_source(fft kernels/fft/radix.h kernels/fft/readwrite.h) make_jit_source(logsumexp) + make_jit_source(cross_entropy) make_jit_source(ternary) make_jit_source(softmax) make_jit_source(scan) @@ -102,6 +103,7 @@ if(MLX_METAL_JIT) kernels/fp4.h) make_jit_source(steel/attn/kernels/steel_attention_nax) + make_jit_source(gated_delta_update_nax) else() message( @@ -134,8 +136,10 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cpp ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gated_delta_update.cpp ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp diff --git a/Source/Cmlx/mlx/mlx/backend/metal/cross_entropy.cpp b/Source/Cmlx/mlx/mlx/backend/metal/cross_entropy.cpp new file mode 100644 index 000000000..cc6f0e171 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/cross_entropy.cpp @@ -0,0 +1,125 @@ +// Copyright © 2026 Apple Inc. +#include + +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/kernels/defines.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" + +namespace mlx::core::fast { + +namespace { + +// Threadgroup shrinks with the axis so a small vocabulary needs one iteration. +inline std::pair cross_entropy_dims( + MTL::ComputePipelineState* kernel, + int axis_size, + int n_rows) { + constexpr size_t simd_size = 32; + size_t tg_max = + (kernel->maxTotalThreadsPerThreadgroup() / simd_size) * simd_size; + size_t needed = + (axis_size + CROSS_ENTROPY_N_READS - 1) / CROSS_ENTROPY_N_READS; + size_t threadgroup_size = std::clamp( + ((needed + simd_size - 1) / simd_size) * simd_size, simd_size, tg_max); + return { + MTL::Size(n_rows * threadgroup_size, 1, 1), + MTL::Size(threadgroup_size, 1, 1)}; +} + +array ensure_row_contiguous( + const array& x, + metal::CommandEncoder& encoder, + const Stream& s) { + if (x.flags().row_contiguous) { + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + encoder.add_temporary(x_copy); + return x_copy; +} + +} // namespace + +bool CrossEntropy::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 2); + auto& in_pre = inputs[0]; + auto& out = outputs[0]; + if (in_pre.size() == 0) { + throw std::invalid_argument("[cross_entropy] Received empty array."); + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + auto in = ensure_row_contiguous(in_pre, compute_encoder, s); + auto targets = ensure_row_contiguous(inputs[1], compute_encoder, s); + out.set_data(allocator::malloc(out.nbytes())); + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + std::string kernel_name = "cross_entropy_" + type_to_name(in); + auto kernel = get_cross_entropy_kernel(d, kernel_name, in); + auto [grid_dims, group_dims] = cross_entropy_dims(kernel, axis_size, n_rows); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_input_array(targets, 1); + compute_encoder.set_output_array(out, 2); + compute_encoder.set_bytes(axis_size, 3); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 4); + auto& in_pre = inputs[0]; + auto& out = outputs[0]; + if (in_pre.size() == 0) { + throw std::invalid_argument("[cross_entropy] Received empty array."); + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + bool donate_in = in_pre.is_donatable() || !in_pre.flags().row_contiguous; + auto in = ensure_row_contiguous(in_pre, compute_encoder, s); + auto targets = ensure_row_contiguous(inputs[1], compute_encoder, s); + auto loss = ensure_row_contiguous(inputs[2], compute_encoder, s); + auto cotan = ensure_row_contiguous(inputs[3], compute_encoder, s); + if (donate_in) { + out.copy_shared_buffer(in); + } else { + out.set_data(allocator::malloc(out.nbytes())); + } + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + std::string kernel_name = "cross_entropy_vjp_" + type_to_name(in); + auto kernel = get_cross_entropy_kernel(d, kernel_name, in); + auto [grid_dims, group_dims] = cross_entropy_dims(kernel, axis_size, n_rows); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_input_array(targets, 1); + compute_encoder.set_input_array(loss, 2); + compute_encoder.set_input_array(cotan, 3); + compute_encoder.set_output_array(out, 4); + compute_encoder.set_bytes(axis_size, 5); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +} // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp index a9a0318b2..3e411445f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp @@ -926,8 +926,13 @@ CommandEncoder& get_command_encoder(Stream s) { } std::unordered_map& get_command_encoders() { - static thread_local std::unordered_map encoders; - return encoders; + if (is_main_thread()) { + static auto* encoders = new std::unordered_map(); + return *encoders; + } else { + static thread_local std::unordered_map encoders; + return encoders; + } } std::unordered_map& get_global_command_encoders() { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp index 70dd0e33b..6594f5507 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp @@ -44,18 +44,20 @@ Fence::Fence(Stream stream) { fence_ = std::make_shared(stream); } -void Fence::wait(Stream stream, const array& x) { +void Fence::wait(Stream stream, const array& x, uint32_t value) { auto& f = *static_cast(fence_.get()); if (!f.use_fast) { - f.event->wait(stream); + auto& event = *f.event; + event.set_value(value); + event.wait(stream); return; } if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable { + scheduler::enqueue(stream, [fence_ = fence_, value]() mutable { auto& f = *static_cast(fence_.get()); - while (f.cpu_value()[0] < count) { + while (f.cpu_value()[0] < value) { } }); return; @@ -74,21 +76,21 @@ void Fence::wait(Stream stream, const array& x) { auto buf = static_cast(f.fence); compute_encoder.set_buffer(buf, 0); - compute_encoder.set_bytes(f.count, 1); + compute_encoder.set_bytes(value, 1); compute_encoder.dispatch_threads(kernel_dims, kernel_dims); compute_encoder.get_command_buffer()->addCompletedHandler( [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); } -void Fence::update(Stream stream, const array& x, bool cross_device) { +uint32_t Fence::update(Stream stream, const array& x, bool cross_device) { auto& f = *static_cast(fence_.get()); f.count++; if (!f.use_fast) { f.event->set_value(f.count); f.event->signal(stream); - return; + return f.count; } if (stream.device == Device::cpu) { @@ -96,7 +98,7 @@ void Fence::update(Stream stream, const array& x, bool cross_device) { auto& f = *static_cast(fence_.get()); f.cpu_value()[0] = count; }); - return; + return f.count; } auto& d = metal::device(stream.device); @@ -130,6 +132,7 @@ void Fence::update(Stream stream, const array& x, bool cross_device) { compute_encoder.get_command_buffer()->addCompletedHandler( [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); + return f.count; } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp b/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp new file mode 100644 index 000000000..45f41677c --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/gated_delta_update.cpp @@ -0,0 +1,167 @@ +// Copyright © 2024 Apple Inc. +#include + +#include "mlx/backend/common/compiled.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/kernels/defines.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" +#include "mlx/utils.h" + +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; + } + + if (has_mask) { + return true; + } + + if (Dk != 128 || Dv != 128) { + return true; + } + + 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; + } + + return false; +} + +inline 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; + } +} + +void GatedDeltaUpdate::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + 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& out = outputs[0]; + auto& hf = outputs[1]; + + 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); + + 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)); + + auto& compute_encoder = metal::get_command_encoder(s); + + out.set_data(allocator::malloc(out.nbytes())); + hf.set_data(allocator::malloc(hf.nbytes())); + + 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 = {}; + + auto delta_kernel = + get_gated_delta_nax_kernel(d, base_name, hash_name, func_consts); + compute_encoder.set_compute_pipeline_state(delta_kernel); + 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 = {}; + + auto delta_kernel = + get_gated_delta_kernel(d, base_name, hash_name, func_consts); + compute_encoder.set_compute_pipeline_state(delta_kernel); + 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 = {}; + + auto delta_kernel = + get_gated_delta_kernel(d, base_name, hash_name, func_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + break; + } + default: { + throw std::runtime_error( + "NYI: Only sequential (C=1) and chunk size 8,16 are supported"); + } + } + 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); +} + +} // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h index 4fb1be111..b4f72e660 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h @@ -18,6 +18,7 @@ const char* unary(); const char* binary(); const char* binary_two(); const char* copy(); +const char* cross_entropy(); const char* fft(); const char* gather_axis(); const char* gather_front(); @@ -59,4 +60,7 @@ const char* fp_quantized_nax(); const char* steel_attention_nax(); +const char* gated_delta_update(); +const char* gated_delta_update_nax(); + } // 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 8dfe30a15..d2ce55bbf 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp @@ -39,6 +39,9 @@ const char* fp_quantized_nax() { const char* steel_attention_nax() { return ""; } +const char* gated_delta_update_nax() { + return ""; +} } // namespace metal #endif // MLX_METAL_NO_NAX @@ -360,6 +363,26 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in) { + // Both kernels live in one library, so the name is not derived from + // kernel_name. + std::string lib_name = "cross_entropy_" + type_to_name(in); + auto lib = d.get_library(lib_name, [&] { + auto t_str = get_type_string(in.dtype()); + std::string kernel_source = metal::utils(); + kernel_source += metal::cross_entropy(); + kernel_source += get_template_definition( + "cross_entropy_" + type_to_name(in), "cross_entropy", t_str); + kernel_source += get_template_definition( + "cross_entropy_vjp_" + type_to_name(in), "cross_entropy_vjp", t_str); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, @@ -1054,28 +1077,44 @@ MTL::ComputePipelineState* get_gather_qmm_kernel( int bk, int wm, int wn, - bool transpose) { + bool transpose, + bool has_global_scale) { const auto& lib_name = kernel_name; auto lib = d.get_library(lib_name, [&]() { std::string kernel_source; concatenate( kernel_source, metal::utils(), metal::quantized_utils(), metal::gemm()); bool is_affine = mode == "affine"; + // Only the fp kernel takes a global scale. + auto template_def = is_affine ? get_template_definition( + lib_name, + "affine_gather_qmm_rhs", + get_type_string(x.dtype()), + group_size, + bits, + bm, + bn, + bk, + wm, + wn, + transpose) + : get_template_definition( + lib_name, + "fp_gather_qmm_rhs", + get_type_string(x.dtype()), + group_size, + bits, + bm, + bn, + bk, + wm, + wn, + transpose, + has_global_scale); concatenate( kernel_source, is_affine ? metal::quantized() : metal::fp_quantized(), - get_template_definition( - lib_name, - (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs"), - get_type_string(x.dtype()), - group_size, - bits, - bm, - bn, - bk, - wm, - wn, - transpose)); + template_def); return kernel_source; }); return d.get_kernel(kernel_name, lib, hash_name, func_consts); @@ -1255,7 +1294,8 @@ MTL::ComputePipelineState* get_gather_qmm_nax_kernel( int bk, int wm, int wn, - bool transpose) { + bool transpose, + bool has_global_scale) { const auto& lib_name = kernel_name; auto lib = d.get_library(lib_name, [&]() { std::string kernel_source; @@ -1265,21 +1305,36 @@ MTL::ComputePipelineState* get_gather_qmm_nax_kernel( metal::gemm_nax(), metal::quantized_utils()); bool is_affine = mode == "affine"; + // Only the fp kernel takes a global scale. + auto template_def = is_affine ? get_template_definition( + lib_name, + "affine_gather_qmm_rhs_nax", + get_type_string(x.dtype()), + group_size, + bits, + bm, + bn, + bk, + wm, + wn, + transpose) + : get_template_definition( + lib_name, + "fp_gather_qmm_rhs_nax", + get_type_string(x.dtype()), + group_size, + bits, + bm, + bn, + bk, + wm, + wn, + transpose, + has_global_scale); concatenate( kernel_source, is_affine ? metal::quantized_nax() : metal::fp_quantized_nax(), - get_template_definition( - lib_name, - (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs_nax"), - get_type_string(x.dtype()), - group_size, - bits, - bm, - bn, - bk, - wm, - wn, - transpose)); + template_def); return kernel_source; }); return d.get_kernel(kernel_name, lib, hash_name, func_consts); @@ -1294,6 +1349,7 @@ MTL::ComputePipelineState* get_steel_attention_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m) { @@ -1311,6 +1367,7 @@ MTL::ComputePipelineState* get_steel_attention_kernel( bq, bk, bd, + bv, wm, wn, get_type_string(m.dtype()))); @@ -1328,6 +1385,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m, @@ -1346,6 +1404,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( bq, bk, bd, + bv, wm, wn, get_type_string(m.dtype()))); @@ -1354,4 +1413,26 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_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 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()); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h index 18a56ebf1..82c3d3d16 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h @@ -64,6 +64,11 @@ MTL::ComputePipelineState* get_logsumexp_kernel( const std::string& kernel_name, const array& out); +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in); + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, @@ -321,7 +326,8 @@ MTL::ComputePipelineState* get_gather_qmm_kernel( int bk, int wm, int wn, - bool transpose); + bool transpose, + bool has_global_scale); MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel( metal::Device& d, @@ -400,7 +406,8 @@ MTL::ComputePipelineState* get_gather_qmm_nax_kernel( int bk, int wm, int wn, - bool transpose); + bool transpose, + bool has_global_scale); MTL::ComputePipelineState* get_steel_attention_kernel( metal::Device& d, @@ -411,6 +418,7 @@ MTL::ComputePipelineState* get_steel_attention_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m); @@ -424,11 +432,24 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m, bool split_d); +MTL::ComputePipelineState* get_gated_delta_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); + // 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 ecabbda55..d5513dd03 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt @@ -50,11 +50,13 @@ endfunction(build_kernel) build_kernel(arg_reduce) build_kernel(conv steel/conv/params.h) build_kernel(dot) +build_kernel(gather_mm_offsets) 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(gated_delta_update gated_delta_update.h) if(MLX_METAL_VERSION GREATER_EQUAL 320) build_kernel(fence) endif() @@ -142,6 +144,7 @@ if(NOT MLX_METAL_JIT) build_kernel(scan scan.h) build_kernel(softmax softmax.h) build_kernel(logsumexp logsumexp.h) + build_kernel(cross_entropy cross_entropy.h) build_kernel(searchsorted searchsorted.h sort.h) build_kernel(sort sort.h) build_kernel(ternary ternary.h ternary_ops.h) @@ -169,6 +172,8 @@ if(NOT MLX_METAL_JIT) ${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}) 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/arg_reduce.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/arg_reduce.metal index 6ce20555f..3ac51f1c9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/arg_reduce.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/arg_reduce.metal @@ -1,11 +1,17 @@ // Copyright © 2023 Apple Inc. +#include #include #include "mlx/backend/metal/kernels/utils.h" using namespace metal; +template , bool> = true> +inline bool isnan(T) { + return false; +} + template struct IndexValPair { uint32_t index; @@ -17,8 +23,11 @@ struct ArgMin { static constexpr constant U init = Limits::max; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val > current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val > current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -29,7 +38,7 @@ struct ArgMin { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -43,8 +52,11 @@ struct ArgMax { static constexpr constant U init = Limits::min; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val < current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val < current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -55,7 +67,7 @@ struct ArgMax { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h index 37650c7d9..968087f8e 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h @@ -261,7 +261,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.h new file mode 100644 index 000000000..bef0bce5f --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.h @@ -0,0 +1,133 @@ +// Copyright © 2026 Apple Inc. + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy( + const device T* in, + const device int32_t* targets, + device float* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + in += gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(in[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = + (offset + i < axis_size) ? AccT(in[offset + i]) : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + uint n_simdgroups = ceildiv(lsize, uint(SIMD_SIZE)); + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = (simd_lane_id < n_simdgroups) ? local_max[simd_lane_id] + : Limits::finite_min; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = + (simd_lane_id < n_simdgroups) ? local_normalizer[simd_lane_id] : AccT(0); + normalizer = simd_sum(normalizer); + + if (lid == 0) { + // Subtract the two logits first; they are the same magnitude, so less is + // lost than in logsumexp - x_t. + AccT gap = maxval - x_t; + out[gid] = isinf(maxval) ? float(gap) : float(log(normalizer) + gap); + } +} + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy_vjp( + const device T* in, + const device int32_t* targets, + const device float* loss, + const device float* cotan, + device T* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + size_t row_offset = gid * size_t(axis_size); + in += row_offset; + out += row_offset; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + AccT g = AccT(cotan[gid]); + AccT l = AccT(loss[gid]); + + // out aliases in when the input is donated, so every thread must read the + // target column before any thread writes. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } else { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + if (col < axis_size) { + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } + } + } +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.metal new file mode 100644 index 000000000..f35fbe060 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/cross_entropy.metal @@ -0,0 +1,18 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +using namespace metal; + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/cross_entropy.h" + +#define instantiate_cross_entropy(name, itype) \ + instantiate_kernel("cross_entropy_" #name, cross_entropy, itype) \ + instantiate_kernel("cross_entropy_vjp_" #name, cross_entropy_vjp, itype) + +instantiate_cross_entropy(float32, float) +instantiate_cross_entropy(float16, half) +instantiate_cross_entropy(bfloat16, bfloat16_t) // clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/defines.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/defines.h index c369adb7e..0e6590616 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/defines.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/defines.h @@ -14,6 +14,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h index 42c5ec128..00efdeb79 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h @@ -30,11 +30,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { 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 7061771d5..b137a6d47 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h @@ -136,13 +136,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { } template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -154,7 +153,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -164,9 +164,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -180,6 +185,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -187,7 +195,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -202,14 +211,19 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales( scales_ + bi * src_ld / group_size + - (bj * pack_factor) / group_size) {} + (bj * pack_factor) / group_size) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -221,21 +235,15 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -788,12 +796,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -831,7 +841,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -850,12 +861,15 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -865,7 +879,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -875,9 +889,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -887,7 +910,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -898,6 +921,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -913,12 +945,14 @@ template < typename T, int group_size, int bits, + bool has_global_scale = false, int BM = 32, int BK = 32, int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -956,7 +990,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; auto wl = (const device uint8_t*)w; @@ -971,16 +1006,23 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -989,9 +1031,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1009,7 +1061,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1018,9 +1074,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1036,8 +1102,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } @@ -1078,11 +1144,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device uint8_t*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -1125,6 +1192,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -1444,12 +1515,14 @@ template < 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> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -1489,8 +1562,22 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1546,8 +1633,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1576,10 +1663,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1635,10 +1723,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1694,10 +1783,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1733,12 +1823,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> [[kernel]] void fp_gather_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1768,10 +1860,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1788,8 +1881,22 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1829,9 +1936,10 @@ 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, x, y, Xs, @@ -1857,6 +1965,7 @@ template < [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1887,10 +1996,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1907,8 +2017,21 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1920,16 +2043,19 @@ template < int BK, int WM, int WN, - bool transpose> + bool transpose, + bool has_global_scale = false> [[kernel]] void fp_gather_qmm_rhs( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1960,7 +2086,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -1973,13 +2100,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1995,112 +2127,61 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } 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 d8d462288..fe508a6aa 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal @@ -47,6 +47,17 @@ bits, \ aligned) +#define instantiate_quantized_aligned_hgs(mode, name, type, aligned, group_size, bits) \ + instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + aligned, \ + true) + #define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched, \ @@ -57,6 +68,17 @@ aligned, \ batched) +#define instantiate_quantized_aligned_batched_hgs(mode, name, type, aligned, batched, group_size, bits) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + aligned, \ + 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, \ @@ -123,7 +145,20 @@ bk, \ wm, \ wn, \ - transpose) + transpose) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn "_hgs", \ + func, \ + type, \ + group_size, \ + bits, \ + bm, \ + bn, \ + bk, \ + wm, \ + wn, \ + transpose, \ + true) #define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \ instantiate_quantized_batched(mode, name, type, 1, group_size, bits) \ @@ -142,13 +177,17 @@ instantiate_quantized(mode, gather_qmm_n, type, group_size, bits) #define instantiate_quantized_all_aligned(type, mode, group_size, bits) \ - instantiate_quantized_aligned(mode, gather_qmm_t, type, true, group_size, bits) \ - instantiate_quantized_aligned(mode, gather_qmm_t, type, false, group_size, bits) \ + instantiate_quantized_aligned_hgs(mode, gather_qmm_t, type, true, group_size, bits) \ + instantiate_quantized_aligned_hgs(mode, gather_qmm_t, type, false, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 1, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 0, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 1, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 0, group_size, bits) +#define instantiate_quantized_qmm_hgs(type) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t, type, true, 0, 16, 4) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t, type, false, 0, 16, 4) + #define instantiate_quantized_all_quad(type, mode, group_size, bits) \ instantiate_quantized_quad(mode, qmv_quad, type, 64, 1, group_size, bits) \ instantiate_quantized_quad(mode, qmv_quad, type, 64, 0, group_size, bits) \ @@ -221,4 +260,5 @@ instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) instantiate_quantized_types(float16_t) +instantiate_quantized_qmm_hgs(bfloat16_t) // clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h index 946bce786..f843c9996 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -60,13 +60,12 @@ struct Dequantize { }; template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -78,7 +77,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -106,6 +106,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -113,7 +116,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -126,7 +130,11 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id) {} + scales(scales_ + bi * src_ld / group_size + group_id) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -135,7 +143,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -149,14 +158,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -165,7 +167,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -191,6 +194,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -200,6 +204,7 @@ template < METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup Wtype* Ws, @@ -229,7 +234,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -245,7 +251,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -335,6 +341,7 @@ template < typename T, const int group_size, const int bits, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -344,6 +351,7 @@ template < METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Ws, @@ -373,7 +381,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -391,7 +400,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -486,11 +495,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device S*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -533,6 +543,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -547,10 +561,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -589,8 +605,19 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -648,8 +675,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -662,10 +689,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -694,10 +723,11 @@ template < threadgroup Wtype Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -714,8 +744,19 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -727,10 +768,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_n_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -761,10 +804,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -781,8 +825,8 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -795,16 +839,19 @@ template < int WM, int WN, bool transpose, + bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_rhs_nax( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -821,7 +868,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -833,13 +881,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -865,10 +918,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -876,149 +929,123 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; - } - } + NAXTile Dtile; + Dtile.clear(); - xn += BK; - loader_w.next(); + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); } - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + xn += BK; + loader_w.next(); + } + + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal index 771b2a963..b6cc78130 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" @@ -24,7 +24,14 @@ type, \ group_size, \ bits, \ - aligned, bm, bk, bn, wm, wn) + aligned, bm, bk, bn, wm, wn) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + aligned, bm, bk, bn, wm, wn, true) #define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \ instantiate_kernel( \ @@ -36,6 +43,16 @@ aligned, \ batched, bm, bk, bn, wm, wn) +#define instantiate_quantized_aligned_batched_hgs(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + aligned, \ + batched, bm, bk, bn, wm, wn, true) + #define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \ @@ -48,7 +65,20 @@ bk, \ wm, \ wn, \ - transpose) + transpose) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn "_hgs", \ + func, \ + type, \ + group_size, \ + bits, \ + bm, \ + bn, \ + bk, \ + wm, \ + wn, \ + transpose, \ + true) #define instantiate_quantized_all_aligned(type, mode, group_size, bits) \ @@ -63,6 +93,12 @@ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 1, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 0, group_size, bits) +#define instantiate_quantized_qmm_hgs(type) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, 16, 4) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, 16, 4) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 0, 16, 4) \ + instantiate_quantized_aligned_batched_hgs(nvfp4, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 0, 16, 4) + #define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \ @@ -82,4 +118,5 @@ instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) instantiate_quantized_types(float16_t) +instantiate_quantized_qmm_hgs(bfloat16_t) // clang-format on 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 new file mode 100644 index 000000000..9553356c7 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.h @@ -0,0 +1,349 @@ +#pragma once + +#include +#include "mlx/backend/metal/kernels/utils.h" + +#define AT(TILE, IDX) TILE.thread_elements()[IDX] +#define SUB(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) - AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) - AT(TILE2, 1); \ + } +#define ADD(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) + AT(TILE2, 1); \ + } +#define FMA(TILE0, S, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = S * AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = S * AT(TILE1, 1) + AT(TILE2, 1); \ + } + +#define SCALE(TILE0, S) \ + { \ + AT(TILE0, 0) *= S; \ + AT(TILE0, 1) *= S; \ + } +#define SCALE2(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= S0; \ + AT(TILE0, 1) *= S1; \ + } +#define SCALE_TRI(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn > fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 > fm ? 0.f : S1; \ + } +#define SCALE_TRIEQ(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn >= fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ + } + +// lambdas are not supported in metal 14 so porting to macros. + +// non transposed +#define LOAD_M(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fm < valid_rows) ? ((SRC)[fm * (LD) + fn]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fm < valid_rows) ? ((SRC)[fm * (LD) + fn + 1]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fm * (LD) + fn]); \ + AT(M, 1) = static_cast((SRC)[fm * (LD) + fn + 1]); \ + } + +// transposed load: sequence is the column -> mask fn / fn+1 +#define LOAD_MT(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fn < valid_rows) ? ((SRC)[fn * (LD) + fm]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fn + 1 < valid_rows) ? ((SRC)[(fn + 1) * (LD) + fm]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fn * (LD) + fm]); \ + AT(M, 1) = static_cast((SRC)[(fn + 1) * (LD) + fm]); \ + } + +#define PROCESS_CHUNK_SG(B, S_tile, VALID) \ + { \ + const short valid_rows = (VALID); \ + \ + 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)) \ + : 0.0f; \ + \ + float gamma_val = simd_prefix_inclusive_sum(g_val); \ + \ + if (thread_index_in_simdgroup < C) { \ + gamma[thread_index_in_simdgroup] = gamma_val; \ + } \ + simdgroup_barrier(mem_flags::mem_threadgroup); \ + \ + float gamma_fm = metal::fast::exp(gamma[fm]); \ + float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); \ + float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); \ + float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); \ + float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); \ + float gamma_C = metal::fast::exp(gamma[C - 1]); \ + \ + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; \ + \ + KKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + LOAD_MT(KT_tile, k_ + kk, Dk * Hk, B) \ + simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); \ + } \ + \ + KKtK_tile = KKt_tile; \ + SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) \ + \ + simdgroup_float8x8 Tinv, P; \ + AT(P, 0) = AT(KKtK_tile, 0); \ + AT(P, 1) = AT(KKtK_tile, 1); \ + SUB(Tinv, I_tile, KKtK_tile) \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int step = 1; (1 << step) < C; step++) { \ + simdgroup_multiply(P, P, P); \ + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); \ + } \ + \ + WS_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + SCALE(K_tile, beta_fm) \ + simdgroup_multiply(W_tile, Tinv, K_tile); \ + SCALE(W_tile, gamma_fm) \ + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); \ + } \ + \ + SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) \ + \ + LOAD_M(V_tile, v_ + dv_idx, Dv * Hv, B) \ + SCALE(V_tile, beta_fm) \ + simdgroup_multiply(U_tile, Tinv, V_tile); \ + SUB(delta_tile, U_tile, WS_tile) \ + \ + tmp_tile = make_filled_simdgroup_matrix(0.f); \ + QKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(Q_tile, q_ + kk, Hk * Dk, B) \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); \ + SCALE(Q_tile, gamma_fm) \ + simdgroup_multiply_accumulate( \ + tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); \ + } \ + \ + SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) \ + \ + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); \ + \ + if (fm < valid_rows) { \ + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); \ + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); \ + } \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) \ + simdgroup_multiply(KD_tile, K_tile, delta_tile); \ + FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) \ + } \ + } + +template +[[kernel]] void gated_delta_fused_chunk( + 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)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + 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); + + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + + ((thread_index_in_simdgroup / 2) % 4); // row coordinate of the held tile + const short fn = (qid & 2) * 2 + + (thread_index_in_simdgroup % 2) * 2; // column coordinate of the held tile + + auto dv_idx = thread_position_in_grid.y * 8; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + +#define OUTPUT(T) \ + if (true) { \ + simdgroup_store(T, y); \ + return; \ + } + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + simdgroup_float8x8 S_tile[Dk / 8]; + + // simdgroup matrices + simdgroup_float8x8 V_tile, K_tile, KT_tile, Q_tile; + simdgroup_float8x8 W_tile, U_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + // tiles for WY form computation + simdgroup_float8x8 KKtK_tile, KKt_tile; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + simdgroup_float8x8 I_tile = make_filled_simdgroup_matrix(0.f); + AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; + AT(I_tile, 1) = (fm == fn + 1) ? 1.0f : 0.0f; + + // load initial state into registers + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); + } + + int t = 0; + MLX_MTL_PRAGMA_UNROLL + for (; t + C <= T; t += C) { + PROCESS_CHUNK_SG(false, S_tile, C); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + PROCESS_CHUNK_SG(true, S_tile, short(T - t)); + } + + MLX_MTL_PRAGMA_UNROLL + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); + } +} + +/* + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + */ +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* 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)]], + 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; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } +} \ No newline at end of file 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 new file mode 100644 index 000000000..21238df43 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update.metal @@ -0,0 +1,42 @@ +#include "mlx/backend/metal/kernels/gated_delta_update.h" +#include "mlx/backend/metal/kernels/utils.h" + +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_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) + +#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) + +#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 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 new file mode 100644 index 000000000..f4ae51826 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.h @@ -0,0 +1,527 @@ +#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; + +// 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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + 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; + + // 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< + decltype(ct_a), + decltype(ct_b), + CType>(); + + // 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 +[[kernel]] void gated_delta_fused_nax( + 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)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + 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); + + 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)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile 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} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + 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); + } + }; + + 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)) + : 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; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + tmp_tile.clear(); + QKt_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) + + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + 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 new file mode 100644 index 000000000..35790fa49 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gated_delta_update_nax.metal @@ -0,0 +1,31 @@ +#include "mlx/backend/metal/kernels/gated_delta_update_nax.h" +#include "mlx/backend/metal/kernels/utils.h" + +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_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16) + +#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 diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gather_mm_offsets.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gather_mm_offsets.metal new file mode 100644 index 000000000..40b933d51 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gather_mm_offsets.metal @@ -0,0 +1,23 @@ +// Copyright © 2026 Apple Inc. + +#include + +using namespace metal; + +[[kernel]] void gather_mm_offsets( + const device uint32_t* indices [[buffer(0)]], + device int32_t* offsets [[buffer(1)]], + const constant int& M [[buffer(2)]], + uint g [[thread_position_in_grid]]) { + int lo = 0; + int hi = M; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (indices[mid] < g) { + lo = mid + 1; + } else { + hi = mid; + } + } + offsets[g] = lo; +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h index 42855e14d..be1ebff59 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h @@ -596,6 +596,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -651,14 +657,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1897,6 +1896,7 @@ template < 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> @@ -2276,6 +2276,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -2437,11 +2438,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2485,13 +2487,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2508,113 +2515,61 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } 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 ed32eb59a..a34774992 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h @@ -653,14 +653,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -793,14 +786,7 @@ struct QuantizedBlockLoader< return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1215,7 +1201,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1334,7 +1321,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_gather_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1477,11 +1465,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1510,13 +1499,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1543,12 +1537,11 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1556,145 +1549,119 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); + NAXTile Dtile; + Dtile.clear(); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); + } - volatile int compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } - } - xn += BK; - loader_w.next(); - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); + (void)compiler_barrier; + } + } - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + xn += BK; + loader_w.next(); + } - volatile int compiler_barrier; + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); + threadgroup_barrier(mem_flags::mem_threadgroup); - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + } - threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); - } + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } 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 187e6fffc..04206c940 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 @@ -28,30 +28,35 @@ using namespace metal; qk_dim, \ value_dim) -#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ - instantiate_kernel( \ - "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ - sdpa_vector_2pass_1_gqa, \ - type, \ - qk_dim, \ - value_dim, \ - 8, \ +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, g, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #g "_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + g, \ hpt) #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ + instantiate_sdpa_vector(type, 96, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ - instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ - instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ + instantiate_sdpa_vector(type, 512, 512) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 8, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 12, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 16, 2) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ instantiate_sdpa_vector_aggregation(type, 192) \ - instantiate_sdpa_vector_aggregation(type, 256) + instantiate_sdpa_vector_aggregation(type, 256) \ + instantiate_sdpa_vector_aggregation(type, 512) instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) 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 2f466a567..4594f0389 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h @@ -512,6 +512,7 @@ template const int num_kv_heads = tpg.x; const int num_q_heads = num_kv_heads * G; const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx; const int chunk = (N + blocks - 1) / blocks; const int kstart = block_idx * chunk; @@ -520,9 +521,9 @@ template const int s0 = kstart + cchunk * sub; const int s1 = min(kend, s0 + sub); - const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + - simd_lid * qk_per_thread; - const device T* vp = values + kv_head_idx * v_head_stride + + const device T* kp = keys + kv_batch_head_idx * k_head_stride + + s0 * k_seq_stride + simd_lid * qk_per_thread; + const device T* vp = values + kv_batch_head_idx * v_head_stride + s0 * v_seq_stride + simd_lid * v_per_thread; U q[HPT][qk_per_thread]; 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 29fa7ba39..d5cb00e42 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 @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "mlx/backend/metal/kernels/steel/attn/attn.h" @@ -63,6 +63,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -114,13 +115,15 @@ template < constexpr short LDQ_tgp = BD + padQ; constexpr short LDK_tgp = BK + padK; - constexpr short LDV_tgp = BD + padV; + constexpr short LDV_tgp = BV + padV; constexpr short tgp_mem_0 = (BK + padK) * (BD); - constexpr short tgp_mem_1 = BK * (BD + padV); + constexpr short tgp_mem_1 = BK * (BV + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = (WN == 2) && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -150,7 +153,7 @@ template < using VBlockLoader = BlockLoaderT< /* typename T = */ T, /* short BROWS = */ BK, - /* short BCOLS = */ BD, + /* short BCOLS = */ BV, /* short kDstStrRow = */ LDV_tgp, /* short kDstStrCol = */ 1, /* short reduction_dim = */ 0, @@ -168,8 +171,15 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; + + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert( + BD % WN == 0 && BV % WN == 0, + "The head dims must split evenly across WN"); - constexpr int kNWarps = WM * WN; + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -180,14 +190,19 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + constexpr int TVh = BV / (kFragSize * WN); + constexpr int BVh = BV / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -195,22 +210,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BVh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -262,17 +300,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -373,7 +445,7 @@ template < // Load V blocks if (!align_K && kb == (params->NK_aligned)) { - loader_v.load_safe(short2(BD, params->kL_rem)); + loader_v.load_safe(short2(BV, params->kL_rem)); } else { loader_v.load_unsafe(); } @@ -425,7 +497,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TVh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -461,10 +533,10 @@ template < threadgroup_barrier(mem_flags::mem_none); // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BVh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index fbd84004f..44ed0b61e 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal @@ -1,31 +1,36 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" #include "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h" -#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ - instantiate_kernel( \ - "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ - "_wm" #wm "_wn" #wn "_mask" #mname, \ - attention, dtype, bq, bk, bd, wm, wn, mtype, float) +#define instantiate_attn(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention, dtype, bq, bk, bd, bv, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) +#define instantiate_attn_helper(iname, itype, bq, bk, bd, bv, wm, wn) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, iname, itype) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, bool_, bool) -#define instantiate_attn_mask_helper(iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, bool_, bool) +#define instantiate_attn_shapes_helper(iname, itype) \ + instantiate_attn_helper(iname, itype, 32, 16, 192, 192, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 16, 128, 128, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 96, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 64, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 80, 80, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 72, 72, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 64, 64, 4, 1) -instantiate_attn_mask_helper(float16, half); -instantiate_attn_mask_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float16, half); +instantiate_attn_shapes_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float32, float); -instantiate_attn_mask_helper(float32, float); +instantiate_attn_helper(float16, half, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(float16, half, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(float32, float, 32, 16, 256, 256, 4, 2); // clang-format on 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 4a5a9716f..7a7099b8f 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 @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "mlx/backend/metal/kernels/steel/attn/nax.h" #include "mlx/backend/metal/kernels/steel/attn/params.h" @@ -76,6 +76,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -136,11 +137,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -421,8 +424,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } @@ -489,15 +492,10 @@ template < // Head-dim split attention kernel /////////////////////////////////////////////////////////////////////////////// -// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup -// accumulator working set of attention_nax (TD output fragments plus the S -// fragments) is what gates tensor-unit throughput, so this kernel splits the -// head dim across the WN = 2 simdgroups of the second warp dimension: each -// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, -// halving its accumulator set. The pair exchanges its partial Q@K.T sums -// through threadgroup memory, then both simdgroups run softmax redundantly on -// the full S tile (the row statistics are cheap) and each accumulates P@V for -// its own half of Dv. +// Split wide heads across WN simdgroups to reduce the accumulator working set. +// Each group owns one D slice. The groups exchange their partial Q@K.T scores +// through threadgroup memory and reduce them in the same order. Each group then +// runs softmax and accumulates its P@V slice. // clang-format off template < @@ -556,9 +554,9 @@ template < constexpr short kU = 16; // The WM simdgroups along the first warp dimension split the Q sequence; - // the WN simdgroups along the second split the head dim. The exchange - // below reduces exactly one peer, so WN is fixed at 2. - static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + // the WN simdgroups along the second split the head dim. + static_assert( + WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4"); constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, @@ -579,20 +577,20 @@ template < constexpr int BDh = BD / WN; static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); - static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + static_assert(TK % 2 == 0, "S fragments are computed pair by pair"); const short row_group = simd_group_id / WN; - const short d_half = simd_group_id % WN; + const short d_group = simd_group_id % WN; using otile_t = NAXTile; otile_t Otile; Otile.clear(); const short tm = kU * TQ * row_group; - Q += tm * int(params->Q_strides[2]) + d_half * BDh; - K += d_half * BDh; - V += d_half * BDh; - O += tm * int(params->O_strides[2]) + d_half * BDh; + Q += tm * int(params->Q_strides[2]) + d_group * BDh; + K += d_group * BDh; + V += d_group * BDh; + O += tm * int(params->O_strides[2]) + d_group * BDh; constexpr short kRowsPT = otile_t::kRowsPerThread; @@ -632,12 +630,11 @@ template < using stile_t = NAXTile; constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; - // One slot per (row group, half): a fragment pair in per-lane-linear - // layout. Both halves share the fragment-to-lane mapping, so the - // exchange needs no coordinate math. - threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + // One slot per (row group, D group) in per-lane-linear layout. All groups + // share the fragment-to-lane mapping. + threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32]; - // Keep the simdgroup's Q half resident in registers for the whole KV + // Keep the simdgroup's Q slice resident in registers for the whole KV // loop: TDh fragments of T are cheap next to the accumulators. NAXTile Qtiles[TDh]; STEEL_PRAGMA_UNROLL @@ -662,7 +659,7 @@ template < stile_t Stile; Stile.clear(); - // S = Q @ K.T, this half of D only, exchanged pair by pair. + // S = Q @ K.T for this D slice. STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { STEEL_PRAGMA_UNROLL @@ -686,26 +683,37 @@ template < Ktile.frag_at(1, 0), metal::true_type{}); } + } - // Exchange the partial pair and reduce. - threadgroup AccumType* slot = s_xchg[row_group][d_half]; - thread auto& s0 = Stile.frag_at(0, ik); - thread auto& s1 = Stile.frag_at(0, ik + 1); - const short base = short(simd_lane_id) * (2 * kEPF); + // Exchange all partial scores and reduce them in a fixed order. + threadgroup AccumType* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - slot[base + i] = s0[i]; - slot[base + kEPF + i] = s1[i]; + slot[base + ik * kEPF + i] = s[i]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - s0[i] += peer[base + i]; - s1[i] += peer[base + kEPF + i]; + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short peer_group = 1; peer_group < WN; peer_group++) { + const threadgroup AccumType* peer = s_xchg[row_group][peer_group]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += peer[base + ik * kEPF + i]; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); // Scale S STEEL_PRAGMA_UNROLL @@ -826,7 +834,7 @@ template < } } - // Do softmax (redundantly per half; the row statistics are cheap) + // Do softmax in each D group; the row statistics are cheap. metal::vec new_max; metal::vec factor; STEEL_PRAGMA_UNROLL @@ -854,7 +862,7 @@ template < simdgroup_barrier(mem_flags::mem_none); - // O = P @ V, this half of Dv only. + // O = P @ V for this D slice. STEEL_PRAGMA_UNROLL for (short id = 0; id < TDh; id += 2) { STEEL_PRAGMA_UNROLL 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 66d55539a..2296187fe 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 @@ -1,29 +1,31 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" #include "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h" -#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ - instantiate_kernel( \ - "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ - "_wm" #wm "_wn" #wn "_mask" #mname, \ - attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float) +#define instantiate_attn(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention_nax, dtype, bq, bk, bd, bv, wm, wn, mtype, float) #define instantiate_attn_dsplit(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ instantiate_kernel( \ - "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ + "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bd \ "_wm" #wm "_wn" #wn "_mask" #mname, \ attention_nax_dsplit, dtype, bq, bk, bd, wm, wn, mtype, float) #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn_dsplit(iname, itype, 32, 32, 512, 2, 4, mname, mtype) \ instantiate_attn_dsplit(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 96, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 64, 64, 4, 1, mname, mtype) + 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, 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) \ @@ -31,6 +33,5 @@ instantiate_attn_mask_helper(float16, half); instantiate_attn_mask_helper(bfloat16, bfloat); - instantiate_attn_mask_helper(float32, float); // clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal index 97cbe4ea1..f2a36fe10 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal @@ -23,6 +23,8 @@ #define instantiate_gemm_shapes_helper(iname, itype, oname, otype) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 256, 2, 2) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 128, 64, 2, 4) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 128, 512, 2, 4) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 64, 512, 4, 2) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 128, 256, 2, 4) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 64, 4, 4) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 256, 4, 4) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h index 4c055e69d..54e45d1e9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h @@ -20,9 +20,10 @@ template < [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { @@ -44,8 +45,19 @@ template < using loader_b_t = typename gemm_kernel::loader_b_t; using mma_t = typename gemm_kernel::mma_t; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } @@ -53,188 +65,127 @@ template < threadgroup T As[gemm_kernel::tgp_mem_size_a]; threadgroup T Bs[gemm_kernel::tgp_mem_size_b]; - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; const size_t c_row_long = size_t(c_row); const size_t c_col_long = size_t(c_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row)); const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col)); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[c_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (rhs_indices[c_row + n] != index) { - offset_next = n; - index_next = rhs_indices[c_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); + // Prepare threadgroup loading operations + thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); + thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id); + + // Prepare iterations + const int gemm_k_iterations = params->gemm_k_iterations_aligned; - // Prepare threadgroup loading operations - thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); - thread loader_b_t loader_b( - B + index * params->batch_stride_b, - params->ldb, + // Do unaligned K iterations first + if (!align_K) { + const int k_last = params->gemm_k_iterations_aligned * BK; + const int k_remain = params->K - k_last; + const size_t k_jump_a = + transpose_a ? params->lda * size_t(k_last) : size_t(k_last); + const size_t k_jump_b = + transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); + + // Move loader source ahead to end + loader_a.src += k_jump_a; + loader_b.src += k_jump_b; + + // Load tile + const short2 tile_dims_A = + transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); + const short2 tile_dims_B = + transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); + + loader_a.load_safe(tile_dims_A); + loader_b.load_safe(tile_dims_B); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Do matmul + mma_op.mma(As, Bs); + + // Reset source back to start + loader_a.src -= k_jump_a; + loader_b.src -= k_jump_b; + } + + const short lbk = 0; + + // Tile aligned don't check + if ((tgp_bm == BM) && (align_N || tgp_bn == BN)) { + gemm_kernel::gemm_loop( + As, Bs, - simd_group_id, - simd_lane_id); - - // Prepare iterations - const int gemm_k_iterations = params->gemm_k_iterations_aligned; - - // Do unaligned K iterations first - if (!align_K) { - const int k_last = params->gemm_k_iterations_aligned * BK; - const int k_remain = params->K - k_last; - const size_t k_jump_a = - transpose_a ? params->lda * size_t(k_last) : size_t(k_last); - const size_t k_jump_b = - transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); - - // Move loader source ahead to end - loader_a.src += k_jump_a; - loader_b.src += k_jump_b; - - // Load tile - const short2 tile_dims_A = - transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); - const short2 tile_dims_B = - transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); - - loader_a.load_safe(tile_dims_A); - loader_b.load_safe(tile_dims_B); + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - threadgroup_barrier(mem_flags::mem_threadgroup); + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Do matmul - mma_op.mma(As, Bs); + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Reset source back to start - loader_a.src -= k_jump_a; - loader_b.src -= k_jump_b; - } + // Nothing aligned so check both rows and cols + else { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Matrix level aligned never check - if (align_M && align_N) { - for (int k = 0; k < gemm_k_iterations; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Load elements into threadgroup - loader_a.load_unsafe(); - loader_b.load_unsafe(); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Multiply and accumulate threadgroup elements - mma_op.mma(As, Bs); - - // Prepare for next iteration - loader_a.next(); - loader_b.next(); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } else { - const short lbk = 0; - - // Tile aligned don't check - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } - - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - - // Nothing aligned so check both rows and cols - else { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - } + // Store results to device memory + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + mma_op.store_result(C, params->ldd); + } else { + mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h index 7fec27060..afd645ff9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h @@ -20,9 +20,11 @@ template < gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], + uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { constexpr short SM = BM / WM; @@ -31,106 +33,66 @@ gather_mm_rhs_nax( constexpr short TM = SM / 16; constexpr short TN = SN / 16; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; - const size_t c_row_long = size_t(c_row); - const size_t c_col_long = size_t(c_col); + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + const short sgp_sm = short(clamp(tgp_bm - tm, 0, int(SM))); + const short sgp_sn = + align_N ? SN : short(clamp(params->N - (c_col + tn), 0, int(SN))); + const size_t c_row_long = size_t(c_row + tm); + const size_t c_col_long = size_t(c_col + tn); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - rhs_indices += c_row; - - const short tm = SM * (simd_group_id / WN); - const short tn = SN * (simd_group_id % WN); - - const int sgp_sm_int = - align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); - const short sgp_sm = short(sgp_sm_int); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); - const int sgp_sn_int = - align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); - const short sgp_sn = short(sgp_sn_int); - const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); + dispatch_bool(align_K, [&](auto kAlignedK) { + dispatch_bool(sgp_sm == SM, [&](auto kAlignedM) { + dispatch_bool(sgp_sn == SN, [&](auto kAlignedN) { + NAXTile Ctile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + params->K, + params->gemm_k_iterations_aligned, + sgp_sm, + sgp_sn); - A += transpose_a ? tm : (tm * params->lda); - B += transpose_b ? (tn * params->ldb) : tn; - C += tm * params->ldd + tn; - rhs_indices += tm; - - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[0]; - short offset_next = 0; - int n = 0; - while (n < sgp_sm) { - n++; - offset = offset_next; - index = index_next; - offset_next = sgp_sm; - for (; n < sgp_sm; n++) { - if (rhs_indices[n] != index) { - offset_next = n; - index_next = rhs_indices[n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - NAXTile Ctile; - - dispatch_bool(align_K, [&](auto kAlignedK) { - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K - T, - SM, - SN, - SK, - BK, - transpose_a, - transpose_b, - kAlignedM.value, - kAlignedN.value, - kAlignedK.value, - AccumType>; - Ctile = do_gemm( - A, - B + index * params->batch_stride_b, - params->lda, - params->ldb, - params->K, - params->gemm_k_iterations_aligned, - sgp_sm, - sgp_sn); - - if constexpr (kAlignedN.value) { - if (offset_next - offset == SM) { - Ctile.store(C, int(params->ldd)); - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(SN, offset_next)); - } - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(sgp_sn, offset_next)); - } - }); + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.store(C, int(params->ldd)); + } else { + Ctile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } }); }); - } + }); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal index b1b6b8962..f848997f4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal @@ -22,6 +22,8 @@ #define instantiate_gemm_splitk_shapes_helper(iname, itype, oname, otype) \ instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 64, 64, 256, 2, 2) \ + instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 64, 128, 512, 2, 4) \ + instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 128, 64, 512, 4, 2) \ instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4) instantiate_gemm_splitk_shapes_helper(float16, half, float32, float); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/unary_ops.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/unary_ops.h index 1337f6b52..4c6088fc4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/unary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/unary_ops.h @@ -35,7 +35,7 @@ struct Abs { return x; }; complex64_t operator()(complex64_t x) thread { - return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; + return {hypot(x.real, x.imag), 0}; }; }; @@ -308,7 +308,7 @@ struct Round { struct Sigmoid { template T operator()(T x) thread { - auto y = 1 / (1 + metal::exp(metal::abs(x))); + auto y = 1 / (1 + metal::precise::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -325,8 +325,8 @@ struct Sign { if (x == complex64_t(0)) { return x; } - return x / - (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag); + auto r = hypot(x.real, x.imag); + return {x.real / r, x.imag / r}; }; }; @@ -378,7 +378,7 @@ struct Sqrt { auto b_abs = metal::precise::sqrt((r - x.real) / 2.0); auto b = metal::copysign(b_abs, x.imag); return {a, b}; - } + }; }; struct Rsqrt { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h index f15928282..405a08281 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h @@ -353,6 +353,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -470,3 +493,50 @@ template <> inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } + +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp index 8c0b46d6d..5fc2c1936 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp @@ -208,11 +208,25 @@ void steel_matmul_regular_axpby_nax( // Temp routing for larger devices char devc = d.get_architecture().back(); - if (devc == 's' || devc == 'c' || devc == 'd') { + if (devc == 's' || devc == 'c') { bk = (K >= 8192 && K > (M + N)) ? 64 : 256; bm = 64; wm = 2; + } else if (devc == 'd') { + if (N > M) { + bm = 128; + bn = 64; + + wm = 4; + wn = 2; + } else { + bm = 64; + bn = 128; + + wm = 2; + wn = 4; + } } // Prepare kernel name @@ -279,8 +293,10 @@ void steel_matmul_regular_axpby_nax( // TODO: Explore device-based tuning for swizzle int swizzle_log = tm <= 3 ? 0 : 1; - if (devc == 's' || devc == 'c' || devc == 'd') { + if (devc == 's' || devc == 'c') { swizzle_log = 2; + } else if (devc == 'd') { + swizzle_log = (M > N) ? 0 : 2; } // Prepare steel matmul params @@ -681,6 +697,33 @@ void steel_gemm_splitk_axpby_nax( int wm = 4, wn = 4; int split_k_partition_size = 4096; + // Temp routing for larger devices + char devc = d.get_architecture().back(); + if (devc == 'd') { + if (M <= 2048 && N <= 2048) { + bm = 64; + bn = 64; + + wm = 2; + wn = 2; + + split_k_partition_size = 2048; + bk = 256; + } else if (N > M) { + bm = 128; + bn = 64; + + wm = 4; + wn = 2; + } else { + bm = 64; + bn = 128; + + wm = 2; + wn = 4; + } + } + if ((M + N) / 2 < 512 || K <= 4096) { bm = bn = 64; bk = 256; @@ -754,6 +797,9 @@ void steel_gemm_splitk_axpby_nax( int tm = (M + bm - 1) / bm; int swizzle_log = tm <= 3 ? 0 : 1; + if (devc == 'd') { + swizzle_log = (M > N) ? 0 : 1; + } // Compute swizzled tile counts int tile = 1 << swizzle_log; @@ -2142,6 +2188,28 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { // GatherMM implementation /////////////////////////////////////////////////////////////////////////////// +array gather_mm_offsets( + const array& indices, + int num_groups, + int M, + metal::Device& d, + const Stream& s) { + array offsets({num_groups}, int32, nullptr, {}); + offsets.set_data(allocator::malloc(offsets.nbytes())); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel("gather_mm_offsets"); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(indices, 0); + compute_encoder.set_output_array(offsets, 1); + compute_encoder.set_bytes(M, 2); + size_t group_size = + std::min(num_groups, kernel->maxTotalThreadsPerThreadgroup()); + compute_encoder.dispatch_threads( + MTL::Size(num_groups, 1, 1), MTL::Size(group_size, 1, 1)); + compute_encoder.add_temporary(offsets); + return offsets; +} + void gather_mm_rhs( const array& a_, const array& b_, @@ -2175,12 +2243,12 @@ void gather_mm_rhs( int M = a.size() / K; int N = b.shape(-1); int lda = K; + int num_groups = b.size() / (b.shape(-2) * b.shape(-1)); // Define the dispatch blocks int bm = 16, bn = 64, bk = 16; int wm = 1, wn = 2; - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -2207,7 +2275,6 @@ void gather_mm_rhs( wn); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -2218,13 +2285,13 @@ void gather_mm_rhs( concatenate( hash_name, base_name, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, num_groups, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_gather_kernel( @@ -2245,6 +2312,7 @@ void gather_mm_rhs( // Prepare the matmul params auto batch_stride_b = b.ndim() > 2 ? b.strides()[b.ndim() - 3] : b.size(); + int tiles_m = std::min(M, (M + bm - 1) / bm + num_groups - 1); steel::GEMMParams params{ /* const int M = */ M, /* const int N = */ N, @@ -2253,7 +2321,7 @@ void gather_mm_rhs( /* const int ldb = */ static_cast(ldb), /* const int ldd = */ N, /* const int tiles_n = */ (N + bn - 1) / bn, - /* const int tiles_m = */ (M + bm - 1) / bm, + /* const int tiles_m = */ tiles_m, /* const int64_t batch_stride_a = */ 0, /* const int64_t batch_stride_b = */ static_cast(batch_stride_b), /* const int64_t batch_stride_d = */ 0, @@ -2268,9 +2336,10 @@ void gather_mm_rhs( // Launch kernel compute_encoder.set_input_array(a, 0); compute_encoder.set_input_array(b, 1); - compute_encoder.set_input_array(indices, 2); + compute_encoder.set_input_array(offsets, 2); compute_encoder.set_output_array(out, 3); compute_encoder.set_bytes(params, 4); + compute_encoder.set_bytes(num_groups, 5); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -2308,14 +2377,14 @@ void gather_mm_rhs_nax( int M = a.size() / K; int N = b.shape(-1); int lda = K; - int E = b.shape(0); + int num_groups = b.size() / (b.shape(-2) * b.shape(-1)); // Define the dispatch blocks int bm, bn = 128, bk = 128, wm, wn = 4; - if (M / E > 48) { + if (M / num_groups > 48) { bm = 64; wm = 2; - } else if (M / E > 24) { + } else if (M / num_groups > 24) { bm = 32l; wm = 1; } else { @@ -2323,7 +2392,6 @@ void gather_mm_rhs_nax( wm = 1; } - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -2350,7 +2418,6 @@ void gather_mm_rhs_nax( wn); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -2361,13 +2428,13 @@ void gather_mm_rhs_nax( concatenate( hash_name, base_name, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, num_groups, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_gather_nax_kernel( @@ -2388,6 +2455,7 @@ void gather_mm_rhs_nax( // Prepare the matmul params auto batch_stride_b = b.ndim() > 2 ? b.strides()[b.ndim() - 3] : b.size(); + int tiles_m = std::min(M, (M + bm - 1) / bm + num_groups - 1); steel::GEMMParams params{ /* const int M = */ M, /* const int N = */ N, @@ -2396,7 +2464,7 @@ void gather_mm_rhs_nax( /* const int ldb = */ static_cast(ldb), /* const int ldd = */ N, /* const int tiles_n = */ (N + bn - 1) / bn, - /* const int tiles_m = */ (M + bm - 1) / bm, + /* const int tiles_m = */ tiles_m, /* const int64_t batch_stride_a = */ 0, /* const int64_t batch_stride_b = */ static_cast(batch_stride_b), /* const int64_t batch_stride_d = */ 0, @@ -2411,9 +2479,10 @@ void gather_mm_rhs_nax( // Launch kernel compute_encoder.set_input_array(a, 0); compute_encoder.set_input_array(b, 1); - compute_encoder.set_input_array(indices, 2); + compute_encoder.set_input_array(offsets, 2); compute_encoder.set_output_array(out, 3); compute_encoder.set_bytes(params, 4); + compute_encoder.set_bytes(num_groups, 5); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/matmul.h b/Source/Cmlx/mlx/mlx/backend/metal/matmul.h index 218664b1f..13ab5338c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/matmul.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/matmul.h @@ -141,4 +141,11 @@ inline void steel_matmul( /* Strides B_batch_stride = */ B_batch_stride); } +array gather_mm_offsets( + const array& indices, + int num_groups, + int M, + metal::Device& d, + const Stream& s); + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp index 3795a6fb2..07e635def 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp @@ -79,6 +79,13 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, @@ -382,6 +389,7 @@ MTL::ComputePipelineState* get_gather_qmm_kernel( int, int, int, + bool, bool) { return d.get_kernel(kernel_name, hash_name, func_consts); } @@ -473,6 +481,7 @@ MTL::ComputePipelineState* get_gather_qmm_nax_kernel( int, int, int, + bool, bool) { return d.get_kernel(kernel_name, hash_name, func_consts); } @@ -488,6 +497,7 @@ MTL::ComputePipelineState* get_steel_attention_kernel( int, int, int, + int, const array&) { return d.get_kernel(kernel_name, hash_name, func_consts); } @@ -503,9 +513,26 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int, int, int, + int, const array&, bool) { return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_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) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp index d1d0e781c..45929e27d 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp @@ -13,7 +13,6 @@ #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/utils.h" #include "mlx/dtype_utils.h" -#include "mlx/fast_primitives.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" #include "mlx/utils.h" @@ -215,26 +214,4 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } -namespace fast { - -// There is no fused Metal cross entropy kernel yet -bool CrossEntropy::use_fallback(Stream s) { - return true; -} - -void CrossEntropy::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error("[CrossEntropy::eval_gpu] Metal cross entropy NYI."); -} - -void CrossEntropyVJP::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error( - "[CrossEntropyVJP::eval_gpu] Metal cross entropy NYI."); -} - -} // namespace fast - } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp index 1ae99a8e5..97c418e87 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp @@ -6,6 +6,7 @@ #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/matmul.h" #include "mlx/backend/metal/reduce.h" #include "mlx/backend/metal/unary.h" #include "mlx/backend/metal/utils.h" @@ -66,6 +67,9 @@ inline array ensure_row_contiguous_matrix( const array& x, metal::Device& d, const Stream& s) { + if (x.flags().row_contiguous) { + return x; + } if (x.ndim() < 2) { if (x.strides()[0] == 1) { return x; @@ -816,6 +820,7 @@ void qmm_nax( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, bool transpose, int group_size, @@ -830,8 +835,9 @@ void qmm_nax( int wm = 2; int wn = 2; - // Use smaller bm when one block covers all of M. - int bm = (M <= 32) ? 32 : 64; + // Use smaller bm when one block covers all of M. Only qmm_t_nax has a 32-row + // instantiation. + int bm = (transpose && M <= 32) ? 32 : 64; int bn = 64; int bk = 64; MTL::Size group_dims(32, wn, wm); @@ -861,7 +867,8 @@ void qmm_nax( "_wn", wn, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); std::string template_def; MTL::ComputePipelineState* kernel; if (transpose) { @@ -879,7 +886,8 @@ void qmm_nax( bk, bn, wm, - wn); + wn, + global_scale.has_value()); } else { kernel = get_qmm_nax_kernel_wrapped( d, @@ -904,6 +912,11 @@ void qmm_nax( compute_encoder.set_input_array(scales, c++); if (biases) { compute_encoder.set_input_array(*biases, c++); + } else if (transpose) { + if (global_scale) { + compute_encoder.set_input_array(*global_scale, c); + } + c++; } compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); @@ -920,6 +933,7 @@ void gather_qmm_nax( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -966,7 +980,8 @@ void gather_qmm_nax( wm, "_wn", wn, - transpose ? (aligned ? "_alN_true" : "_alN_false") : ""); + transpose ? (aligned ? "_alN_true" : "_alN_false") : "", + global_scale ? "_hgs" : ""); MTL::ComputePipelineState* kernel; if (transpose) { kernel = get_qmm_nax_kernel_wrapped( @@ -982,7 +997,8 @@ void gather_qmm_nax( bk, bn, wm, - wn); + wn, + global_scale.has_value()); } else { kernel = get_qmm_nax_kernel_wrapped( d, @@ -1002,12 +1018,14 @@ void gather_qmm_nax( auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -1026,6 +1044,7 @@ void qmm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, bool transpose, int group_size, @@ -1036,15 +1055,21 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { + assert( + (!global_scale || (transpose && !biases)) && + "qmm global scale requires transposed weights without biases"); + bool has_nax_kernel = metal::is_nax_available() && (transpose || mode == "affine"); - if (has_nax_kernel && transpose && (K % 64 == 0) && + bool nax_aligned = (K % 64 == 0) && (transpose || N % 64 == 0); + if (has_nax_kernel && nax_aligned && (env::enable_tf32() || x.dtype() != float32)) { return qmm_nax( /* const array& x = */ x, /* const array& w = */ w, /* const array& scales = */ scales, /* const std::optional& biases = */ biases, + /* const std::optional& global_scale = */ global_scale, /* array& out = */ out, /* bool transpose = */ transpose, /* int group_size = */ group_size, @@ -1080,7 +1105,8 @@ void qmm( "_b_", bits, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); std::string template_def; MTL::ComputePipelineState* kernel; if (transpose) { @@ -1093,7 +1119,8 @@ void qmm( group_size, bits, aligned, - batched); + batched, + global_scale.has_value()); } else { kernel = get_quantized_kernel_wrapped( d, kname, "qmm_n", mode, type_string, group_size, bits, batched); @@ -1106,6 +1133,11 @@ void qmm( compute_encoder.set_input_array(scales, c++); if (biases) { compute_encoder.set_input_array(*biases, c++); + } else if (transpose) { + if (global_scale) { + compute_encoder.set_input_array(*global_scale, c); + } + c++; } compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); @@ -1151,7 +1183,21 @@ void qmm_splitk( } if (split_k <= 1) { return qmm( - x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode); + x, + w, + scales, + biases, + std::nullopt, + out, + true, + group_size, + bits, + M, + N, + K, + d, + s, + mode); } int k_partition_size = K / split_k; @@ -1221,6 +1267,7 @@ void gather_qmm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -1240,6 +1287,7 @@ void gather_qmm( /* const array& w = */ w, /* const array& scales = */ scales, /* const std::optional& biases = */ biases, + /* const std::optional& global_scale = */ global_scale, /* const array& lhs_indices = */ lhs_indices, /* const array& rhs_indices = */ rhs_indices, /* array& out = */ out, @@ -1275,25 +1323,43 @@ void gather_qmm( group_size, "_b_", bits, - transpose ? (aligned ? "_alN_true" : "_alN_false") : ""); + transpose ? (aligned ? "_alN_true" : "_alN_false") : "", + global_scale ? "_hgs" : ""); MTL::ComputePipelineState* kernel; if (transpose) { kernel = get_quantized_kernel_wrapped( - d, kname, "gather_qmm_t", mode, type_string, group_size, bits, aligned); + d, + kname, + "gather_qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + global_scale.has_value()); } else { kernel = get_quantized_kernel_wrapped( - d, kname, "gather_qmm_n", mode, type_string, group_size, bits); + d, + kname, + "gather_qmm_n", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); } auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -1452,6 +1518,7 @@ void gather_qmm_rhs_nax( const array& w_, const array& scales_, const std::optional& biases_, + const std::optional& global_scale, const array& indices_, array& out, bool transpose, @@ -1491,6 +1558,10 @@ void gather_qmm_rhs_nax( if (biases_) { biases = ensure_row_contiguous(*biases_, d, s); } + std::optional gs; + if (global_scale) { + gs = ensure_row_contiguous(*global_scale, d, s); + } // Use smaller bm for many experts and few tokens. int E = w.size() / w.shape(-1) / w.shape(-2); @@ -1498,7 +1569,6 @@ void gather_qmm_rhs_nax( int bn = 64, bk = 64; int wm = 2, wn = 2; - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -1524,10 +1594,10 @@ void gather_qmm_rhs_nax( "_wm_", wm, "_wn_", - wn); + wn, + global_scale ? "_hgs" : ""); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -1538,13 +1608,13 @@ void gather_qmm_rhs_nax( concatenate( hash_name, kname, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, E, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_nax_kernel( @@ -1561,24 +1631,29 @@ void gather_qmm_rhs_nax( bk, wm, wn, - transpose); + transpose, + global_scale.has_value()); compute_encoder.set_compute_pipeline_state(kernel); MTL::Size group_dims(32, wn, wm); - MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + MTL::Size grid_dims( + (N + bn - 1) / bn, std::min(M, (M + bm - 1) / bm + E - 1), 1); - int c = 0; - compute_encoder.set_input_array(x, c++); - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(x, 0); + compute_encoder.set_input_array(w, 1); + compute_encoder.set_input_array(scales, 2); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 3); + } else if (gs) { + compute_encoder.set_input_array(*gs, 3); } - compute_encoder.set_input_array(indices, c++); + int c = 4; + compute_encoder.set_input_array(offsets, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(M, c++); compute_encoder.set_bytes(N, c++); compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(E, c++); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -1588,6 +1663,7 @@ void gather_qmm_rhs( const array& w_, const array& scales_, const std::optional& biases_, + const std::optional& global_scale, const array& indices_, array& out, bool transpose, @@ -1606,6 +1682,7 @@ void gather_qmm_rhs( /* const array& w_ = */ w_, /* const array& scales_ = */ scales_, /* const std::optional& biases_ = */ biases_, + /* const std::optional& global_scale = */ global_scale, /* const array& indices_ = */ indices_, /* array& out = */ out, /* bool transpose = */ transpose, @@ -1647,12 +1724,16 @@ void gather_qmm_rhs( if (biases_) { biases = ensure_row_contiguous(*biases_, d, s); } + std::optional gs; + if (global_scale) { + gs = ensure_row_contiguous(*global_scale, d, s); + } // TODO: Tune the block sizes + int E = w.size() / w.shape(-1) / w.shape(-2); int bm = 16, bn = 32, bk = 32; int wm = 1, wn = 2; - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -1677,10 +1758,10 @@ void gather_qmm_rhs( "_wm_", wm, "_wn_", - wn); + wn, + global_scale ? "_hgs" : ""); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -1691,13 +1772,13 @@ void gather_qmm_rhs( concatenate( hash_name, kname, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, E, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_kernel( @@ -1714,24 +1795,29 @@ void gather_qmm_rhs( bk, wm, wn, - transpose); + transpose, + global_scale.has_value()); compute_encoder.set_compute_pipeline_state(kernel); MTL::Size group_dims(32, wn, wm); - MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + MTL::Size grid_dims( + (N + bn - 1) / bn, std::min(M, (M + bm - 1) / bm + E - 1), 1); - int c = 0; - compute_encoder.set_input_array(x, c++); - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(x, 0); + compute_encoder.set_input_array(w, 1); + compute_encoder.set_input_array(scales, 2); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 3); + } else if (gs) { + compute_encoder.set_input_array(*gs, 3); } - compute_encoder.set_input_array(indices, c++); + int c = 4; + compute_encoder.set_input_array(offsets, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(M, c++); compute_encoder.set_bytes(N, c++); compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(E, c++); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -1816,6 +1902,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, out, transpose_, group_size_, @@ -1883,9 +1970,13 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { array x = ensure_row_contiguous_matrix(inputs[0], d, s); array w = ensure_row_contiguous_matrix(inputs[1], d, s); array scales = ensure_row_contiguous_matrix(inputs[2], d, s); + // Affine gets biases at index 3, nvfp4 an optional global scale. std::optional biases = std::nullopt; - if (inputs.size() == 6) { + std::optional global_scale = std::nullopt; + if (mode_ == QuantizationMode::Affine) { biases = ensure_row_contiguous_matrix(inputs[3], d, s); + } else if (inputs.size() == 6) { + global_scale = ensure_row_contiguous(inputs[3], d, s); } const array& lhs_indices = inputs[inputs.size() - 2]; const array& rhs_indices = inputs[inputs.size() - 1]; @@ -1908,6 +1999,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + global_scale, rhs_indices, out, transpose_, @@ -1929,6 +2021,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + global_scale, lhs_indices, rhs_indices, out, @@ -1950,7 +2043,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, - std::nullopt, + global_scale, lhs_indices, rhs_indices, out, @@ -1970,7 +2063,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, - std::nullopt, + global_scale, lhs_indices, rhs_indices, out, @@ -2026,6 +2119,26 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = x.shape(-1); int M = non_batched ? x.size() / K : x.shape(-2); int N = out.shape(-1); + if (has_global_scales && w_quantized && non_batched && + x.dtype() == bfloat16 && K % 32 == 0 && + M >= get_qmv_batch_limit(K, N, d)) { + qmm(x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + true, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } dispatch_qmv( x, w_q, @@ -2087,6 +2200,66 @@ void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { int K = x.shape(-1); int M = non_batched ? x.size() / K : x.shape(-2); int N = out.shape(-1); + + // temporary, until we add proper scaling for gather qqmm + if (global_scale_w) { + int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2); + array gs_e(Shape{E}, float32, nullptr, {}); + broadcast(*global_scale_w, gs_e); + global_scale_w = ensure_row_contiguous(gs_e, d, s); + } + + int B = out.size() / M / N; + bool use_matrix_kernels = has_global_scales && w_quantized && + x.dtype() == bfloat16 && w_q.ndim() == 3 && + K % (metal::is_nax_available() ? 64 : 32) == 0; + + if (use_matrix_kernels && M == 1 && B >= 16 && right_sorted_) { + int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2); + if (B / E >= 4) { + gather_qmm_rhs( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + rhs_indices, + out, + true, + group_size_, + bits_, + x.size() / K, + N, + K, + d, + s, + mode); + return; + } + } + + if (use_matrix_kernels && M >= get_qmv_batch_limit(K, N, d)) { + gather_qmm( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + true, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + gather_qmv( x, w_q, 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 981d2b181..3f5ed8336 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,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include #include "mlx/backend/common/compiled.h" @@ -29,13 +29,13 @@ void sdpa_full_self_attention_nax( using namespace mlx::steel; int bd = q.shape(-1); - int bq = 64; + int bv = v.shape(-1); + int bq = bd == 512 ? 32 : 64; int bk = 32; - bool split_d = bd == 256; - int wm = 4; - int wn = split_d ? 2 : 1; - + 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); int H = q.shape(1); int D = q.shape(3); @@ -98,6 +98,8 @@ void sdpa_full_self_attention_nax( bk, "_bd", bd, + "_bv", + bv, "_wm", wm, "_wn", @@ -131,6 +133,7 @@ void sdpa_full_self_attention_nax( bq, bk, bd, + bv, wm, wn, (has_mask ? *mask : q), @@ -216,7 +219,7 @@ void sdpa_full_self_attention_metal( int kL = k.shape(2); if (metal::is_nax_available() && - (D == 64 || D == 96 || D == 128 || D == 256) && + (D == 64 || D == 96 || D == 128 || D == 256 || D == 512) && (env::enable_tf32() || q.dtype() != float32)) { return sdpa_full_self_attention_nax( /* const Stream& s = */ s, @@ -231,14 +234,88 @@ void sdpa_full_self_attention_metal( /* const std::optional& sinks = */ sinks); } - using namespace mlx::steel; + // 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. + bool pad_default = qL >= 512 && kL >= 512 && !do_causal_ && !mask && !sinks; + if ((D == 72 || D == 80) && 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; + auto& enc = metal::get_command_encoder(s); + array zero = array(0, q.dtype()); + + auto pad_head_dim = [&](const array& x) { + Shape padded_shape = x.shape(); + padded_shape.back() = pad_to; + array xp(std::move(padded_shape), x.dtype(), nullptr, {}); + fill_gpu(zero, xp, s); + // Explicit output strides: a [.., D] view over xp would be + // non-contiguous with span > size, so it cannot carry xp's flags. + copy_gpu_inplace( + /* const array& in = */ x, + /* array& out = */ xp, + /* const Shape& data_shape = */ x.shape(), + /* const Strides& i_strides = */ x.strides(), + /* const Strides& o_strides = */ xp.strides(), + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ 0, + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ s); + enc.add_temporary(xp); + return xp; + }; - int wm = 4; - int wn = 1; + array qp = pad_head_dim(q); + array kp = pad_head_dim(k); + array vp = pad_head_dim(v); + + Shape padded_out = o.shape(); + padded_out.back() = pad_to; + array op(std::move(padded_out), o.dtype(), nullptr, {}); + op.set_data(allocator::malloc(op.nbytes())); + + sdpa_full_self_attention_nax( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& q = */ qp, + /* const array& k = */ kp, + /* const array& v = */ vp, + /* const float scale = */ scale, + /* array& o = */ op, + /* bool do_causal_ = */ do_causal_, + /* const std::optional& mask = */ mask, + /* const std::optional& sinks = */ sinks); + // Slice the padded lanes back out into o's caller-chosen strides. + copy_gpu_inplace( + /* const array& in = */ op, + /* array& out = */ o, + /* const Shape& data_shape = */ o.shape(), + /* const Strides& i_strides = */ op.strides(), + /* const Strides& o_strides = */ o.strides(), + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ 0, + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ s); + enc.add_temporary(op); + enc.add_temporary(zero); + return; + } + + using namespace mlx::steel; + + char devc = d.get_architecture().back(); int bd = q.shape(-1); + int bv = v.shape(-1); int bq = 32; - int bk = bd < 128 ? 32 : 16; + int bk = (bd == 256) && (q.dtype() != float32) && (devc == 'd') + ? 32 + : (bd < 128 ? 32 : 16); + + int wm = 4; + int wn = (bd == 256) ? 2 : 1; const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; @@ -264,6 +341,8 @@ void sdpa_full_self_attention_metal( bk, "_bd", bd, + "_bv", + bv, "_wm", wm, "_wn", @@ -297,6 +376,7 @@ void sdpa_full_self_attention_metal( bq, bk, bd, + bv, wm, wn, (has_mask ? *mask : q)); @@ -466,10 +546,15 @@ void sdpa_vector_2pass( std::string kname; kname.reserve(64); kname += "sdpa_vector_2pass_1"; - if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) && - q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && - k.shape(2) >= 8192) { - kname += "_gqa"; + int gqa_factor = q.shape(1) / k.shape(1); + bool gqa_dims = + (gqa_factor == 8 && (q.shape(-1) == 64 || q.shape(-1) == 128)) || + ((gqa_factor == 12 || gqa_factor == 16) && q.shape(-1) == 128); + if (!mask && !sinks && q.shape(2) == 1 && + q.shape(1) == gqa_factor * k.shape(1) && q.shape(-1) == v.shape(-1) && + gqa_dims && k.shape(2) >= 8192) { + kname += "_gqa_"; + kname += std::to_string(gqa_factor); } kname += "_"; kname += get_type_string(q.dtype()); @@ -479,7 +564,6 @@ void sdpa_vector_2pass( kname += std::to_string(v.shape(-1)); // Compute the necessary sizes - int gqa_factor = q.shape(1) / k.shape(1); int n_simds = gqa_factor * q.shape(2); char devc = d.get_architecture().back(); @@ -659,15 +743,22 @@ std::tuple has_fused_kernel( std::ostringstream msg; if (query_sequence_length > 8) { - const bool supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 || - query_head_dim == 96 || query_head_dim == 128 || - query_head_dim == 192 || query_head_dim == 256); + const bool supports_d512 = metal::is_nax_available() && + (env::enable_tf32() || q.dtype() != float32); + const bool asymmetric = query_head_dim == 96 && value_head_dim == 64; + const bool supported_head_dim = asymmetric || + (query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 72 || + query_head_dim == 80 || query_head_dim == 96 || + query_head_dim == 128 || query_head_dim == 192 || + query_head_dim == 256 || (query_head_dim == 512 && supports_d512))); if (!supported_head_dim) { msg << "the full attention kernel supports head dims " - << "{64, 72, 80, 96, 128, 192, 256} with matching query/value head " - << "dims; got query head dim " << query_head_dim - << " and value head dim " << value_head_dim << "."; + << "{64, 72, 80, 96, 128, 192, 256} with matching query/value dims, " + << "or (query, value) head dims (96, 64), " + << "plus head dim 512 on NAX GPUs (float32 also requires TF32); got " + << "query head dim " << query_head_dim << " and value head dim " + << value_head_dim << "."; return {false, msg.str()}; } if (has_mask && !has_arr_mask && @@ -683,13 +774,14 @@ std::tuple has_fused_kernel( (query_head_dim == value_head_dim && (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || query_head_dim == 192 || - query_head_dim == 256)) || - (query_head_dim == 192 && value_head_dim == 128); + query_head_dim == 256 || query_head_dim == 512)) || + (query_head_dim == 192 && value_head_dim == 128) || + (query_head_dim == 96 && value_head_dim == 64); if (!supported_head_dim) { msg << "the vector attention kernel supports head dims " - << "{64, 96, 128, 192, 256} with matching query/value head dims, " - << "or query head dim 192 with value head dim 128; got query head " - << "dim " << query_head_dim << " and value head dim " + << "{64, 96, 128, 192, 256, 512} with matching query/value head " + << "dims, or (query, value) head dims (96, 64) or (192, 128); got " + << "query head dim " << query_head_dim << " and value head dim " << value_head_dim << "."; return {false, msg.str()}; } @@ -706,6 +798,36 @@ std::tuple has_fused_kernel( << query_sequence_length << " and GQA factor " << gqa_factor << "."; return {false, msg.str()}; } + // Head dim 512 uses the generic vector kernel, reading K/V is its main + // cost. By default, use it only for one query, GQA factor 8, and no array + // mask. + if (query_head_dim == 512) { + int min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", 1024); + bool always = (min_key_sequence_length == 0); + if (!always && query_sequence_length != 1) { + msg << "the vector attention kernel for head dim 512 defaults to " + "single-token queries; got query length " + << query_sequence_length << "."; + return {false, msg.str()}; + } + if (!always && gqa_factor != 8) { + msg << "the vector attention kernel for head dim 512 defaults to " + "GQA factor 8; got GQA factor " + << gqa_factor << "."; + return {false, msg.str()}; + } + if (!always && has_arr_mask) { + msg << "the vector attention kernel for head dim 512 defaults to " + "calls without array masks."; + return {false, msg.str()}; + } + if (key_sequence_length < min_key_sequence_length) { + msg << "the vector attention kernel for head dim 512 requires at " + << "least " << min_key_sequence_length << " keys; got key length " + << key_sequence_length << "."; + return {false, msg.str()}; + } + } } return {true, ""}; } @@ -749,12 +871,30 @@ bool ScaledDotProductAttention::use_fallback( const int query_head_dim = q.shape(-1); const int value_head_dim = v.shape(-1); + if (query_head_dim == 512 && query_sequence_length > 8) { + constexpr int64_t min_query_blocks = 1024; + const int64_t query_blocks = int64_t(q.shape(0)) * q.shape(1) * + int64_t(ceildiv(query_sequence_length, 32)); + // The D512 kernel needs this many query blocks to match fallback speed. + const bool eligible = metal::is_nax_available() && + (env::enable_tf32() || q.dtype() != float32) && + query_sequence_length >= 1024 && do_causal && !has_arr_mask && + query_blocks >= min_query_blocks; + return !eligible; + } + // Use headdim-split kernel when NAX is enabled and there are enough query // blocks to fill the machine. if (metal::is_nax_available() && (env::enable_tf32() || q.dtype() != float32) && - query_sequence_length >= 1024 && query_head_dim == 256 && do_causal && - !has_arr_mask) { + query_sequence_length >= 1024 && query_head_dim == 256 && + (do_causal || has_arr_mask)) { + return false; + } + + if (!metal::is_nax_available() && q.dtype() != float32 && do_causal && + query_head_dim == 256 && query_sequence_length >= 2048 && + query_sequence_length == k.shape(2)) { return false; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp index 18935191f..02d707f2b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp @@ -116,6 +116,11 @@ void scan_gpu_inplace( void Scan::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); + if (out.size() == 0) { + out.set_data(allocator::malloc(0)); + return; + } + auto in = inputs[0]; if (in.flags().contiguous && in.strides()[axis_] != 0) { if (in.is_donatable() && in.itemsize() == out.itemsize()) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp index 65f144c02..15d6bdfb4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp @@ -278,6 +278,10 @@ void gpu_merge_sort( array& out, int axis_, bool argsort) { + if (out.size() == 0) { + return; + } + // Get size info int axis = axis_ < 0 ? axis_ + in.ndim() : axis_; int size_sorted_axis = in.shape(axis); diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp index 05852c860..2e027032d 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp @@ -16,15 +16,18 @@ Fence::Fence(Stream s) { fence_ = std::make_shared(0, s); } -void Fence::wait(Stream s, const array&) { - cast().event.wait(s); +void Fence::wait(Stream s, const array&, uint32_t value) { + auto event = cast().event; + event.set_value(value); + event.wait(s); } -void Fence::update(Stream s, const array&, bool) { +uint32_t Fence::update(Stream s, const array&, bool) { auto& f = cast(); f.count++; f.event.set_value(f.count); f.event.signal(s); + return f.count; } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp index f17f12cfa..c37515f9e 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp @@ -52,6 +52,16 @@ bool fast::ScaledDotProductAttention::supports_bool_mask() { return false; } +bool fast::GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + return true; +} + NO_GPU(Abs) NO_GPU(Add) NO_GPU(AddMM) @@ -179,6 +189,7 @@ NO_GPU_MULTI(RMSNormVJP) NO_GPU_USE_FALLBACK(RoPE) NO_GPU_MULTI(ScaledDotProductAttention) NO_GPU_MULTI(ScaledDotProductAttentionVJP) +NO_GPU_MULTI(GatedDeltaUpdate) 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 bb17c4496..3cab4c05d 100644 --- a/Source/Cmlx/mlx/mlx/compile.cpp +++ b/Source/Cmlx/mlx/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -302,9 +303,9 @@ std::uintptr_t get_function_address(const std::function& fun) { class CompileCache { public: struct CacheEntry { - CacheEntry(Stream stream, bool shapeless) - : stream(stream), shapeless(shapeless) {}; - Stream stream; + CacheEntry(std::optional stream, bool shapeless) + : stream(stream), shapeless(shapeless) {} + std::optional stream; bool shapeless; std::vector inputs; std::vector outputs; @@ -370,7 +371,7 @@ class CompileCache { // Loop over entries and check: // - Default stream and device match the entry's default stream // - Inputs match i.e. shapes and types must be equal. - auto stream = default_stream(default_device()); + auto stream = peek_default_stream(default_device()); for (CacheEntry& entry : entries) { // Check that the default stream and device match if (entry.stream != stream) { 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 15b81a217..d6305d9d4 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 @@ -198,7 +198,11 @@ class RingImpl { size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); dispatch_wires(n_wires, [&](int lw) { all_gather_wire( - out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + out_ptr, + n_bytes, + static_cast(n_bytes_per_wire), + n_wires, + lw); }); } @@ -215,14 +219,21 @@ class RingImpl { char* out_ptr, int64_t n_bytes, int64_t n_bytes_per_wire, + int n_wires, int lw) { - // Both directions send the same contiguous slice of each rank's region. - int64_t slice = static_cast(lw) * n_bytes_per_wire; - int64_t wire_offset[2] = {slice, slice}; - // Clamp to this wire's slice end so the last frame can't spill into the - // next wire's slice. - int64_t wire_end_bytes = std::min(n_bytes, slice + n_bytes_per_wire); - int64_t wire_end[2] = {wire_end_bytes, wire_end_bytes}; + // Each rank's region is divided into 2 directional halves, each split + // into n_wires contiguous slices - mirroring all_reduce_wire. Without + // the lr term both directions cover the same lower-half slice and the + // upper half of every peer's region is never transferred. + int64_t wire_offset[2]; + int64_t wire_end[2]; + for (int lr = 0; lr < 2; lr++) { + wire_offset[lr] = lr * n_wires * n_bytes_per_wire + + static_cast(lw) * n_bytes_per_wire; + int64_t region_end = + std::min(n_bytes, (lr + 1) * n_wires * n_bytes_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + n_bytes_per_wire); + } int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; int64_t recv_offset[2] = { ((rank_ + size_ - 1) % size_) * n_bytes, @@ -230,7 +241,7 @@ class RingImpl { ring_pass<2, char>( lw, - n_bytes, + n_bytes * size_, n_bytes, n_bytes * size_, n_bytes_per_wire, diff --git a/Source/Cmlx/mlx/mlx/export.cpp b/Source/Cmlx/mlx/mlx/export.cpp index bcc0ad161..55ccbe187 100644 --- a/Source/Cmlx/mlx/mlx/export.cpp +++ b/Source/Cmlx/mlx/mlx/export.cpp @@ -702,6 +702,7 @@ FunctionExporter::FunctionExporter( fun(std::move(fun)), ftable(std::make_shared(shapeless)), metadata_(std::move(metadata)) { + os.open(); if (!os.is_open()) { throw std::runtime_error("[export_function] Failed to open " + file); } diff --git a/Source/Cmlx/mlx/mlx/fast.cpp b/Source/Cmlx/mlx/mlx/fast.cpp index d6321ed36..101fe127c 100644 --- a/Source/Cmlx/mlx/mlx/fast.cpp +++ b/Source/Cmlx/mlx/mlx/fast.cpp @@ -1032,6 +1032,133 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { has_sinks_ == a_other.has_sinks_; } +std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state, /* = std::nullopt */ + const std::optional& mask_, /* = std::nullopt */ + StreamOrDevice s_ /* = {} */) { + // determine output dtype + auto s = to_stream(s_); + + auto promoted = promote_types(queries.dtype(), keys.dtype()); + auto out_dtype = issubdtype(promoted, floating) + ? promoted + : promote_types(promoted, float32); + + // cast all inputs + auto q = astype(queries, out_dtype, s); + auto k = astype(keys, out_dtype, s); + auto v = astype(values, out_dtype, s); + auto g = astype(gates, out_dtype, s); + auto beta = astype(beta_, out_dtype, 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); + + auto h0 = initial_state.has_value() ? astype(*initial_state, float32, s) + : zeros({B, Hv, Dv, Dk}, float32, s); + + bool has_mask = mask_.has_value(); + auto mask = has_mask ? astype(*mask_, bool_, s) : array(false); + + auto fallback = [B, T, Hk, Dk, Hv, Dv, has_mask, s]( + std::vector inputs) { + auto out_dtype = inputs[0].dtype(); + auto q = inputs[0]; + auto k = inputs[1]; + auto v = inputs[2]; + auto g = inputs[3]; + auto beta = inputs[4]; + auto state = astype(inputs[5], float32, s); + + if (Hv != Hk) { + int repeat_factor = Hv / Hk; + q = repeat(q, repeat_factor, 2, s); + k = repeat(k, repeat_factor, 2, s); + } + + array mask = has_mask ? astype(inputs[6], bool_, s) : array(false); + const array zero = array(0.0f, float32); + + std::vector outputs; + for (int t = 0; t < T; t++) { + auto get_t = [&](const array& a, int t) { + Shape start(a.ndim(), 0), stop = a.shape(); + start[1] = t; + stop[1] = t + 1; + return squeeze(slice(a, start, stop, s), 1, s); + }; + auto q_t = get_t(q, t); + auto k_t = get_t(k, t); + auto v_t = get_t(v, t); + auto g_t = get_t(g, t); + auto beta_t = get_t(beta, t); + array mask_t = has_mask ? get_t(mask, t) : array(true); + + auto state_prev = state; + + auto decay = (g_t.ndim() == 2) + ? expand_dims(g_t, {-1, -2}, s) // [B,H,1,1] + : expand_dims(g_t, -2, s); // [B,H,1,Dk] + + auto state_next = multiply(state_prev, decay, s); + + auto kv = + sum(multiply(state_next, expand_dims(k_t, -2, s), s), + -1, + false, + s); // [B,H,Dv] + + auto delta = subtract(v_t, kv, s); + delta = multiply(delta, expand_dims(beta_t, -1, s), s); + + state_next = + add(state_next, + multiply(expand_dims(delta, -1, s), expand_dims(k_t, -2, s), s), + s); + + if (has_mask) { + auto state_mask = expand_dims(mask_t, {-1, -2, -3}, s); + state = where(state_mask, state_next, state_prev, s); + } else { + state = state_next; + } + + auto o_t = sum(multiply(state, expand_dims(q_t, -2, s), s), -1, false, s); + + if (has_mask) { + auto out_mask = expand_dims(mask_t, {-1, -2}, s); + o_t = where(out_mask, o_t, zero, s); + } + outputs.push_back(astype(o_t, out_dtype, s)); + } + auto out = stack(outputs, 1, s); + return std::vector{out, state}; + }; + + if (!GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { + auto result = array::make_arrays( + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, + /* dtypes */ {out_dtype, float32}, + /* primitive */ + std::make_shared(s, fallback), + /* inputs */ {q, k, v, g, beta, h0}); + + return result; + } + + auto result = fallback({q, k, v, g, beta, h0, mask}); + return result; +} + bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); return ( diff --git a/Source/Cmlx/mlx/mlx/fast.h b/Source/Cmlx/mlx/mlx/fast.h index e79f300ce..072a3ad0c 100644 --- a/Source/Cmlx/mlx/mlx/fast.h +++ b/Source/Cmlx/mlx/mlx/fast.h @@ -60,6 +60,16 @@ MLX_API array scaled_dot_product_attention( bool force_fused = false, StreamOrDevice s = {}); +MLX_API std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state = std::nullopt, + const std::optional& mask = std::nullopt, + StreamOrDevice s = {}); + using TemplateArg = std::variant; using ScalarArg = std::variant; diff --git a/Source/Cmlx/mlx/mlx/fast_primitives.h b/Source/Cmlx/mlx/mlx/fast_primitives.h index 63021cd0c..cef13bb0f 100644 --- a/Source/Cmlx/mlx/mlx/fast_primitives.h +++ b/Source/Cmlx/mlx/mlx/fast_primitives.h @@ -395,6 +395,38 @@ class ConvertFP8 : public Primitive { bool to_fp8_; }; +class GatedDeltaUpdate : public Custom { + public: + GatedDeltaUpdate( + 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, + const bool has_mask, + 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(GatedDeltaUpdate); + DEFINE_INPUT_OUTPUT_SHAPE() + auto state() const { + return std::make_tuple(nullptr); /* TODO */ + } + + private: +}; + class Quantize : public Custom { public: explicit Quantize( diff --git a/Source/Cmlx/mlx/mlx/fence.h b/Source/Cmlx/mlx/mlx/fence.h index 3fd5da333..fd6bc116a 100644 --- a/Source/Cmlx/mlx/mlx/fence.h +++ b/Source/Cmlx/mlx/mlx/fence.h @@ -1,6 +1,6 @@ // Copyright © 2024 Apple Inc. -#include +#include #include "mlx/array.h" @@ -8,12 +8,9 @@ namespace mlx::core { /* A fence to be used for synchronizing work between streams. * - * Calls to `wait` wait in the given stream until all previous calls to update - * are complete on their given stream. - * - * The array passed to `update` is computed and visible after the call to - * `wait` returns. The array passed to `wait` will not be read until all - * previous calls to `update` have completed. + * `update` returns a value that marks when its array is computed and visible. + * `wait` orders work in the consumer stream after that value is signaled. + * Later updates do not extend an earlier array's wait. * * Note, calls to `update` should always be from the same thread or explicitly * synchronized so that they occur in sequence. Calls to `wait` can be on any @@ -29,8 +26,8 @@ class Fence { Fence() {}; explicit Fence(Stream stream); - void update(Stream stream, const array& x, bool cross_device); - void wait(Stream stream, const array& x); + uint32_t update(Stream stream, const array& x, bool cross_device); + void wait(Stream stream, const array& x, uint32_t value); template auto& cast() const { diff --git a/Source/Cmlx/mlx/mlx/fft.cpp b/Source/Cmlx/mlx/mlx/fft.cpp index 06860a0e3..55596472a 100644 --- a/Source/Cmlx/mlx/mlx/fft.cpp +++ b/Source/Cmlx/mlx/mlx/fft.cpp @@ -61,7 +61,7 @@ array fft_impl( std::ostringstream msg; msg << "[fftn] Invalid axis received for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // In the following shape manipulations there are three cases to consider: @@ -262,7 +262,7 @@ array fftshift_impl( std::ostringstream msg; msg << "[" << name << "] Invalid axis " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Match NumPy's implementation int shift = a.shape(axis) / 2; diff --git a/Source/Cmlx/mlx/mlx/io/gguf.cpp b/Source/Cmlx/mlx/mlx/io/gguf.cpp index 6c27c4698..dddcbfc77 100644 --- a/Source/Cmlx/mlx/mlx/io/gguf.cpp +++ b/Source/Cmlx/mlx/mlx/io/gguf.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "mlx/io/gguf.h" @@ -47,11 +48,30 @@ std::optional gguf_type_to_dtype(const uint32_t& gguf_type) { } } +[[noreturn]] void tensor_error( + const gguf_tensor& tensor, + const std::string& what) { + std::ostringstream msg; + msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) + << "' " << what << ". Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); +} + Shape get_shape(const gguf_tensor& tensor) { Shape shape; + uint64_t num_weights = 1; // The dimension order in GGML is the reverse of the order used in MLX. for (int i = tensor.ndim - 1; i >= 0; i--) { - shape.push_back(tensor.dim[i]); + auto dim = tensor.dim[i]; + if (dim > std::numeric_limits::max()) { + tensor_error(tensor, "has a dimension that is too large"); + } + if (dim != 0 && num_weights > std::numeric_limits::max() / dim) { + tensor_error( + tensor, "has a dimension product that does not fit in 64 bits"); + } + num_weights *= dim; + shape.push_back(dim); } return shape; } @@ -303,20 +323,14 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { } void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) { - auto fail = [&tensor](const std::string& what) { - std::ostringstream msg; - msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) - << "' " << what << ". Perhaps an incomplete download or corrupt file?"; - throw std::runtime_error(msg.str()); - }; if (tensor.offset < ctx->data_off) { - fail("has a data offset that overflows the data section"); + tensor_error(tensor, "has a data offset that overflows the data section"); } if (tensor.offset > ctx->size) { - fail("has a data offset past the end of the file"); + tensor_error(tensor, "has a data offset past the end of the file"); } if (tensor.bsize > ctx->size - tensor.offset) { - fail("extends past the end of the file"); + tensor_error(tensor, "extends past the end of the file"); } } diff --git a/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp b/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp index 83dfe1f82..1717be560 100644 --- a/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp +++ b/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp @@ -112,6 +112,11 @@ void gguf_load_quantized( std::string name(tensor.name, tensor.namelen); auto shape = get_shape(tensor); + if (shape.empty()) { + std::ostringstream msg; + msg << "[load_gguf] quantized tensor " << name << " has no dimensions"; + throw std::runtime_error(msg.str()); + } const uint64_t weights_per_block = 32; if (shape[shape.size() - 1] % weights_per_block != 0) { std::ostringstream msg; @@ -125,15 +130,17 @@ void gguf_load_quantized( auto w_nbytes = uint32.size() * std::accumulate(weights_shape.begin(), weights_shape.end(), - 1, + size_t{1}, std::multiplies()); array weights(allocator::malloc(w_nbytes), std::move(weights_shape), uint32); // For scales and bias shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block; - auto sb_nbytes = float16.size() * - std::accumulate(shape.begin(), shape.end(), 1, std::multiplies()); + auto sb_nbytes = + float16.size() * + std::accumulate( + shape.begin(), shape.end(), size_t{1}, std::multiplies()); array scales(allocator::malloc(sb_nbytes), shape, float16); array biases(allocator::malloc(sb_nbytes), std::move(shape), float16); diff --git a/Source/Cmlx/mlx/mlx/io/load.cpp b/Source/Cmlx/mlx/mlx/io/load.cpp index 10a9f72a9..cd34ddc9c 100644 --- a/Source/Cmlx/mlx/mlx/io/load.cpp +++ b/Source/Cmlx/mlx/mlx/io/load.cpp @@ -4,6 +4,7 @@ #include #include #include +#include // Detect iOS for thread pool sizing. #if defined(__APPLE__) @@ -162,6 +163,10 @@ void save(std::shared_ptr out_stream, array a) { a = contiguous(a, true); a.eval(); + // Open only after the inputs are evaluated: opening a FileWriter + // truncates the file and lazy inputs may still read from it. + out_stream->open(); + //////////////////////////////////////////////////////// // Check file if (!out_stream->good() || !out_stream->is_open()) { @@ -348,6 +353,44 @@ array load(std::string file, StreamOrDevice s) { namespace io { +namespace { + +// The reads to keep in flight while reading a file in batches. The disk, not +// the cpu, is the limit here: a thread waits for it most of the time, but a +// queue deeper than the device can serve only adds latency. Half the cores +// saturates the storage of a machine, which is faster on the machines that +// have more cores. +size_t batch_thread_count() { +#if defined(TARGET_OS_IPHONE) && TARGET_OS_IPHONE + // Fork: one reader on iOS, see MLX_IO_THREAD_COUNT above. + return MLX_IO_THREAD_COUNT; +#endif + size_t n_cores = std::thread::hardware_concurrency(); + return std::clamp(n_cores / 2, 4, 16); +} + +// The pool of the batched reads. Every reader shares it, because the disk is +// one resource: a pool per file would multiply the reads in flight by the +// number of open files, which is how a sharded model is read. +// +// It is held with a weak reference, so the threads start with the first +// batched read and stop once the readers that use them are destroyed. Reading +// a file is a startup task, and this keeps no thread for the rest of the run. +std::shared_ptr acquire_batch_pool() { + static std::mutex mtx; + static std::weak_ptr weak_pool; + + std::lock_guard lock(mtx); + auto pool = weak_pool.lock(); + if (!pool) { + pool = std::make_shared(batch_thread_count()); + weak_pool = pool; + } + return pool; +} + +} // namespace + ThreadPool& thread_pool() { // Leak - see Scheduler singleton comment in scheduler.cpp. static ThreadPool* pool_ = new ThreadPool{MLX_IO_THREAD_COUNT}; @@ -355,9 +398,8 @@ ThreadPool& thread_pool() { } ThreadPool& ParallelFileReader::thread_pool() { - // Leak - see Scheduler singleton comment in scheduler.cpp. - static ThreadPool* thread_pool = new ThreadPool{MLX_IO_THREAD_COUNT}; - return *thread_pool; + std::call_once(pool_once_, [this]() { pool_ = acquire_batch_pool(); }); + return *pool_; } void ParallelFileReader::read(char* data, size_t n) { @@ -385,26 +427,22 @@ void ParallelFileReader::read(char* data, size_t n, size_t offset) { } return true; }; + // Read the whole batches in the pool and what is left in this thread. std::vector> futs; - while (n != 0) { - if (n < batch_size_) { - if (!readfn(offset, n, data)) { - throw std::runtime_error("[read] Unable to read from file."); - } - break; - } else { - size_t m = batch_size_; - futs.emplace_back( - ParallelFileReader::thread_pool().enqueue(readfn, offset, m, data)); - data += m; - n -= m; - offset += m; - } + while (n >= batch_size_) { + futs.emplace_back(thread_pool().enqueue(readfn, offset, batch_size_, data)); + data += batch_size_; + offset += batch_size_; + n -= batch_size_; } + bool ok = n == 0 || readfn(offset, n, data); + + // Wait for every batch, they all write into |data|. for (auto& f : futs) { - if (!f.get()) { - throw std::runtime_error("[read] Unable to read from file."); - } + ok = f.get() && ok; + } + if (!ok) { + throw std::runtime_error("[read] Unable to read from file."); } } diff --git a/Source/Cmlx/mlx/mlx/io/load.h b/Source/Cmlx/mlx/mlx/io/load.h index 00c3ff1d4..0a480d981 100644 --- a/Source/Cmlx/mlx/mlx/io/load.h +++ b/Source/Cmlx/mlx/mlx/io/load.h @@ -3,11 +3,13 @@ #pragma once #include +#include #include -#include // SIZE_MAX +#include // SIZE_MAX +#include #include -#ifdef _MSC_VER +#ifdef _WIN32 #include #else #include @@ -66,6 +68,8 @@ class Writer { std::ios_base::seekdir way = std::ios_base::beg) = 0; virtual void write(const char* data, size_t n) = 0; virtual std::string label() const = 0; + virtual void open() {} + virtual ~Writer() = default; }; @@ -88,20 +92,30 @@ class ParallelFileReader : public Reader { } size_t tell() override { +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } // Warning: do not use this function from multiple threads as // it advances the file descriptor void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } // Warning: do not use this function from multiple threads as @@ -115,42 +129,56 @@ class ParallelFileReader : public Reader { } private: - // On iOS, always use sequential pread() — no std::async thread spawning. - // std::async wakeup bursts during model loading exceed iOS's per-process - // wakeup rate limit → EXC_RESOURCE (RESOURCE_TYPE_WAKEUPS) crash. - // On macOS, 32MB parallel batches maximize NVMe throughput. + // Reads larger than this are split in batches and read in parallel. + // Fork: on iOS, always use sequential pread() (no thread spawning). Wakeup + // bursts during model loading exceed iOS's per-process wakeup rate limit + // (EXC_RESOURCE, RESOURCE_TYPE_WAKEUPS). #ifdef MLX_IOS_SEQUENTIAL_IO static constexpr size_t batch_size_ = SIZE_MAX; #else - static constexpr size_t batch_size_ = 1 << 25; // 32MB + static constexpr size_t batch_size_ = 1 << 25; #endif - static ThreadPool& thread_pool(); + + // The pool that reads the batches, held from the first batched read until + // the reader is destroyed. It cannot be io::thread_pool(), the tasks that + // run there wait for these batches and would deadlock in the same pool. + ThreadPool& thread_pool(); + int fd_; std::string label_; + std::once_flag pool_once_; + std::shared_ptr pool_; }; class FileWriter : public Writer { public: explicit FileWriter() {} explicit FileWriter(std::string file_path) - : fd_(open( - file_path.c_str(), - O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, - 0644)), - label_(std::move(file_path)) {} + : file_path_(std::move(file_path)) {} FileWriter(const FileWriter&) = delete; FileWriter& operator=(const FileWriter&) = delete; - FileWriter(FileWriter&& other) { - std::swap(fd_, other.fd_); + FileWriter(FileWriter&& other) + : fd_(std::exchange(other.fd_, -1)), + file_path_(std::move(other.file_path_)) { + other.file_path_.clear(); } ~FileWriter() override { - if (fd_ != 0) { + if (fd_ >= 0) { close(fd_); } } + // Kept separate from construction so lazy inputs can be evaluated first, + // they may still read from the file. + void open() override { + if (fd_ < 0 && !file_path_.empty()) { + fd_ = ::open( + file_path_.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, 0644); + } + } + bool is_open() const override { return fd_ >= 0; } @@ -160,21 +188,34 @@ class FileWriter : public Writer { } size_t tell() override { + check_open(); +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + check_open(); + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } void write(const char* data, size_t n) override { + check_open(); while (n != 0) { auto m = ::write(fd_, data, std::min(n, static_cast(INT32_MAX))); if (m <= 0) { @@ -188,12 +229,18 @@ class FileWriter : public Writer { } std::string label() const override { - return "file " + label_; + return "file " + file_path_; } private: - int fd_{0}; - std::string label_; + void check_open() const { + if (!is_open()) { + throw std::runtime_error("[write] File " + file_path_ + " is not open."); + } + } + + int fd_{-1}; + std::string file_path_; }; } // namespace io diff --git a/Source/Cmlx/mlx/mlx/io/safetensors.cpp b/Source/Cmlx/mlx/mlx/io/safetensors.cpp index 75c8a6ddc..319579c48 100644 --- a/Source/Cmlx/mlx/mlx/io/safetensors.cpp +++ b/Source/Cmlx/mlx/mlx/io/safetensors.cpp @@ -221,6 +221,20 @@ void save_safetensors( std::shared_ptr out_stream, std::unordered_map a, std::unordered_map metadata /* = {} */) { + { + std::vector to_eval; + to_eval.reserve(a.size()); + for (auto& p : a) { + p.second = contiguous(p.second); + to_eval.push_back(p.second); + } + eval(std::move(to_eval)); + } + + // Open only after the inputs are evaluated: opening a FileWriter + // truncates the file and lazy inputs may still read from it. + out_stream->open(); + //////////////////////////////////////////////////////// // Check file if (!out_stream->good() || !out_stream->is_open()) { @@ -238,16 +252,6 @@ void save_safetensors( } parent["__metadata__"] = _metadata; - { - std::vector to_eval; - to_eval.reserve(a.size()); - for (auto& p : a) { - p.second = contiguous(p.second); - to_eval.push_back(p.second); - } - eval(std::move(to_eval)); - } - size_t offset = 0; for (auto& [key, arr] : a) { json child; diff --git a/Source/Cmlx/mlx/mlx/linalg.cpp b/Source/Cmlx/mlx/mlx/linalg.cpp index 0d06b5646..c4213c93b 100644 --- a/Source/Cmlx/mlx/mlx/linalg.cpp +++ b/Source/Cmlx/mlx/mlx/linalg.cpp @@ -4,6 +4,7 @@ #include #include +#include "mlx/backend/cuda/cuda.h" #include "mlx/linalg.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -18,6 +19,18 @@ void check_cpu_stream(const StreamOrDevice& s, const std::string& prefix) { "Explicitly pass a CPU stream to run it."); } } + +// For ops that have a CUDA implementation but no Metal one yet. +void check_cpu_or_cuda_stream( + const StreamOrDevice& s, + const std::string& prefix) { + if (to_stream(s).device == Device::gpu && !cu::is_available()) { + throw std::invalid_argument( + prefix + + " This op is not yet supported on the GPU. " + "Explicitly pass a CPU stream to run it."); + } +} void check_float(Dtype dtype, const std::string& prefix) { if (dtype != float32 && dtype != float64) { std::ostringstream msg; @@ -336,7 +349,7 @@ array cholesky( const array& a, bool upper /* = false */, StreamOrDevice s /* = {} */) { - check_cpu_stream(s, "[linalg::cholesky]"); + check_cpu_or_cuda_stream(s, "[linalg::cholesky]"); check_float(a.dtype(), "[linalg::cholesky]"); if (a.ndim() < 2) { std::ostringstream msg; diff --git a/Source/Cmlx/mlx/mlx/memory.cpp b/Source/Cmlx/mlx/mlx/memory.cpp new file mode 100644 index 000000000..700cc2729 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/memory.cpp @@ -0,0 +1,30 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +#include "mlx/memory.h" + +namespace mlx::core { + +size_t get_array_buffer_size(const std::vector& arrays) { + std::unordered_set buffers; + buffers.reserve(arrays.size()); + + size_t total = 0; + for (const auto& arr : arrays) { + if (arr.status() == array::Status::unscheduled) { + throw std::invalid_argument( + "[get_array_buffer_size] Arrays must be evaluated before querying " + "buffer size."); + } + + const auto& buffer = arr.buffer(); + if (buffer.ptr() != nullptr && buffers.insert(buffer.ptr()).second) { + total += arr.buffer_size(); + } + } + return total; +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/memory.h b/Source/Cmlx/mlx/mlx/memory.h index f4eabc997..14c099f1b 100644 --- a/Source/Cmlx/mlx/mlx/memory.h +++ b/Source/Cmlx/mlx/mlx/memory.h @@ -1,10 +1,12 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #pragma once #include +#include #include "mlx/api.h" +#include "mlx/array.h" namespace mlx::core { @@ -15,6 +17,14 @@ namespace mlx::core { * */ MLX_API size_t get_active_memory(); +/* 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. + * */ +MLX_API size_t get_array_buffer_size(const std::vector& arrays); + /* Get the peak amount of used memory in bytes. * * The maximum memory used recorded from the beginning of the program @@ -42,6 +52,9 @@ MLX_API size_t get_cache_memory(); * When Metal is available the memory limit defaults to 1.5 times the maximum * recommended working set size reported by the device. * + * On native Windows with CUDA the memory limit is also bounded by the current + * video memory budget reported by the operating system. + * * Returns the previous memory limit. * */ MLX_API size_t set_memory_limit(size_t limit); diff --git a/Source/Cmlx/mlx/mlx/ops.cpp b/Source/Cmlx/mlx/mlx/ops.cpp index 5fe3d9be5..83094d25b 100644 --- a/Source/Cmlx/mlx/mlx/ops.cpp +++ b/Source/Cmlx/mlx/mlx/ops.cpp @@ -482,7 +482,7 @@ array unflatten( std::ostringstream msg; msg << "[unflatten] Invalid axes " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } size_t size = 1; @@ -817,7 +817,7 @@ void normalize_dynamic_slice_inputs( std::ostringstream msg; msg << prefix << " Invalid axis " << ax << " for array with dimension " << a.ndim() << "."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } ax = new_ax; } @@ -1465,7 +1465,7 @@ array reflect_pad( Shape starts(a.ndim(), 0); auto stops = a.shape(); for (size_t i = 0; i < axes.size(); i++) { - int ax = axes[i]; + int ax = normalize_axis_index(axes[i], a.ndim(), "[pad] "); starts[ax] = low_pad_size[i]; stops[ax] += low_pad_size[i]; } @@ -1473,7 +1473,7 @@ array reflect_pad( array padded = slice_update(out, a, starts, stops, s); for (size_t i = 0; i < axes.size(); i++) { - int ax = axes[i]; + int ax = normalize_axis_index(axes[i], a.ndim(), "[pad] "); int n = a.shape(ax); int L = low_pad_size[i]; int H = high_pad_size[i]; @@ -1546,37 +1546,41 @@ array edge_pad( const Shape& out_shape, StreamOrDevice s /* = {}*/) { array out = zeros(out_shape, a.dtype(), s); - auto stops = a.shape(); - for (int i = 0; i < stops.size(); i++) { - stops[i] += low_pad_size[i]; + Shape in_starts(a.ndim(), 0); + auto in_stops = a.shape(); + for (size_t i = 0; i < axes.size(); i++) { + int ax = normalize_axis_index(axes[i], a.ndim(), "[pad] "); + in_starts[ax] = low_pad_size[i]; + in_stops[ax] += low_pad_size[i]; } // Copy over values from the unpadded array - array padded = slice_update(out, a, low_pad_size, stops, s); + array padded = slice_update(out, a, in_starts, in_stops, s); - for (int axis = 0; axis < a.ndim(); axis++) { - if (low_pad_size[axis] > 0) { + for (size_t i = 0; i < axes.size(); i++) { + int ax = normalize_axis_index(axes[i], a.ndim(), "[pad] "); + if (low_pad_size[i] > 0) { Shape starts(a.ndim(), 0); - starts[axis] = low_pad_size[axis]; + starts[ax] = low_pad_size[i]; auto stops = out.shape(); - stops[axis] = low_pad_size[axis] + 1; + stops[ax] = low_pad_size[i] + 1; // Fetch edge values array edge_value = slice(padded, starts, stops, s); - starts[axis] = 0; - stops[axis] = low_pad_size[axis]; + starts[ax] = 0; + stops[ax] = low_pad_size[i]; // Update edge values in the padded array padded = slice_update(padded, edge_value, starts, stops, s); } - if (high_pad_size[axis] > 0) { + if (high_pad_size[i] > 0) { Shape starts(a.ndim(), 0); - starts[axis] = -high_pad_size[axis] - 1; + starts[ax] = -high_pad_size[i] - 1; auto stops = out.shape(); - stops[axis] = -high_pad_size[axis]; + stops[ax] = -high_pad_size[i]; array edge_value = slice(padded, starts, stops, s); - starts[axis] = -high_pad_size[axis]; - stops[axis] = out.shape(axis); + starts[ax] = -high_pad_size[i]; + stops[ax] = out.shape(ax); padded = slice_update(padded, edge_value, starts, stops, s); } } @@ -1618,7 +1622,7 @@ array pad( throw std::invalid_argument(msg.str()); } - auto ax = axes[i] < 0 ? a.ndim() + axes[i] : axes[i]; + auto ax = normalize_axis_index(axes[i], a.ndim(), "[pad] "); out_shape[ax] = safe_cast( static_cast(out_shape[ax]) + low_pad_size[i] + high_pad_size[i], @@ -1754,7 +1758,7 @@ array transpose( } if (axes.size() != a.ndim()) { std::ostringstream msg; - msg << "[transpose] Recived " << axes.size() << " axes for array with " + msg << "[transpose] Received " << axes.size() << " axes for array with " << a.ndim() << " dimensions."; throw std::invalid_argument(msg.str()); } @@ -1766,7 +1770,7 @@ array transpose( std::ostringstream msg; msg << "[transpose] Invalid axis (" << ax << ") for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } if (shape[ax] != 0) { throw std::invalid_argument("[transpose] Repeat axes not allowed."); @@ -1790,6 +1794,17 @@ array transpose(const array& a, StreamOrDevice s /* = {} */) { return transpose(a, std::move(axes), to_stream(s)); } +array matrix_transpose(const array& a, StreamOrDevice s /* = {} */) { + if (a.ndim() < 2) { + throw std::invalid_argument( + "[matrix_transpose] Input array must have at least 2 dimensions."); + } + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + std::swap(axes[a.ndim() - 1], axes[a.ndim() - 2]); + return transpose(a, std::move(axes), to_stream(s)); +} + array broadcast_to( const array& a, const Shape& shape, @@ -2329,7 +2344,7 @@ array mean( std::ostringstream msg; msg << "[mean] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } } auto dtype = at_least_float(a.dtype()); @@ -2363,7 +2378,7 @@ array median( std::ostringstream msg; msg << "[median] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } set_axes.insert(axis < 0 ? axis + ndim : axis); } @@ -3150,8 +3165,14 @@ array floor_divide( auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); auto shape = inputs[0].shape(); - return array( + auto quotient = 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); } array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { @@ -3673,7 +3694,7 @@ array take( std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -3718,7 +3739,7 @@ array take(const array& a, int index, int axis, StreamOrDevice s /* = {} */) { std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -4144,14 +4165,7 @@ array cumsum( bool inclusive /* = true*/, std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { - int ndim = a.ndim(); - if (axis >= ndim || axis < -ndim) { - std::ostringstream msg; - msg << "[cumsum] Axis " << axis << " is out of bounds for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - axis = (axis + a.ndim()) % a.ndim(); + axis = normalize_axis_index(axis, a.ndim(), "[cumsum] "); auto x = dtype ? astype(a, *dtype, s) : a; auto out_type = x.dtype() == bool_ ? int32 : x.dtype(); return array( @@ -4179,14 +4193,7 @@ array cumprod( bool inclusive /* = true*/, std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { - int ndim = a.ndim(); - if (axis >= ndim || axis < -ndim) { - std::ostringstream msg; - msg << "[cumprod] Axis " << axis << " is out of bounds for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - axis = (axis + a.ndim()) % a.ndim(); + axis = normalize_axis_index(axis, a.ndim(), "[cumprod] "); auto x = dtype ? astype(a, *dtype, s) : a; return array( x.shape(), @@ -4211,14 +4218,7 @@ array cummax( bool reverse /* = false*/, bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { - int ndim = a.ndim(); - if (axis >= ndim || axis < -ndim) { - std::ostringstream msg; - msg << "[cummax] Axis " << axis << " is out of bounds for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - axis = (axis + a.ndim()) % a.ndim(); + axis = normalize_axis_index(axis, a.ndim(), "[cummax] "); return array( a.shape(), a.dtype(), @@ -4241,14 +4241,7 @@ array cummin( bool reverse /* = false*/, bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { - int ndim = a.ndim(); - if (axis >= ndim || axis < -ndim) { - std::ostringstream msg; - msg << "[cummin] Axis " << axis << " is out of bounds for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - axis = (axis + a.ndim()) % a.ndim(); + axis = normalize_axis_index(axis, a.ndim(), "[cummin] "); return array( a.shape(), a.dtype(), @@ -4273,7 +4266,7 @@ array diff( int ndim = static_cast(a.ndim()); int ax = axis < 0 ? axis + ndim : axis; if (ax < 0 || ax >= ndim) { - throw std::invalid_argument("[diff] Axis is out of bounds for the array."); + throw std::out_of_range("[diff] Axis is out of bounds for the array."); } if (n < 0) { throw std::invalid_argument("[diff] Order `n` must be non-negative."); @@ -4298,14 +4291,7 @@ array logcumsumexp( bool reverse /* = false*/, bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { - int ndim = a.ndim(); - if (axis >= ndim || axis < -ndim) { - std::ostringstream msg; - msg << "[logcumsumexp] Axis " << axis << " is out of bounds for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - axis = (axis + a.ndim()) % a.ndim(); + axis = normalize_axis_index(axis, a.ndim(), "[logcumsumexp] "); return array( a.shape(), a.dtype(), @@ -5567,9 +5553,14 @@ array gather_qmm( std::optional group_size_ /* = std::nullopt */, std::optional bits_ /* = std::nullopt */, const std::string& mode /* = "affine" */, + const std::optional& global_scale /* = std::nullopt */, bool sorted_indices /* = false */, StreamOrDevice s /* = {} */) { if (!lhs_indices_ && !rhs_indices_) { + if (global_scale) { + throw std::invalid_argument( + "[gather_qmm] Global scale is not supported without indices."); + } return quantized_matmul( x, w, scales, biases, transpose, group_size_, bits_, mode, s); } @@ -5580,6 +5571,32 @@ array gather_qmm( quantization_params_from_mode(qmode, group_size_, bits_); auto [w_inner_dims, w_outer_dims] = extract_quantized_matmul_dims( "gather_qmm", x, w, scales, biases, transpose, group_size, bits); + if (global_scale) { + if (qmode != QuantizationMode::Nvfp4) { + throw std::invalid_argument( + "[gather_qmm] Global scale is only supported for 'nvfp4' " + "quantization mode."); + } + if (global_scale->dtype() != float32) { + std::ostringstream msg; + msg << "[gather_qmm] Global scale must have dtype float32 but got " + << global_scale->dtype() << "."; + throw std::invalid_argument(msg.str()); + } + // One scale per expert, so it matches the batch dimensions of w. + Shape expected(w.shape().begin(), w.shape().end() - 2); + if (global_scale->shape() != expected) { + std::ostringstream msg; + msg << "[gather_qmm] Global scale must have one entry per expert with " + << "shape " << expected << " but got " << global_scale->shape() + << "."; + throw std::invalid_argument(msg.str()); + } + if (to_stream(s).device != Device::gpu) { + throw std::invalid_argument( + "[gather_qmm] Global scale is only supported on the GPU."); + } + } if (qmode == QuantizationMode::Affine) { out_type = promote_types(x.dtype(), out_type); } else { @@ -5613,12 +5630,12 @@ array gather_qmm( std::move(lhs_indices), std::move(rhs_indices)}; } else { - inputs = { - astype(x, out_type, s), - std::move(w), - std::move(scales), - std::move(lhs_indices), - std::move(rhs_indices)}; + inputs = {astype(x, out_type, s), std::move(w), std::move(scales)}; + if (global_scale) { + inputs.push_back(*global_scale); + } + inputs.push_back(std::move(lhs_indices)); + inputs.push_back(std::move(rhs_indices)); } return array( std::move(out_shape), @@ -5680,7 +5697,13 @@ array gather_qqmm( return array( std::move(out_shape), x.dtype(), - std::make_shared(stream, group_size, bits, qmode), + std::make_shared( + stream, + group_size, + bits, + qmode, + sorted_indices && !rhs_indices_, + sorted_indices && !lhs_indices_), std::move(inputs)); } @@ -5796,7 +5819,7 @@ array vecdot( } int ax = axis < 0 ? axis + a.ndim() : axis; if (ax < 0 || ax >= a.ndim()) { - throw std::invalid_argument("[vecdot] axis is out of bounds."); + throw std::out_of_range("[vecdot] axis is out of bounds."); } if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) { throw std::invalid_argument("[vecdot] axis is out of bounds."); @@ -6677,7 +6700,7 @@ array roll( std::ostringstream msg; msg << "[roll] Invalid axis " << axes[i] << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } auto sh = shift[i]; diff --git a/Source/Cmlx/mlx/mlx/ops.h b/Source/Cmlx/mlx/mlx/ops.h index f597753b1..ab676ca6c 100644 --- a/Source/Cmlx/mlx/mlx/ops.h +++ b/Source/Cmlx/mlx/mlx/ops.h @@ -381,6 +381,9 @@ inline array transpose( return transpose(a, std::vector(axes), s); } +/** Permutes last two dimensions of an array. */ +MLX_API array matrix_transpose(const array& a, StreamOrDevice s = {}); + /** Swap two axes of an array. */ MLX_API array swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {}); @@ -1626,6 +1629,7 @@ MLX_API array gather_qmm( std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "affine", + const std::optional& global_scale = std::nullopt, bool sorted_indices = false, StreamOrDevice s = {}); diff --git a/Source/Cmlx/mlx/mlx/primitives.cpp b/Source/Cmlx/mlx/mlx/primitives.cpp index 9a1771394..f046670ee 100644 --- a/Source/Cmlx/mlx/mlx/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/primitives.cpp @@ -3335,9 +3335,10 @@ std::vector Pad::vjp( Shape start(cotan.ndim(), 0); auto stop = cotan.shape(); - for (auto i : axes_) { - start[i] = low_pad_size_[i]; - stop[i] -= high_pad_size_[i]; + for (size_t i = 0; i < axes_.size(); i++) { + auto ax = normalize_axis_index(axes_[i], cotan.ndim(), "[pad] "); + start[ax] = low_pad_size_[i]; + stop[ax] -= high_pad_size_[i]; } auto out = slice(cotan, start, stop, stream()); @@ -3732,9 +3733,13 @@ std::vector GatherQMM::vjp( auto biases = (mode_ == QuantizationMode::Affine) ? std::optional(primals[3]) : std::nullopt; + auto global_scale = (mode_ != QuantizationMode::Affine && primals.size() == 6) + ? std::optional(primals[3]) + : std::nullopt; int M = cotan.shape(-2); int K = x.shape(-1); + int first_index_arg = primals.size() - 2; bool sorted = left_sorted_ || right_sorted_; bool no_broadcast = rhs_indices.size() * M * K == x.size(); @@ -3754,6 +3759,7 @@ std::vector GatherQMM::vjp( group_size_, bits_, quantization_mode_to_string(mode_), + global_scale, sorted, stream()); if (sorted && no_broadcast) { @@ -3772,15 +3778,18 @@ std::vector GatherQMM::vjp( } // gradient wrt to the indices is undefined - else if (arg > 3) { + else if (arg >= first_index_arg) { throw std::runtime_error( "[GatherQMM::vjp] cannot compute the gradient wrt the indices."); } - // gradient wrt to w_q, scales or biases + // gradient wrt to w_q, scales, biases or the global scale else if (arg == 1) { throw std::runtime_error( "[GatherQMM::vjp] no gradient wrt the quantized weights."); + } else if (global_scale && arg == 3) { + throw std::runtime_error( + "[GatherQMM::vjp] no gradient wrt the global scale."); } else { if (mode_ != QuantizationMode::Affine) { std::ostringstream msg; @@ -3852,8 +3861,7 @@ bool GatherQMM::is_equivalent(const Primitive& other) const { std::vector GatherQMM::output_shapes(const std::vector& inputs) { const auto& x = inputs[0]; const auto& w = inputs[1]; - const auto& lhs_indices = - (mode_ == QuantizationMode::Affine) ? inputs[4] : inputs[3]; + const auto& lhs_indices = inputs[inputs.size() - 2]; int w_outer = transpose_ ? w.shape(-2) : w.shape(-1) * 32 / bits_; auto out_shape = lhs_indices.shape(); out_shape.push_back(x.shape(-2)); diff --git a/Source/Cmlx/mlx/mlx/primitives.h b/Source/Cmlx/mlx/mlx/primitives.h index 0cfc71bf0..f5ffc7d86 100644 --- a/Source/Cmlx/mlx/mlx/primitives.h +++ b/Source/Cmlx/mlx/mlx/primitives.h @@ -1888,6 +1888,7 @@ class Scan : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS(); + DEFINE_INPUT_OUTPUT_SHAPE() const char* name() const override { switch (reduce_type_) { diff --git a/Source/Cmlx/mlx/mlx/random.cpp b/Source/Cmlx/mlx/mlx/random.cpp index 2932b3eaf..9a49853ff 100644 --- a/Source/Cmlx/mlx/mlx/random.cpp +++ b/Source/Cmlx/mlx/mlx/random.cpp @@ -383,7 +383,7 @@ int get_valid_axis(int axis, int ndim) { std::ostringstream msg; msg << "[categorical] Invalid axis " << axis << " for logits with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } return ax; } diff --git a/Source/Cmlx/mlx/mlx/stream.cpp b/Source/Cmlx/mlx/mlx/stream.cpp index b78ee67d6..902015cbf 100644 --- a/Source/Cmlx/mlx/mlx/stream.cpp +++ b/Source/Cmlx/mlx/mlx/stream.cpp @@ -50,6 +50,10 @@ Stream default_stream(Device d) { return s.value(); } +std::optional peek_default_stream(Device d) { + return default_stream_storage(d); +} + void set_default_stream(Stream s) { if (!gpu::is_available() && s.device == Device::gpu) { throw std::invalid_argument( diff --git a/Source/Cmlx/mlx/mlx/stream.h b/Source/Cmlx/mlx/mlx/stream.h index 3099ee268..331af60a6 100644 --- a/Source/Cmlx/mlx/mlx/stream.h +++ b/Source/Cmlx/mlx/mlx/stream.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include @@ -29,6 +30,9 @@ struct MLX_API ThreadLocalStream : public Stream { /** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); +/** Get the default stream of current thread if it exists. */ +MLX_API std::optional peek_default_stream(Device d); + /** Make the stream the default for its device on current thread. */ MLX_API void set_default_stream(Stream s); diff --git a/Source/Cmlx/mlx/mlx/transforms.cpp b/Source/Cmlx/mlx/mlx/transforms.cpp index 694bca53d..77d81e11e 100644 --- a/Source/Cmlx/mlx/mlx/transforms.cpp +++ b/Source/Cmlx/mlx/mlx/transforms.cpp @@ -80,17 +80,23 @@ thread_local int detail::RetainGraph::tracing_counter{0}; array eval_impl(std::vector outputs, bool async) { std::deque tape; - // Make an effort to choose a good output stream - Stream stream = default_stream(default_device()); - for (auto& o : outputs) { - if (o.status() == array::Status::unscheduled && o.has_primitive()) { - stream = o.primitive().stream(); - break; + // Make an effort to choose a good output stream, and only create the default + // stream when there is no other choice. + Stream stream = [&outputs]() { + for (auto& o : outputs) { + if (o.status() == array::Status::unscheduled && o.has_primitive()) { + return o.primitive().stream(); + } } - } + return default_stream(default_device()); + }(); - // Map of array id that needs fence and stream it's computed on - std::unordered_map> needs_fence; + struct FenceInfo { + int stream_index; + bool cross_device; + uint32_t value{0}; + }; + std::unordered_map needs_fence; auto synchronizer = array( {}, bool_, std::make_shared(stream), std::move(outputs)); @@ -146,9 +152,9 @@ array eval_impl(std::vector outputs, bool async) { a.primitive().stream().device != in.primitive().stream().device; auto [it, inserted] = needs_fence.emplace( in.id(), - std::make_pair(in.primitive().stream().index, device_switch)); + FenceInfo{in.primitive().stream().index, device_switch}); if (!inserted) { - it->second.second |= device_switch; + it->second.cross_device |= device_switch; } } } @@ -251,7 +257,8 @@ array eval_impl(std::vector outputs, bool async) { // Use fence to wait within a single eval // Get the input array's stream fence and wait on the // output arrays stream - fences[it->second.first].wait(stream, in); + auto& info = it->second; + fences.at(info.stream_index).wait(stream, in, info.value); } else if (in.event().valid()) { if (in.event().is_signaled()) { in.detach_event(); @@ -291,7 +298,8 @@ array eval_impl(std::vector outputs, bool async) { if (it == fences.end()) { it = fences.emplace(stream.index, Fence{stream}).first; } - it->second.update(stream, a, nf->second.second); + nf->second.value = + it->second.update(stream, a, nf->second.cross_device); } }; diff --git a/Source/Cmlx/mlx/mlx/types/complex.h b/Source/Cmlx/mlx/mlx/types/complex.h index 51101cc97..65370ca89 100644 --- a/Source/Cmlx/mlx/mlx/types/complex.h +++ b/Source/Cmlx/mlx/mlx/types/complex.h @@ -106,8 +106,23 @@ inline complex64_t operator-(const complex64_t& v) { complex_binop_helper(_op_, _operator_, float16_t) \ complex_binop_helper(_op_, _operator_, bfloat16_t) \ complex_binop_helper(_op_, _operator_, float) -// clang-format on complex_binop(+, operator+) +template +constexpr bool is_complex = false; + +template +constexpr bool is_complex().real())>> = true; +// clang-format on + +template +inline bool isnan(T v) { + if constexpr (is_complex) { + return std::isnan(std::real(v)) || std::isnan(std::imag(v)); + } else { + return std::isnan(v); + } +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/types/half_types.h b/Source/Cmlx/mlx/mlx/types/half_types.h index 560f73735..aa1e56690 100644 --- a/Source/Cmlx/mlx/mlx/types/half_types.h +++ b/Source/Cmlx/mlx/mlx/types/half_types.h @@ -56,13 +56,6 @@ fp16_bf16_binop_helper(/, operator/) #endif -template -constexpr bool is_complex = false; - -template -constexpr bool is_complex().real())>> = - true; - template constexpr bool is_signed_v = std::is_signed_v || std::is_same_v || std::is_same_v; diff --git a/Source/Cmlx/mlx/mlx/utils.cpp b/Source/Cmlx/mlx/mlx/utils.cpp index 4e29e8847..07130fa69 100644 --- a/Source/Cmlx/mlx/mlx/utils.cpp +++ b/Source/Cmlx/mlx/mlx/utils.cpp @@ -179,7 +179,7 @@ int normalize_axis_index( std::ostringstream msg; msg << msg_prefix << "Axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } return axis < 0 ? axis + ndim : axis; } diff --git a/Source/Cmlx/mlx/mlx/utils.h b/Source/Cmlx/mlx/mlx/utils.h index 0c62ba30d..4b7e6d12d 100644 --- a/Source/Cmlx/mlx/mlx/utils.h +++ b/Source/Cmlx/mlx/mlx/utils.h @@ -7,6 +7,7 @@ #include #include #include +#include #include #include "mlx/api.h" @@ -28,7 +29,9 @@ MLX_API Stream to_stream(StreamOrDevice s, Device default_); struct StreamContext { public: - StreamContext(StreamOrDevice s) : _stream(default_stream(default_device())) { + StreamContext(StreamOrDevice s) + : _stream(default_stream(default_device())), + _owner(std::this_thread::get_id()) { if (std::holds_alternative(s)) { throw std::runtime_error( "[StreamContext] Invalid argument, please specify a stream or device."); @@ -38,13 +41,21 @@ struct StreamContext { set_default_stream(_s); } - ~StreamContext() { + // noexcept(false) so the thread check below can throw. A destructor that + // throws while another exception unwinds this thread calls std::terminate. + ~StreamContext() noexcept(false) { + if (std::this_thread::get_id() != _owner) { + throw std::runtime_error( + "[StreamContext] Destroyed in a different thread than it was " + "constructed in."); + } set_default_device(_stream.device); set_default_stream(_stream); } private: Stream _stream; + std::thread::id _owner; }; struct MLX_API PrintOptions { diff --git a/Source/Cmlx/mlx/mlx/version.h b/Source/Cmlx/mlx/mlx/version.h index 9de9f2e55..a403f6f8b 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 2 +#define MLX_VERSION_PATCH 3 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) 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 49f300d4b..52a6733f2 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py @@ -127,14 +127,34 @@ def __call__(self, x): class ALiBi(Module): + """Implements Attention with Linear Biases (ALiBi). + + 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 `_. + """ + @staticmethod def create_alibi_matrix( q_sequence_length: int, k_sequence_length: int, num_heads: int, - offset: int, - dtype=mx.float32, - ): + offset: int = 0, + dtype: mx.Dtype = mx.float32, + ) -> mx.array: + """Create the ALiBi bias matrix. + + Args: + q_sequence_length (int): The query sequence length. + 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``. + + Returns: + array: The ALiBi bias matrix. + """ x1 = mx.arange(offset, q_sequence_length) x2 = mx.arange(0, k_sequence_length) distance_matrix = -mx.abs( @@ -145,7 +165,17 @@ def create_alibi_matrix( return alibi_mask @staticmethod - def create_alibi_slope(num_heads, dtype): + def create_alibi_slope(num_heads: int, dtype: mx.Dtype = mx.float32) -> mx.array: + """Create the geometric slopes for ALiBi across attention heads. + + Args: + num_heads (int): The number of attention heads. + dtype (Dtype, optional): Data type of the output array. Default: ``mx.float32``. + + Returns: + array: The slopes array expanded for head broadcasting. + """ + def get_slopes(n: int): if math.log2(n).is_integer(): start = 2 ** (-(2 ** -(math.log2(n) - 3))) @@ -161,7 +191,22 @@ def get_slopes(n: int): out = mx.array(slopes, dtype=dtype) return mx.expand_dims(out, axis=(-1, -2)) - def __call__(self, attention_scores, offset=0, mask=None): + def __call__( + self, + attention_scores: mx.array, + offset: int = 0, + mask: Optional[mx.array] = None, + ) -> mx.array: + """Apply ALiBi matrix to attention scores. + + Args: + attention_scores (array): The input attention scores. + offset (int, optional): The position offset. Default: ``0``. + mask (array, optional): An optional attention mask. Default: ``None``. + + Returns: + array: The updated attention scores with ALiBi bias added. + """ alibi_mask = ALiBi.create_alibi_matrix( q_sequence_length=attention_scores.shape[-2] + offset, k_sequence_length=attention_scores.shape[-1], diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py index 0be9f87bf..b04f6e9ba 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py @@ -302,6 +302,13 @@ def upsample_linear( indices_fn=_linear_aa_indices, align_corners=align_corners, ) + if not antialias and not align_corners: + return _interpolate_separable( + x=x, + scale_factor=scale_factor, + indices_fn=_linear_indices, + align_corners=align_corners, + ) return _interpolate( x=x, scale_factor=scale_factor, diff --git a/Source/Cmlx/mlx/python/mlx/nn/losses.py b/Source/Cmlx/mlx/python/mlx/nn/losses.py index b98d2765d..b72b91816 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/losses.py +++ b/Source/Cmlx/mlx/python/mlx/nn/losses.py @@ -64,15 +64,15 @@ def cross_entropy( >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) >>> - >>> # Half precision logits with class indices as targets. On CUDA a + >>> # Half precision logits with class indices as targets. On the GPU a >>> # fused kernel accumulates the reduction in float32: >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16) >>> targets = mx.array([0, 1]) >>> nn.losses.cross_entropy(logits, targets) - array([0.0485873, 0.0485873], dtype=float32) + array([0.048584, 0.048584], dtype=bfloat16) >>> - >>> # Metal and the CPU reduce in the dtype of the logits, so upcast - >>> # them to get the same accuracy: + >>> # The CPU reduces in the dtype of the logits, so upcast them to get + >>> # the same accuracy: >>> nn.losses.cross_entropy(logits.astype(mx.float32), targets) array([0.0485873, 0.0485873], dtype=float32) """ @@ -96,8 +96,7 @@ def _drop_dim(shape, axis): ) use_fast = ( - mx.cuda.is_available() - and mx.default_device() == mx.gpu + mx.default_device() == mx.gpu and not targets_as_probs and label_smoothing == 0 and axis in (-1, logits.ndim - 1) diff --git a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py index 8be344247..3966f961f 100644 --- a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py +++ b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py @@ -274,7 +274,7 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): optimizer state.""" if self.weight_decay != 0: - gradient += self.weight_decay * parameter + gradient = gradient + self.weight_decay * parameter if self.momentum <= 0: return parameter - self.learning_rate.astype(gradient.dtype) * gradient @@ -856,7 +856,7 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): update = exp_avg if self.weight_decay != 0: - parameter += parameter * (-self.weight_decay * learning_rate) + parameter = parameter + parameter * (-self.weight_decay * learning_rate) return parameter - update diff --git a/Source/Cmlx/mlx/python/src/array.cpp b/Source/Cmlx/mlx/python/src/array.cpp index fbab84618..0f13293a1 100644 --- a/Source/Cmlx/mlx/python/src/array.cpp +++ b/Source/Cmlx/mlx/python/src/array.cpp @@ -76,6 +76,9 @@ class ArrayAt { class ArrayPythonIterator { public: ArrayPythonIterator(mx::array x) : idx_(0), x_(std::move(x)) { + if (x_.ndim() == 0) { + throw nb::type_error("iter() 0-dimensional array."); + } if (x_.shape(0) > 0 && x_.shape(0) < 10) { splits_ = mx::split(x_, x_.shape(0)); } @@ -1034,6 +1037,23 @@ void init_array(nb::module_& m) { [](mx::array& a) { return nb::cast>(to_scalar(a)); }) + .def( + "__index__", + [](mx::array& a) { + if (!mx::issubdtype(a.dtype(), mx::integer) || a.ndim() != 0) { + throw nb::type_error( + "Only 0-dimensional integer arrays can be converted to an index."); + } + return nb::int_(to_scalar(a)); + }) + .def( + "__bytes__", + [](mx::array& a) { + auto c = mx::contiguous(a); + c.eval(); + return nb::bytes( + reinterpret_cast(c.data()), c.nbytes()); + }) .def( "__format__", [](mx::array& a, nb::object format_spec) { @@ -1248,6 +1268,10 @@ void init_array(nb::module_& m) { "T", [](const mx::array& a) { return mx::transpose(a); }, "Equivalent to calling ``self.transpose()`` with no arguments.") + .def_prop_ro( + "mT", + [](const mx::array& a) { return mx::matrix_transpose(a); }, + "Equivalent to calling ``self.transpose()`` with the last two axes swapped.") .def( "sum", [](const mx::array& a, @@ -1352,7 +1376,9 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::std( a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, @@ -1360,6 +1386,7 @@ void init_array(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), "See :func:`std`.") .def( @@ -1368,7 +1395,9 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::var( a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, @@ -1376,6 +1405,7 @@ void init_array(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), "See :func:`var`.") .def( diff --git a/Source/Cmlx/mlx/python/src/convert.cpp b/Source/Cmlx/mlx/python/src/convert.cpp index a3da76ef5..ad1bc61a5 100644 --- a/Source/Cmlx/mlx/python/src/convert.cpp +++ b/Source/Cmlx/mlx/python/src/convert.cpp @@ -250,7 +250,10 @@ mx::array metal_nd_array_to_mlx( if (copy) { auto result = mx::astype(out, dst_dtype, true, mx::Device::gpu); - result.eval(); + { + nb::gil_scoped_release nogil; + result.eval(); + } return result; } return out; diff --git a/Source/Cmlx/mlx/python/src/convert.h b/Source/Cmlx/mlx/python/src/convert.h index 133f443ed..7c767a594 100644 --- a/Source/Cmlx/mlx/python/src/convert.h +++ b/Source/Cmlx/mlx/python/src/convert.h @@ -14,24 +14,6 @@ namespace nb = nanobind; namespace nanobind { -template <> -struct ndarray_traits { - static constexpr bool is_complex = false; - static constexpr bool is_float = true; - static constexpr bool is_bool = false; - static constexpr bool is_int = false; - static constexpr bool is_signed = true; -}; - -template <> -struct ndarray_traits { - static constexpr bool is_complex = false; - static constexpr bool is_float = true; - static constexpr bool is_bool = false; - static constexpr bool is_int = false; - static constexpr bool is_signed = true; -}; - namespace detail { template <> diff --git a/Source/Cmlx/mlx/python/src/fast.cpp b/Source/Cmlx/mlx/python/src/fast.cpp index 67c3442cf..f9c472ab9 100644 --- a/Source/Cmlx/mlx/python/src/fast.cpp +++ b/Source/Cmlx/mlx/python/src/fast.cpp @@ -189,8 +189,8 @@ void init_fast(nb::module_& parent_module) { Computes ``logsumexp(logits, axis=-1) - logits[..., target]`` in a fused kernel with accumulation in float32. - Note: Currently is implemented only on CUDA, fallback to unfused version with - manual casting on Metal and CPU. + Note: The fused kernel is available on Metal and CUDA. The CPU falls + back to the unfused version, which reduces in the dtype of the logits. Args: logits (array): The unnormalized logits. The loss is computed over @@ -386,6 +386,32 @@ void init_fast(nb::module_& parent_module) { out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal") )pbdoc"); + m.def( + "gated_delta_update", + &mlx::core::fast::gated_delta_update, + "q"_a, + "k"_a, + "v"_a, + "gamma"_a, + "beta"_a, + "initial_state"_a = nb::none(), // optional, defaults to None + "mask"_a = nb::none(), // optional, defaults to None + "stream"_a = nb::none(), // optional, defaults to None + R"( + Chunked gated delta network forward pass. + + Args: + q: Queries [B, T, Hk, Dk] + k: Keys [B, T, Hk, Dk] + v: Values [B, T, Hv, Dv] + gamma: Decay rate in linear space [B, T, Hv] + beta: Delta update rates [B, T, Hv] + initial_state: Optional initial hidden state [B, T, Hk, Dv] + mask: Optional + Returns: + Tuple of (output [B, T, Hv, Dv], final_state [B, T, Hk, Dv]) + )"); + m.def( "metal_kernel", [](const std::string& name, diff --git a/Source/Cmlx/mlx/python/src/indexing.cpp b/Source/Cmlx/mlx/python/src/indexing.cpp index 1dce53788..5bf467750 100644 --- a/Source/Cmlx/mlx/python/src/indexing.cpp +++ b/Source/Cmlx/mlx/python/src/indexing.cpp @@ -5,6 +5,7 @@ #include +#include "mlx/dtype.h" #include "mlx/ops.h" #include "python/src/convert.h" #include "python/src/indexing.h" @@ -20,6 +21,9 @@ bool is_index_scalar(const nb::object& obj) { if (nb::isinstance(obj)) { return false; } + if (nb::isinstance(obj)) { + return false; + } if (!PyIndex_Check(obj.ptr())) { return false; } @@ -253,8 +257,14 @@ auto mlx_expand_ellipsis(const mx::Shape& shape, const nb::tuple& entries) { // Expand ellipsis if (has_ellipsis) { + int ndim = static_cast(shape.size()); + if (non_none_indices > ndim) { + std::ostringstream msg; + msg << "Too many indices for array with " << ndim << " dimensions."; + throw std::invalid_argument(msg.str()); + } for (int axis = non_none_indices_before; - axis < shape.size() - non_none_indices_after; + axis < ndim - non_none_indices_after; axis++) { indices.push_back( nb::slice(mx::ShapeElem{0}, shape[axis], mx::ShapeElem{1})); @@ -907,7 +917,7 @@ mlx_compute_slice_update_args( upd_ax--; } else if (is_index_scalar(pyidx)) { int st = safe_to_int32(pyidx); - st = (st < 0) ? st + src.shape(i) : st; + st = (st < 0) ? st + src.shape(ax) : st; starts[ax] = st; stops[ax] = st + 1; if (upd_ax >= 0) { diff --git a/Source/Cmlx/mlx/python/src/memory.cpp b/Source/Cmlx/mlx/python/src/memory.cpp index 5ce9a765b..529995539 100644 --- a/Source/Cmlx/mlx/python/src/memory.cpp +++ b/Source/Cmlx/mlx/python/src/memory.cpp @@ -1,8 +1,10 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/memory.h" #include +#include "python/src/trees.h" + namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; @@ -17,6 +19,30 @@ void init_memory(nb::module_& m) { Note, this will not always match memory use reported by the system because it does not include cached memory buffers. )pbdoc"); + m.def( + "get_array_buffer_size", + [](const nb::args& args) { + std::vector arrays = tree_flatten(args, false); + nb::gil_scoped_release nogil; + return mx::get_array_buffer_size(arrays); + }, + nb::arg(), + nb::sig("def get_array_buffer_size(*args) -> int"), + R"pbdoc( + Get the size of the buffers backing 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. Arrays + must be evaluated before calling this function. The result does not + include graph objects, cached buffers, or other non-buffer memory. + + Args: + *args (arrays or trees of arrays): Each argument can be a single array + or a tree of arrays. Leaves which are not arrays are ignored. + + Returns: + int: The size of the unique buffers in bytes. + )pbdoc"); m.def( "get_peak_memory", &mx::get_peak_memory, @@ -62,6 +88,12 @@ void init_memory(nb::module_& m) { Returns: int: The previous memory limit in bytes. )pbdoc"); + m.def( + "get_memory_limit", + &mx::get_memory_limit, + R"pbdoc( + Get the current memory limit in bytes. + )pbdoc"); m.def( "set_cache_limit", &mx::set_cache_limit, diff --git a/Source/Cmlx/mlx/python/src/ops.cpp b/Source/Cmlx/mlx/python/src/ops.cpp index 677d998af..fe63bbf71 100644 --- a/Source/Cmlx/mlx/python/src/ops.cpp +++ b/Source/Cmlx/mlx/python/src/ops.cpp @@ -213,14 +213,14 @@ void init_ops(nb::module_& m) { m.def( "unstack", [](const mx::array& a, int axis, mx::StreamOrDevice s) { - return mx::unstack(a, axis, s); + return nb::tuple(nb::cast(mx::unstack(a, axis, s))); }, nb::arg(), nb::kw_only(), "axis"_a = 0, "stream"_a = nb::none(), nb::sig( - "def unstack(x: array, /, *, axis: int = 0, stream: StreamOrDevice = None) -> list[array]"), + "def unstack(x: array, /, *, axis: int = 0, stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Split an array into a sequence of arrays along the given axis. @@ -232,7 +232,7 @@ void init_ops(nb::module_& m) { axis (int, optional): Axis along which to unstack. Default: ``0``. Returns: - list(array): A list of arrays, one for each index along ``axis``. + tuple(array): A tuple of arrays, one for each index along ``axis``. )pbdoc"); m.def( "expand_dims", @@ -1768,11 +1768,11 @@ void init_ops(nb::module_& m) { }, nb::arg(), "indices"_a, - "axis"_a = nb::none(), + "axis"_a.none() = -1, nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def take_along_axis(a: array, /, indices: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), + "def take_along_axis(a: array, /, indices: array, axis: int | None = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Take values along an axis at the specified indices. @@ -1782,7 +1782,7 @@ void init_ops(nb::module_& m) { the input array excluding the `axis` dimension. axis (int or None): Axis in the input to take the values from. If ``axis == None`` the array is flattened to 1D prior to the indexing - operation. + operation. Defaults to ``-1``. Returns: array: The output array. @@ -1807,7 +1807,7 @@ void init_ops(nb::module_& m) { nb::arg(), "indices"_a, "values"_a, - "axis"_a.none(), + "axis"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( @@ -2106,7 +2106,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tri(n: int, m: int, k: int, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), + "def tri(n: int, m: int | None = None, k: int = 0, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array with ones at and below the given diagonal and zeros elsewhere. @@ -2128,7 +2128,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tril(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), + "def tril(x: array, k: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array above the given diagonal. @@ -2148,7 +2148,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def triu(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), + "def triu(x: array, k: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array below the given diagonal. @@ -2549,6 +2549,16 @@ void init_ops(nb::module_& m) { Returns: array: The transposed array. )pbdoc"); + m.def( + "matrix_transpose", + &mx::matrix_transpose, + nb::arg(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def matrix_transpose(a: array, /, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Transpose the last two dimensions of the array.)pbdoc"); m.def( "permute_dims", [](const mx::array& a, @@ -2856,7 +2866,9 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), @@ -2864,6 +2876,7 @@ void init_ops(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), nb::sig( "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), @@ -2889,7 +2902,9 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), @@ -2897,6 +2912,7 @@ void init_ops(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), nb::sig( "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), @@ -4823,11 +4839,12 @@ void init_ops(nb::module_& m) { "group_size"_a = nb::none(), "bits"_a = nb::none(), "mode"_a = "affine", + "global_scale"_a = nb::none(), nb::kw_only(), "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_qmm(x: array, w: array, /, scales: array, biases: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), + "def gather_qmm(x: array, w: array, /, scales: array, biases: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', global_scale: array | None = None, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform quantized matrix multiplication with matrix-level gather. @@ -4857,6 +4874,9 @@ void init_ops(nb::module_& m) { ``w`` in the quantized array. See supported values and defaults in the :ref:`table of quantization modes `. Default: ``None``. mode (str, optional): The quantization mode. Default: ``"affine"``. + global_scale (array, optional): The per-input float32 scale used for + ``nvfp4`` quantization of ``w``. Only supported on the GPU. + Default: ``None``. sorted_indices (bool, optional): May allow a faster implementation if the passed indices are sorted. Default: ``False``. @@ -5143,7 +5163,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), + "def gather_mm(a: array, b: array, /, lhs_indices: array | None = None, rhs_indices: array | None = None, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with matrix-level gather. @@ -6176,6 +6196,5 @@ void init_ops(nb::module_& m) { m.attr("cumulative_sum") = m.attr("cumsum"); m.attr("empty") = m.attr("zeros"); m.attr("empty_like") = m.attr("zeros_like"); - m.attr("matrix_transpose") = m.attr("transpose"); m.attr("pow") = m.attr("power"); } diff --git a/Source/Cmlx/mlx/python/src/small_vector.h b/Source/Cmlx/mlx/python/src/small_vector.h index 0b2e6bb2e..1a07abe1b 100644 --- a/Source/Cmlx/mlx/python/src/small_vector.h +++ b/Source/Cmlx/mlx/python/src/small_vector.h @@ -31,7 +31,7 @@ struct type_caster<::mlx::core::SmallVector> { // Not noexcept: on overflow of a narrow integer element we raise // OverflowError so nanobind surfaces a clean error to the user. - bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) { + bool from_python(handle src, uint32_t flags, cleanup_list* cleanup) { size_t size; PyObject* temp; diff --git a/Source/Cmlx/mlx/python/src/stream.cpp b/Source/Cmlx/mlx/python/src/stream.cpp index 467518e99..2ddb99179 100644 --- a/Source/Cmlx/mlx/python/src/stream.cpp +++ b/Source/Cmlx/mlx/python/src/stream.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include +#include #include #include @@ -15,10 +16,11 @@ namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; -// Create the StreamContext on enter and delete on exit. +// Create a StreamContext on enter and delete it on exit. The contexts are +// stacked, so the same object can be entered more than once. class PyStreamContext { public: - PyStreamContext(mx::StreamOrDevice s) : _inner(nullptr) { + PyStreamContext(mx::StreamOrDevice s) { if (std::holds_alternative(s)) { throw std::runtime_error( "[StreamContext] Invalid argument, please specify a stream or device."); @@ -27,19 +29,33 @@ class PyStreamContext { } void enter() { - _inner = new mx::StreamContext(_s); + _contexts.push_back(new mx::StreamContext(_s)); } void exit() { - if (_inner != nullptr) { - delete _inner; - _inner = nullptr; + if (_contexts.empty()) { + return; + } + auto* inner = _contexts.back(); + _contexts.pop_back(); + // the destructor can throw, inner is freed regardless. + delete inner; + } + + // ~StreamContext throws when destroyed on a different thread, which Python + // cannot control + ~PyStreamContext() { + while (!_contexts.empty()) { + try { + exit(); + } catch (...) { + } } } private: mx::StreamOrDevice _s; - mx::StreamContext* _inner; + std::vector _contexts; }; void init_stream(nb::module_& m) { @@ -140,6 +156,7 @@ void init_stream(nb::module_& m) { "clear_streams", []() { reset_random_state(); + nb::gil_scoped_release nogil; mx::clear_streams(); }, R"pbdoc(Destroy all streams created in current thread.)pbdoc"); @@ -190,6 +207,7 @@ void init_stream(nb::module_& m) { m.def( "synchronize", [](mx::StreamOrDevice s) { + nb::gil_scoped_release nogil; if (std::holds_alternative(s)) { mx::synchronize(); } else { diff --git a/Source/Cmlx/mlx/python/src/transforms.cpp b/Source/Cmlx/mlx/python/src/transforms.cpp index e6a778eca..6c2299b7f 100644 --- a/Source/Cmlx/mlx/python/src/transforms.cpp +++ b/Source/Cmlx/mlx/python/src/transforms.cpp @@ -334,7 +334,7 @@ auto py_vmap( msg << "[vmap] Invalid" << (output_axes ? " output " : " ") << "vectorization axis " << axis << " for array with shape " << x.shape(); - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } flat_axes.push_back(axis); } else if (nb::isinstance(inputs[1])) { @@ -351,7 +351,7 @@ auto py_vmap( msg << "[vmap] Invalid" << (output_axes ? " output " : " ") << "vectorization axis " << axis << " for array with shape " << x.shape(); - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } flat_axes.push_back(axis); } else if (l.size() == 1 && l[0].is_none()) { @@ -814,8 +814,8 @@ class PyCustomFunction { } int array_index = 0; int tangent_index = 0; - auto new_tangents = - nb::cast(tree_map(args, [&](nb::handle element) { + auto new_tangents = nb::cast( + tree_map(args, [&](nb::handle element) -> nb::object { if (nb::isinstance(element) && have_tangents[array_index++]) { return nb::cast(tangents[tangent_index++]); @@ -861,8 +861,8 @@ class PyCustomFunction { } int arr_index = 0; - auto new_axes = - nb::cast(tree_map(args, [&](nb::handle element) { + auto new_axes = nb::cast( + tree_map(args, [&](nb::handle element) -> nb::object { int axis = axes[arr_index++]; if (nb::isinstance(element) && axis >= 0) { return nb::cast(axis); diff --git a/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py b/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py index cbb966304..7d6e56aaa 100644 --- a/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py +++ b/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py @@ -325,6 +325,22 @@ def test_all_gather(self): self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + def test_all_gather_rank_data(self): + # Every rank contributes distinct data so that a missing, stale, or + # misplaced peer region is detected (an all-ones gather cannot tell). + # Sizes include one large enough that backends which slice transfers + # across directions/wires (e.g. the jaccl ring) exercise every slice, + # and odd sizes that exercise the tail clamping. + world = mx.distributed.init() + for size in [7, 1024, 1000003, 4 * 1024 * 1024]: + base = mx.arange(size, dtype=mx.int32) % 100003 + x = base + world.rank() + y = mx.distributed.all_gather(x) + self.assertEqual(y.shape, (world.size() * size,)) + regions = y.reshape(world.size(), size) + for r in range(world.size()): + self.assertTrue(mx.all(regions[r] == base + r).item()) + def test_clip_grad_norm_sharded(self): world = mx.distributed.init() N = world.size() diff --git a/Source/Cmlx/mlx/python/tests/mlx_tests.py b/Source/Cmlx/mlx/python/tests/mlx_tests.py index 2b60f4615..ba5595389 100644 --- a/Source/Cmlx/mlx/python/tests/mlx_tests.py +++ b/Source/Cmlx/mlx/python/tests/mlx_tests.py @@ -1,5 +1,6 @@ # Copyright © 2023 Apple Inc. +import contextlib import os import platform import sys @@ -10,6 +11,26 @@ import numpy as np +@contextlib.contextmanager +def scoped_env(**environ): + """ + Temporarily set the process environment variables. + + Passing a value of None removes the variable for the duration of the context. + """ + old_environ = dict(os.environ) + for key, value in environ.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + os.environ.clear() + os.environ.update(old_environ) + + class MLXTestRunner(unittest.TestProgram): def __init__(self, *args, **kwargs): # Do not exit in runTests diff --git a/Source/Cmlx/mlx/python/tests/test_array.py b/Source/Cmlx/mlx/python/tests/test_array.py index 80459fabe..e9d781e52 100644 --- a/Source/Cmlx/mlx/python/tests/test_array.py +++ b/Source/Cmlx/mlx/python/tests/test_array.py @@ -140,6 +140,30 @@ def test_scalar_conversion(self): self.assertListEqual(list(z.shape), list(x.shape)) self.assertListEqual(list(z.shape), list(y.shape)) + def test_index_conversion(self): + for dtype in [ + mx.uint8, + mx.uint16, + mx.uint32, + mx.uint64, + mx.int8, + mx.int16, + mx.int32, + mx.int64, + ]: + with self.subTest(dtype=dtype): + self.assertEqual(operator.index(mx.array(2, dtype)), 2) + self.assertEqual(list(range(mx.array(3, dtype))), [0, 1, 2]) + + def test_index_conversion_invalid(self): + for dtype in [mx.float16, mx.float32, mx.bfloat16, mx.complex64, mx.bool_]: + with self.subTest(dtype=dtype): + with self.assertRaises(TypeError): + operator.index(mx.array(2, dtype)) + + with self.assertRaises(TypeError): + list(range(mx.array(3, dtype))) + def test_finfo(self): with self.assertRaises(ValueError): mx.finfo(mx.int32) @@ -468,6 +492,27 @@ def test_array_basics(self): self.assertEqual(x.dtype, mx.int32) self.assertEqual(x.tolist(), [1, 2, 3]) + def test_matrix_transpose(self): + x = mx.array([[1, 2], [3, 4]]) + self.assertEqual(x.mT.tolist(), [[1, 3], [2, 4]]) + self.assertEqual(mx.matrix_transpose(x).tolist(), [[1, 3], [2, 4]]) + + x = mx.arange(24).reshape((2, 3, 4)) + self.assertEqual(x.mT.shape, (2, 4, 3)) + self.assertEqualArray(x.mT, x.transpose((0, 2, 1))) + self.assertEqualArray(x.mT.mT, x) + + self.assertEqual(mx.matrix_transpose(x).shape, (2, 4, 3)) + self.assertEqualArray(mx.matrix_transpose(x), x.transpose((0, 2, 1))) + + x = mx.array([1, 2, 3]) + + with self.assertRaises(ValueError): + x.mT + + with self.assertRaises(ValueError): + mx.matrix_transpose(x) + def test_bool_conversion(self): x = mx.array(True) self.assertTrue(x) @@ -1044,6 +1089,10 @@ def test_array_iteration(self): self.assertEqual(y.tolist(), [3.0, 4.0]) self.assertEqual(z.tolist(), [5.0, 6.0]) + a = mx.array(3) + with self.assertRaises(TypeError): + list(a) + def test_array_pickle(self): dtypes = [ mx.int8, @@ -1267,6 +1316,13 @@ def check_slices(arr_np, *idx_np): a_mlx = mx.array(a_np) self.assertTrue(np.array_equal(a_np[2:-1, 0], np.array(a_mlx[2:-1, 0]))) + # Ellipsis with more trailing indices than dimensions + a_mlx = mx.array([1, 2, 3]) + with self.assertRaises(ValueError): + a_mlx[..., 0, 0] + with self.assertRaises(ValueError): + a_mlx[..., 0, 0] = 5 + def test_indexing_grad(self): x = mx.array([[1, 2], [3, 4]]).astype(mx.float32) ind = mx.array([0, 1, 0]).astype(mx.float32) @@ -1341,6 +1397,13 @@ def test_setitem(self): a[0:1] = mx.array([1]) self.assertEqual(a.tolist(), [1, 4, 4]) + # Regression test: a negative integer index after a None + # (newaxis) used to be normalized against the wrong axis size, + # silently writing nothing instead of updating the last row. + b = mx.zeros((3, 4)) + b[None, -1] = 9 + self.assertEqual(b.tolist(), [[0, 0, 0, 0], [0, 0, 0, 0], [9, 9, 9, 9]]) + with self.assertRaises(ValueError): a[0:1] = mx.array([2, 3]) @@ -2004,6 +2067,15 @@ def test_buffer_protocol(self): self.assertEqual(b"aaaaaaaaaa", ab[::2]) self.assertEqual(b"abcdefghij", ab[1::2]) + # Test bytes on non-contiguous arrays + a = mx.arange(10, dtype=mx.uint8) + self.assertEqual(bytes(a[::2]), b"\x00\x02\x04\x06\x08") + self.assertEqual(bytes(a[::-1]), b"\x09\x08\x07\x06\x05\x04\x03\x02\x01\x00") + b = mx.arange(6, dtype=mx.int32).reshape(2, 3).T + self.assertEqual(bytes(b), np.array(b).tobytes()) + c = mx.broadcast_to(mx.array([1, 2], dtype=mx.uint8), (3, 2)) + self.assertEqual(bytes(c), np.array(c).tobytes()) + def test_buffer_protocol_ref_counting(self): a = mx.arange(3) wr = weakref.ref(a) diff --git a/Source/Cmlx/mlx/python/tests/test_autograd.py b/Source/Cmlx/mlx/python/tests/test_autograd.py index c85fcf09a..13cc3c42a 100644 --- a/Source/Cmlx/mlx/python/tests/test_autograd.py +++ b/Source/Cmlx/mlx/python/tests/test_autograd.py @@ -1319,6 +1319,17 @@ def fun(arrs): self.assertNotEqual(id(a_1), id(arrs[1])) self.assertEqual(id(a_2), id(arrs[2])) + def test_conv_second_order_grad_with_padding(self): + x = mx.random.normal((1, 4, 4, 1)) + w = mx.random.normal((1, 2, 2, 1)) + + def wgrad_sum(x_): + fn = lambda w_: mx.conv2d(x_, w_, padding=(1, 2)).sum() + return mx.grad(fn)(w).sum() + + dx = mx.grad(wgrad_sum)(x) + self.assertEqual(dx.shape, x.shape) + def test_grad_with_inplace_update(self): def loss_fn(model): model[1] = mx.array(2.0) diff --git a/Source/Cmlx/mlx/python/tests/test_bf16.py b/Source/Cmlx/mlx/python/tests/test_bf16.py index d25364c14..1f90e1e51 100644 --- a/Source/Cmlx/mlx/python/tests/test_bf16.py +++ b/Source/Cmlx/mlx/python/tests/test_bf16.py @@ -141,6 +141,21 @@ def test_reduction_ops(self): torch_op="a" + op, ) + def test_arithmetic_reduction_ops(self): + cases = { + "sum": np.ones((4096, 4), dtype=np.float32), + "prod": np.full((128, 4), 1.0078125, dtype=np.float32), + } + for op, values in cases.items(): + with self.subTest(op=op): + x = mx.array(values, dtype=mx.bfloat16) + expected = mx.array( + getattr(np, op)(values, axis=0, dtype=np.float32), + dtype=mx.bfloat16, + ) + actual = getattr(mx, op)(x, axis=0, stream=mx.cpu) + self.assertEqual(actual.tolist(), expected.tolist()) + def test_arg_reduction_ops(self): data = np.random.rand(10, 12, 13).astype(np.float32) x = mx.array(data).astype(mx.bfloat16) diff --git a/Source/Cmlx/mlx/python/tests/test_compile.py b/Source/Cmlx/mlx/python/tests/test_compile.py index 663ec295b..185019342 100644 --- a/Source/Cmlx/mlx/python/tests/test_compile.py +++ b/Source/Cmlx/mlx/python/tests/test_compile.py @@ -7,6 +7,7 @@ import threading from functools import partial, wraps from io import StringIO +from itertools import product import mlx.core as mx import mlx_tests @@ -45,6 +46,22 @@ def fun(x, y): self.assertEqual(out.dtype, mx.int32) self.assertTrue(mx.array_equal(out, mx.array([2, 4]))) + def test_compile_sigmoid_matches_eager(self): + # Regression test: `Sigmoid` spelled an unqualified `metal::exp`, which + # the prebuilt metallib (-fno-fast-math) resolves to the precise + # implementation but a runtime-compiled kernel resolves to the fast + # approximation -- so a fused compiled chain returned a different, less + # accurate sigmoid than the eager primitive. `Exp` in the same header + # already spells `metal::precise::exp`. + x = mx.linspace(-12, 12, 4001) + + for dtype in (mx.float32, mx.bfloat16): + xd = x.astype(dtype) + expected = mx.sigmoid(xd) + fused = mx.compile(lambda a: mx.sigmoid(a) * 1.0)(xd) + mx.eval(expected, fused) + self.assertTrue(mx.array_equal(fused, expected)) + def test_compile_nonfinite_constants(self): # Regression test: a non-finite scalar constant (NaN / infinity) baked # into a fused compiled kernel used to stream a bare token (e.g. `nan`) @@ -71,6 +88,12 @@ def test_compile_nonfinite_constants(self): self.assertEqual(out[0].item(), 1.0) self.assertEqual(out[1].item(), float("-inf")) + 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 + self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + def test_compile_tuple_output_in_thread(self): @mx.compile def fun(x): @@ -635,6 +658,33 @@ def fun(x): cfun = mx.compile(fun, shapeless=True) self.assertTrue(mx.array_equal(fun(x2), cfun(x2))) + def test_shapeless_compile_scan(self): + ops = (mx.cumsum, mx.cumprod, mx.cummin, mx.cummax, mx.logcumsumexp) + for op, axis, reverse, inclusive in product( + ops, (0, 1, -1), (False, True), (False, True) + ): + with self.subTest( + op=op.__name__, axis=axis, reverse=reverse, inclusive=inclusive + ): + scan = partial(op, axis=axis, reverse=reverse, inclusive=inclusive) + trace_count = 0 + + def fun(x): + nonlocal trace_count + trace_count += 1 + return scan(x) + + cfun = mx.compile(fun, shapeless=True) + for shape in ((3, 4), (5, 7), (1, 2), (2, 1)): + with self.subTest(shape=shape): + x = mx.arange(math.prod(shape), dtype=mx.float32) + x = ((x % 7 - 3) / 4).reshape(shape) + expected = scan(x) + actual = cfun(x) + self.assertEqual(actual.shape, x.shape) + self.assertEqualArray(actual, expected, atol=1e-6, rtol=1e-6) + self.assertEqual(trace_count, 1) + def test_shapeless_compile_unflatten(self): x = mx.zeros((1, 1, 4 * 32)) @@ -1293,14 +1343,13 @@ def f(x): def test_double_constant(self): with mx.stream(mx.cpu): - x = mx.array(1.0, dtype=mx.float64) + x = mx.array([1.0], dtype=mx.float64) + constant = math.nextafter(1.0, 2.0) def fun(x): - return (x + math.pi) * 2.0 + return (x * x) * constant - y = fun(x).item() - y_compiled = mx.compile(fun)(x).item() - self.assertEqual(y, y_compiled) + self.assertTrue(mx.array_equal(fun(x), mx.compile(fun)(x))) def test_shared_broadcast(self): def fun(x, y, z): diff --git a/Source/Cmlx/mlx/python/tests/test_conv.py b/Source/Cmlx/mlx/python/tests/test_conv.py index 062841c33..2a317e5b0 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv.py +++ b/Source/Cmlx/mlx/python/tests/test_conv.py @@ -1220,7 +1220,6 @@ def test_conv2d_winograd_batch_tiling(self): # Use envs to test tiling without allocating large buffers. tile_key = "MLX_CONV_WINOGRAD_TILE_BATCH" ws_key = "MLX_CONV_WINOGRAD_WORKING_SET" - prev = {k: os.environ.get(k) for k in (tile_key, ws_key)} # Winograd needs 3x3 stride-1, channels in multiples of 32, # C + O >= 256 and N * iH * iW >= 4096. @@ -1231,81 +1230,68 @@ def test_conv2d_winograd_batch_tiling(self): ) def run(x, w, env={}): - for k in (tile_key, ws_key): - os.environ.pop(k, None) - os.environ.update(env) - y = mx.conv2d(x, w, padding=1) - mx.eval(y) - return np.array(y) - - try: - for in_shape, wt_shape in cases: - np.random.seed(0) - x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) - # Small weights keep the output near unit scale. - w = mx.array( - (np.random.normal(size=wt_shape) * 0.05).astype(np.float32) - ) - b = mx.zeros((wt_shape[0],)) - mx.eval(x, w, b) - cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu)) - - untiled = run(x, w) - self.assertGreater(np.abs(untiled).max(), 0) - self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3)) - - # Tiled winograd keeps the same per-element reduction order, - # so it is bit-identical to untiled; the implicit gemm - # fallback never is. Exact equality pins each run to its path. - # 3 divides none of the batches, so it also covers a short - # final tile. - for tile in (1, 3): - with self.subTest(in_shape=in_shape, tile=tile): - tiled = run(x, w, {tile_key: str(tile)}) - self.assertTrue(np.array_equal(untiled, tiled)) - - # A consumer op checks the output is fenced across - # command encoders. - os.environ[tile_key] = str(tile) - fused = mx.conv2d(x, w, padding=1) + b - mx.eval(fused) - self.assertTrue(np.allclose(untiled, fused, atol=1e-4)) - os.environ.pop(tile_key, None) - - # 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] - 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 - ) - used = (n * iH * iW * (C + O) + 64 * C * O) * 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}) + with mlx_tests.scoped_env(**env): + y = mx.conv2d(x, w, padding=1) + mx.eval(y) + return np.array(y) + + for in_shape, wt_shape in cases: + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + # Small weights keep the output near unit scale. + w = mx.array((np.random.normal(size=wt_shape) * 0.05).astype(np.float32)) + b = mx.zeros((wt_shape[0],)) + mx.eval(x, w, b) + cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu)) + + untiled = run(x, w) + self.assertGreater(np.abs(untiled).max(), 0) + self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3)) + + # Tiled winograd keeps the same per-element reduction order, + # so it is bit-identical to untiled; the implicit gemm + # fallback never is. Exact equality pins each run to its path. + # 3 divides none of the batches, so it also covers a short + # final tile. + for tile in (1, 3): + with self.subTest(in_shape=in_shape, tile=tile): + tiled = run(x, w, {tile_key: str(tile)}) self.assertTrue(np.array_equal(untiled, tiled)) - # Too small for even one batch element: must fall back. - with self.subTest(in_shape=in_shape, budget="infeasible"): - fallback = run(x, w, {ws_key: "1"}) - self.assertFalse(np.array_equal(untiled, fallback)) - self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3)) - - # A forced tile is capped by the budget, so this must still - # fall back. - with self.subTest(in_shape=in_shape, budget="forced+infeasible"): - capped = run(x, w, {ws_key: "1", tile_key: "1"}) - self.assertFalse(np.array_equal(untiled, capped)) - self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3)) - finally: - for k, v in prev.items(): - if v is None: - os.environ.pop(k, None) - else: - os.environ[k] = v + # A consumer op checks the output is fenced across + # command encoders. + with mlx_tests.scoped_env(MLX_CONV_WINOGRAD_TILE_BATCH=str(tile)): + fused = mx.conv2d(x, w, padding=1) + b + mx.eval(fused) + self.assertTrue(np.allclose(untiled, fused, atol=1e-4)) + + # 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] + 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 + ) + used = (n * iH * iW * (C + O) + 64 * C * O) * 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}) + self.assertTrue(np.array_equal(untiled, tiled)) + + # Too small for even one batch element: must fall back. + with self.subTest(in_shape=in_shape, budget="infeasible"): + fallback = run(x, w, {ws_key: "1"}) + self.assertFalse(np.array_equal(untiled, fallback)) + self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3)) + + # A forced tile is capped by the budget, so this must still + # fall back. + with self.subTest(in_shape=in_shape, budget="forced+infeasible"): + capped = run(x, w, {ws_key: "1", tile_key: "1"}) + self.assertFalse(np.array_equal(untiled, capped)) + self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3)) def test_conv2d_large_filter_small_channels(self): x = mx.random.normal(shape=(1, 181, 181, 1)) diff --git a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py index 7d12cb16f..28593a530 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py +++ b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py @@ -813,37 +813,27 @@ def test_conv_transpose_unfold_tiling(self): # The explicit-GEMM conv path unfolds into one buffer that can exceed # maxBufferLength for large outputs; it unfolds and runs the gemm in row # tiles instead (issue #3082). - key = "MLX_CONV_UNFOLD_TILE_ROWS" - prev = os.environ.get(key) cases = ( (mx.conv_transpose1d, (2, 9, 4), (5, 3, 4), {"stride": 2}), (mx.conv_transpose2d, (2, 5, 5, 4), (5, 3, 3, 4), {"stride": 2}), (mx.conv_transpose3d, (1, 4, 4, 4, 2), (3, 2, 2, 2, 2), {"stride": 2}), (mx.conv_transpose1d, (2, 9, 4), (6, 3, 2), {"stride": 2, "groups": 2}), ) - try: - # tile_rows=1 forces uniform tiles; 7 does not divide any output, so - # it exercises a partial final tile too. - for conv, in_shape, wt_shape, kwargs in cases: - for tile_rows in (1, 7): - with self.subTest( - conv=conv.__name__, tile_rows=tile_rows, **kwargs - ): - np.random.seed(0) - x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) - w = mx.array(np.random.normal(size=wt_shape).astype(np.float32)) - os.environ.pop(key, None) + # tile_rows=1 forces uniform tiles; 7 does not divide any output, so + # it exercises a partial final tile too. + for conv, in_shape, wt_shape, kwargs in cases: + for tile_rows in (1, 7): + with self.subTest(conv=conv.__name__, tile_rows=tile_rows, **kwargs): + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + w = mx.array(np.random.normal(size=wt_shape).astype(np.float32)) + with mlx_tests.scoped_env(MLX_CONV_UNFOLD_TILE_ROWS=None): untiled = conv(x, w, **kwargs) mx.eval(untiled) - os.environ[key] = str(tile_rows) + with mlx_tests.scoped_env(MLX_CONV_UNFOLD_TILE_ROWS=str(tile_rows)): tiled = conv(x, w, **kwargs) mx.eval(tiled) - self.assertTrue(np.allclose(untiled, tiled, atol=1e-4)) - finally: - if prev is None: - os.environ.pop(key, None) - else: - os.environ[key] = prev + self.assertTrue(np.allclose(untiled, tiled, atol=1e-4)) if __name__ == "__main__": diff --git a/Source/Cmlx/mlx/python/tests/test_device.py b/Source/Cmlx/mlx/python/tests/test_device.py index 6cd6df81f..c0ed9ca90 100644 --- a/Source/Cmlx/mlx/python/tests/test_device.py +++ b/Source/Cmlx/mlx/python/tests/test_device.py @@ -1,5 +1,7 @@ # Copyright © 2023 Apple Inc. +import concurrent.futures +import gc import unittest import mlx.core as mx @@ -112,6 +114,44 @@ def test_op_on_stream(self): b = mx.add(x, y, stream=s_cpu) self.assertEqual(a.item(), b.item()) + def test_reentrant_stream_context(self): + default = mx.default_stream(mx.default_device()) + s = mx.new_stream(mx.default_device()) + context = mx.stream(s) + + with context: + with context: + self.assertEqual(mx.default_stream(mx.default_device()), s) + self.assertEqual(mx.default_stream(mx.default_device()), s) + self.assertEqual(mx.default_stream(mx.default_device()), default) + + def test_stream_context_exited_on_another_thread(self): + # A generator suspended inside `mx.stream(...)` and abandoned is + # finalized by the GC on whatever thread drops the last reference. + # Restoring there would install a foreign, thread-affine stream as + # that thread's default and break every later op on it. + abandoned = {} + + def on_worker(): + def gen(): + with mx.stream(mx.new_stream(mx.default_device())): + yield 1 + yield 2 + + g = gen() + next(g) + abandoned["gen"] = g + + before = mx.default_stream(mx.default_device()) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + pool.submit(on_worker).result() + + del abandoned["gen"] + gc.collect() + + self.assertEqual(mx.default_stream(mx.default_device()), before) + mx.eval(mx.add(mx.array(1.0), mx.array(1.0))) + class TestDeviceInfo(mlx_tests.MLXTestCase): def test_device_count(self): diff --git a/Source/Cmlx/mlx/python/tests/test_fast.py b/Source/Cmlx/mlx/python/tests/test_fast.py index 200781d37..56f6211ef 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast.py +++ b/Source/Cmlx/mlx/python/tests/test_fast.py @@ -535,6 +535,9 @@ def cross_entropy_ref(logits, targets): ) tolerances = {mx.float32: 1e-5, mx.float16: 3e-2, mx.bfloat16: 3e-1} + # The fused GPU kernels accumulate in float32 for every input dtype. + if mx.default_device() == mx.gpu: + tolerances = {mx.float32: 1e-5, mx.float16: 1e-3, mx.bfloat16: 1e-3} for V in [7, 32, 128, 255, 256, 1000, 4096, 8192]: for dtype in [mx.float32, mx.float16, mx.bfloat16]: diff --git a/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py b/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py new file mode 100644 index 000000000..d7a33c67a --- /dev/null +++ b/Source/Cmlx/mlx/python/tests/test_fast_gated_delta.py @@ -0,0 +1,286 @@ +import os +import unittest + +import mlx.core as mx +import mlx_tests +import numpy as np + +try: + import torch + + has_torch = True +except ImportError as e: + has_torch = False + + +def gated_delta_oracle( + q, + k, + v, + beta, + g, + scale=None, + initial_state=None, + output_final_state=False, +): + """ + Reference PyTorch implementation of recurrent gated delta rule. + Taken from: https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/naive.py + + Args: + q: [B, T, H, K] + 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. + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, beta, g = map( + lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g] + ) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum("bhd,bhdm->bhm", b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def runner(dims, stream=mx.gpu, reference=True): + B, Hk, Hv, T, Dk, Dv = dims + + assert Hv % Hk == 0 + repeat_factor = Hv // Hk + + 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.random.uniform(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) + + if reference: + # Prepare reference inputs + qpt = torch.from_numpy(np.array(q)) + kpt = torch.from_numpy(np.array(k)) + vpt = torch.from_numpy(np.array(v)) + bpt = torch.from_numpy(np.array(b)) + gpt = torch.from_numpy(np.array(g)) + h0pt = torch.from_numpy(np.array(h0)).transpose(-1, -2).contiguous() + + if repeat_factor > 1: + qpt = qpt.repeat_interleave(repeat_factor, dim=2) + kpt = kpt.repeat_interleave(repeat_factor, dim=2) + + out_on_py, hf_on_py = gated_delta_oracle( + qpt, + kpt, + vpt, + bpt, + torch.log(gpt), + scale=1.0, + initial_state=h0pt, + output_final_state=True, + ) + + out_on = mx.array(out_on_py.detach().cpu().numpy()) # [B, T, Hv, Dv] + hf_on = mx.swapaxes( + mx.array(hf_on_py.detach().cpu().numpy()), -1, -2 + ) # -> [B, Hv, Dv, Dk] + 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 + mask = mx.ones((B, T)) + + out_ref, hf_ref = mx.fast.gated_delta_update( + q, + k, + v, + g, + b, + initial_state=h0, + mask=mask, + ) + + mx.eval(out_ref, hf_ref) + + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, stream=stream) + + mx.eval(out, hf) + return (out, hf), (out_ref, hf_ref) + + +class TestGatedDelta(mlx_tests.MLXTestCase): + base_dims = (1, 32, 32, 1, 128, 128) + unaligned_dims = (1, 32, 32, 33, 128, 128) + big_batch_dims = (128, 32, 32, 16, 128, 128) + large_t_dims = (2, 32, 32, 1111, 128, 128) + diff_heads = (1, 16, 32, 33, 128, 128) + diff_heads2 = (1, 16, 48, 33, 128, 128) + unsupported_heads1 = (1, 8, 32, 33, 128, 128) + unsupported_heads2 = (1, 24, 48, 33, 128, 128) + + fallback_dims = [unsupported_heads1, unsupported_heads2] + gpu_dims = [ + base_dims, + unaligned_dims, + big_batch_dims, + diff_heads, + diff_heads2, + large_t_dims, + ] + + @unittest.skipIf(not has_torch, "requires Torch") + def test_gated_delta_fallback(self): + for dims in self.fallback_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, mx.cpu) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + 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)) + 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.random.uniform(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) + + # make a mask + lengths = mx.random.randint(1, T + 1, shape=(B,)) + mask = mx.arange(T)[None, :] < lengths[:, None] + # mask one input in python + mask_float = mask.astype(q.dtype) + km = k * mask_float[..., None, None] + vm = v * mask_float[..., None, None] + qm = q * mask_float[..., None, None] + bm = b * mask_float[..., None] + gm = mx.where(mask[..., None], g, 1.0) + + out_ref, hf_ref = mx.fast.gated_delta_update( + qm, + km, + vm, + gm, + bm, + initial_state=h0, + ) + + mx.eval(out_ref, hf_ref) + out, hf = mx.fast.gated_delta_update( + q, + k, + v, + g, + b, + initial_state=h0, + mask=mask, + ) + mx.eval(out, hf) + + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + 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) + k = mx.random.normal(shape=(B, T, Hk, Dk), dtype=dtype) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, Hv, Dv), dtype=dtype) + g = mx.random.uniform(shape=(B, T, Hv), dtype=dtype) + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv), dtype=dtype)) + h0 = mx.random.normal((B, Hv, Dv, Dk), dtype=mx.float32) + + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0) + + msg = f"Output dtype mismatch on Dimensions: {dims}" + self.assertTrue(dtype == out.dtype, msg="Out " + msg) + self.assertTrue(hf.dtype == mx.float32, msg="State " + msg) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_gated_delta_sequential(self): + os.environ["GATED_DELTA_CHUNK"] = "0" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_gated_delta_simdgroup(self): + os.environ["GATED_DELTA_CHUNK"] = "8" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_gated_delta_nax(self): + os.environ["GATED_DELTA_CHUNK"] = "16" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-1, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-1, rtol=1e-4), msg="State " + msg + ) + + +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 b11ce884e..142d34195 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py +++ b/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py @@ -144,6 +144,60 @@ def test_sdpa_head_dim_72(self): diff = mx.abs(out - ref) - atol * mx.abs(ref) self.assertLessEqual(mx.max(diff).item(), atol) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_pad_head_dim_opt_in(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + with mlx_tests.scoped_env(MLX_SDPA_PAD_HEAD_DIM="1"): + self.test_sdpa_head_dim_72() + self.test_sdpa_head_dim_80() + self.test_sdpa_head_dim_72_80_sinks() + + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_80(self): + B, D, qH, kH = (1, 80, 8, 2) + for qL, kL, dtype, mask_str in product( + (64, 65), + (128, 127), + (mx.float16, mx.bfloat16, mx.float32), + (None, "additive", "bool", "causal"), + ): + with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, False, dtype + ) + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + + if dtype == mx.float32: + atol = 1e-5 + elif dtype == mx.bfloat16: + atol = 5e-3 + else: + atol = 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_72_80_sinks(self): + B, qH, kH, qL, kL = (1, 8, 2, 64, 128) + for D, dtype in product((72, 80), (mx.float16, mx.bfloat16)): + with self.subTest(D=D, dtype=dtype): + q, k, v, scale, _ = prepare_inputs( + B, qL, kL, D, qH, kH, None, False, dtype + ) + sinks = 10 * mx.random.normal(shape=(qH,), dtype=dtype) + ref = mlx_ref_attn(q, k, v, scale, sinks=sinks) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, sinks=sinks + ) + + atol = 5e-3 if dtype == mx.bfloat16 else 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") def test_sdpa_head_dim_96(self): B, D, qH, kH = (1, 96, 8, 2) @@ -171,6 +225,75 @@ def test_sdpa_head_dim_96(self): diff = mx.abs(out - ref) - atol * mx.abs(ref) self.assertLessEqual(mx.max(diff).item(), atol) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_full_d96_v64(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + mx.random.seed(0) + scale = 96**-0.5 + cells = [ + (1, 40, 40, 9, 33), + (1, 4, 4, 64, 64), + (2, 4, 2, 65, 97), + (1, 4, 4, 16, 9), + (1, 4, 2, 32, 31), + (1, 4, 2, 33, 32), + ] + for dtype, (B, qH, kH, qL, kL) in product( + (mx.float16, mx.bfloat16, mx.float32), cells + ): + q = ( + (0.5 * mx.random.normal((B, qL, qH, 96))) + .astype(dtype) + .transpose(0, 2, 1, 3) + ) + k = (0.5 * mx.random.normal((B, kH, kL + 32, 96))).astype(dtype)[:, :, :kL] + masks = [ + None, + mx.arange(kL) < kL - 3, + mx.random.normal((qL, kL)).astype(dtype), + ] + if qL <= kL: + masks.append("causal") + for dv, mask, with_sinks in product((64, 96), masks, (False, True)): + with self.subTest( + dtype=dtype, + cell=(B, qH, kH, qL, kL), + dv=dv, + mask=type(mask).__name__, + sinks=with_sinks, + ): + v = (0.5 * mx.random.normal((B, kH, kL + 32, dv))).astype(dtype)[ + :, :, :kL + ] + if qL == 9: + # V sliced from the interleaved K/V projection. + v = (0.5 * mx.random.normal((B, kL, kH, 2 * dv))).astype(dtype) + v = v.transpose(0, 2, 1, 3)[..., dv:] + sinks = ( + mx.random.normal((qH,)).astype(dtype) if with_sinks else None + ) + ref = mlx_ref_attn( + q.astype(mx.float32), + k.astype(mx.float32), + v.astype(mx.float32), + scale, + mask, + sinks, + ) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + mask=mask, + sinks=sinks, + force_fused=True, + ) + self.assertEqual(out.shape, (B, qH, qL, dv)) + tol = 5e-3 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") def test_sdpa_full_head_dim_256(self): # On NAX devices, large nearly-square causal blocks take the fused @@ -238,6 +361,116 @@ def test_sdpa_full_head_dim_256(self): tol = 5e-3 self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + @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: + self.skipTest("requires GPU") + + D = 512 + scale = D**-0.5 + mx.random.seed(0) + + q_probe = mx.random.normal(shape=(1, 2, 9, D)).astype(mx.bfloat16) + k_probe = mx.random.normal(shape=(1, 1, 31, D)).astype(mx.bfloat16) + v_probe = mx.random.normal(shape=(1, 1, 31, D)).astype(mx.bfloat16) + try: + mx.eval( + mx.fast.scaled_dot_product_attention( + q_probe, + k_probe, + v_probe, + scale=scale, + force_fused=True, + ) + ) + except ValueError as e: + if "plus head dim 512 on NAX GPUs" not in str(e): + raise + self.skipTest("D512 full attention requires NAX") + + cases = [ + # Multiple query tiles with GQA. + (1, 8, 2, 257, 513, None, True), + # Multiple batches and a ragged final tile. + (2, 4, 1, 513, 777, None, True), + # Gemma 4 31B default route at the work threshold with sliced K/V. + (1, 32, 16, 1024, 1025, 1056, False), + ] + for dtype in (mx.float16, mx.bfloat16): + for B, Nq, Nkv, qL, kL, cache_len, force_fused in cases: + with self.subTest( + dtype=dtype, + B=B, + Nq=Nq, + Nkv=Nkv, + qL=qL, + kL=kL, + cache_len=cache_len, + ): + q = mx.random.normal(shape=(B, Nq, qL, D)).astype(dtype) + if cache_len is None: + k = mx.random.normal(shape=(B, Nkv, kL, D)).astype(dtype) + v = mx.random.normal(shape=(B, Nkv, kL, D)).astype(dtype) + else: + k = mx.random.normal(shape=(B, Nkv, cache_len, D)).astype( + dtype + )[..., :kL, :] + v = mx.random.normal(shape=(B, Nkv, cache_len, D)).astype( + dtype + )[..., :kL, :] + + ref = mlx_ref_attn(q, k, v, scale, "causal") + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + mask="causal", + force_fused=force_fused, + ) + tol = 2e-3 if dtype != mx.bfloat16 else 5e-2 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + q = mx.random.normal(shape=(1, 8, 33, D)).astype(mx.bfloat16) + k = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16) + v = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-2, rtol=1e-2)) + + bool_mask = mx.random.uniform(shape=(1, 1, 33, 67)) > 0.2 + additive_mask = mx.where(bool_mask, 0.0, -1e4).astype(mx.bfloat16) + for mask in (bool_mask, additive_mask): + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-2, rtol=1e-2)) + + q_same = mx.random.normal(shape=(1, 8, 33, D)).astype(mx.bfloat16) + k_same = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16) + v_slice = mx.random.normal(shape=(1, 2, 67, D // 4)).astype(mx.bfloat16) + v_same = mx.tile(v_slice, (1, 1, 1, 4)) + out_same = mx.fast.scaled_dot_product_attention( + q_same, k_same, v_same, scale=scale, force_fused=True + ) + for d_group in range(1, 4): + self.assertTrue( + mx.array_equal( + out_same[..., : D // 4], + out_same[..., d_group * D // 4 : (d_group + 1) * D // 4], + ) + ) + + sinks = mx.random.normal(shape=(8,)).astype(mx.bfloat16) + ref = mlx_ref_attn(q, k, v, scale, sinks=sinks) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, sinks=sinks, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-2, rtol=1e-2)) + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 @@ -345,17 +578,137 @@ def test_sdpa_vector(self): def test_sdpa_vector_gqa_long(self): scale = 1.0 mx.random.seed(0) - for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]: - for L in [8192, 8201]: - q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D)) - k = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] - v = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64), (48, 4, 128), (64, 4, 128)]: + for B, L in [(1, 8192), (1, 8201), (2, 8192)]: + q = 5e-1 * mx.random.normal(shape=(B, Nq, 1, D)) + k = 5e-1 * mx.random.normal(shape=(B, Nkv, L + 32, D))[:, :, :L] + v = 5e-1 * mx.random.normal(shape=(B, Nkv, L + 32, D))[:, :, :L] kr = mx.repeat(k, Nq // Nkv, axis=1) vr = mx.repeat(v, Nq // Nkv, axis=1) ref = mlx_primitives_sdpa(q, kr, vr, scale) out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) 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_d96_v64(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + mx.random.seed(0) + scale = 96**-0.5 + # Include MiniCPM3 heads and GQA to reach two-pass on smaller GPUs. + cells = [(1, 40, 40, 1, 127), (1, 40, 40, 8, 1025), (2, 4, 2, 4, 4097)] + for dtype, (B, qH, kH, qL, kL) in product( + (mx.float32, mx.float16, mx.bfloat16), cells + ): + q = mx.random.normal((B, qL, qH, 96)).astype(dtype) + q = q.transpose(0, 2, 1, 3) + k = mx.random.normal((B, kH, kL + 32, 96)).astype(dtype)[:, :, :kL] + v = mx.random.normal((B, kH, kL + 32, 64)).astype(dtype)[:, :, :kL] + masks = ( + None, + "causal", + mx.arange(kL) < kL - 3, + mx.random.normal((qL, kL)).astype(dtype), + ) + for mask in masks: + with self.subTest( + dtype=dtype, cell=(B, qH, kH, qL, kL), mask=type(mask).__name__ + ): + # GPU float32 GEMM may use TF32; keep the reference exact. + with mx.stream(mx.cpu): + ref = mlx_ref_attn( + q.astype(mx.float32), + k.astype(mx.float32), + v.astype(mx.float32), + scale, + mask, + ) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + self.assertEqual(out.shape, (B, qH, qL, 64)) + tol = ( + 1e-5 + if dtype == mx.float32 + else (3e-3 if dtype == mx.float16 else 2e-2) + ) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_vector_head_dim_512(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + # gemma-4 global attention: 32 query heads over 4 key/value heads. + D = 512 + Nq, Nkv = 32, 4 + scale = D**-0.5 + mx.random.seed(0) + + with mlx_tests.scoped_env(MLX_SDPA_D512_MIN_KL="0"): + # Test 1-pass kernel. + for dtype in (mx.float32, mx.float16, mx.bfloat16): + with self.subTest(L=128, dtype=dtype, threshold="off"): + q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype) + k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype) + v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + + # Test 2-pass kernel. + for L, dtype in product( + (8192, 8201), (mx.float32, mx.float16, mx.bfloat16) + ): + with self.subTest(L=L, dtype=dtype): + q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype) + k = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype) + v = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + + # Test other heads. + for q_heads, kv_heads in ((8, 8), (32, 8)): + with self.subTest(q_heads=q_heads, kv_heads=kv_heads): + q = mx.random.normal(shape=(1, q_heads, 1, D)) + k = mx.random.normal(shape=(1, kv_heads, L, D)) + v = mx.random.normal(shape=(1, kv_heads, L, D)) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + force_fused=True, + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + + # Test batched. + B = 2 + sinks = 10 * mx.random.normal(shape=(Nq,)) + q = mx.random.normal(shape=(B, Nq, 1, D)) + for L in (256, 8192): + k = mx.random.normal(shape=(B, Nkv, L, D)) + v = mx.random.normal(shape=(B, Nkv, L, D)) + for s_in in (None, sinks): + with self.subTest(L=L, sinks=s_in is not None): + ref = mlx_ref_attn(q, k, v, scale, sinks=s_in) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + sinks=s_in, + force_fused=True, + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_fully_masked(self): Lkv = 8 mask = mx.array(False) @@ -505,13 +858,10 @@ def test_sdpa_blocks_env_override(self): k = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) v = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) ref = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) - try: - for blocks in (16, 33, 48, 100): - os.environ["MLX_SDPA_BLOCKS"] = str(blocks) + for blocks in (16, 33, 48, 100): + with mlx_tests.scoped_env(MLX_SDPA_BLOCKS=str(blocks)): out = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) - finally: - del os.environ["MLX_SDPA_BLOCKS"] @unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU") def test_sdpa(self): @@ -568,15 +918,16 @@ def test_sdpa(self): out_ref = do_attention(mlx_ref_attn, q, k, v, scale, mask, t) - out_fst = do_attention( - mx.fast.scaled_dot_product_attention, - q, - k, - v, - scale, - mask, - t, - ) + with mlx_tests.scoped_env(MLX_ENABLE_TF32="0"): + out_fst = do_attention( + mx.fast.scaled_dot_product_attention, + q, + k, + v, + scale, + mask, + t, + ) # For causal mask when qL > kL, first qL-kL rows are undefined # Compare only the valid portion @@ -807,40 +1158,49 @@ def test_grad(slow, fast, args): def test_sdpa_force_fused_metal(self): if mx.default_device() != mx.gpu: self.skipTest("requires GPU") + mx.random.seed(0) - def make_qkv(qL, kL, D, qH=8, kH=8): - q = mx.random.normal((1, qH, qL, D), mx.float16) - k = mx.random.normal((1, kH, kL, D), mx.float16) - v = mx.random.normal((1, kH, kL, D), mx.float16) + def make_qkv(qL, kL, D, qH=8, kH=8, dtype=mx.float16): + q = mx.random.normal((1, qH, qL, D), dtype) + k = mx.random.normal((1, kH, kL, D), dtype) + v = mx.random.normal((1, kH, kL, D), dtype) return q, k, v # Full attention kernel. - for D, qL, mask in product((192, 256), (9, 16), (None, "causal")): - with self.subTest(head_dim=D, qL=qL, mask=mask): - q, k, v = make_qkv(qL, 512, D, 8, 4) + for D, qL, kL, mask, dtype in product( + (192, 256), + (9, 16), + (31, 512), + (None, "causal"), + (mx.bfloat16, mx.float32), + ): + with self.subTest(head_dim=D, qL=qL, kL=kL, mask=mask): + q, k, v = make_qkv(qL, kL, D, 8, 4) scale = D**-0.5 ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) out = mx.fast.scaled_dot_product_attention( q, k, v, scale=scale, mask=mask, force_fused=True ) - self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) # Vector attention kernel. - for D in (192, 256): + for D in (192, 256, 512): with self.subTest(head_dim=D): - q, k, v = make_qkv(4, 16385, D, 4, 2) - scale = D**-0.5 - ref = mlx_ref_attn(q, k, v, scale=scale) - out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, force_fused=True - ) - self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + with mlx_tests.scoped_env(MLX_SDPA_D512_MIN_KL="0"): + q, k, v = make_qkv(4, 16385, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) # No full attention fused kernels. with self.assertRaisesRegex(ValueError, "supports head dims"): - q, k, v = make_qkv(16, 512, 512) + q, k, v = make_qkv(16, 512, 1024) mx.fast.scaled_dot_product_attention( - q, k, v, scale=512**-0.5, force_fused=True + q, k, v, scale=1024**-0.5, force_fused=True ) with self.assertRaisesRegex( ValueError, "query sequence to be no longer than the key sequence" @@ -866,6 +1226,12 @@ def make_qkv(qL, kL, D, qH=8, kH=8): mx.fast.scaled_dot_product_attention( q, k, v, scale=64**-0.5, force_fused=True ) + with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): + with mlx_tests.scoped_env(MLX_SDPA_D512_MIN_KL=None): + q, k, v = make_qkv(1, 512, 512, qH=32, kH=4) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) # No CPU fused kernel. with mx.stream(mx.cpu): diff --git a/Source/Cmlx/mlx/python/tests/test_fft.py b/Source/Cmlx/mlx/python/tests/test_fft.py index 1f96aad56..60f7610fe 100644 --- a/Source/Cmlx/mlx/python/tests/test_fft.py +++ b/Source/Cmlx/mlx/python/tests/test_fft.py @@ -385,9 +385,9 @@ def test_ifftshift(self): def test_fftshift_errors(self): # Test invalid axes x = mx.array(np.random.rand(4, 4).astype(np.float32)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.fft.fftshift(x, axes=[2]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.fft.fftshift(x, axes=[-3]) # Test empty array diff --git a/Source/Cmlx/mlx/python/tests/test_linalg.py b/Source/Cmlx/mlx/python/tests/test_linalg.py index 52e02e4d9..2da3781b4 100644 --- a/Source/Cmlx/mlx/python/tests/test_linalg.py +++ b/Source/Cmlx/mlx/python/tests/test_linalg.py @@ -306,6 +306,35 @@ def test_cholesky(self): for M, L in zip(AB, Ls): self.assertTrue(mx.allclose(L @ L.T, M, rtol=1e-5, atol=1e-7)) + if mx.cuda.is_available(): + # The matrix above is singular, so its factor is undefined past the + # rank boundary and the GPU need not match the CPU there. Use + # positive definite inputs, sized to hit both cuSOLVER paths. + mx.random.seed(7) + for batch, n in [((), 3), ((2,), 2048), ((16,), 8)]: + x = mx.random.normal(batch + (n, n)) + A = x @ mx.swapaxes(x, -1, -2) + n * mx.eye(n) + for upper in [False, True]: + G = mx.linalg.cholesky(A, upper=upper, stream=mx.gpu) + C = mx.linalg.cholesky(A, upper=upper, stream=mx.cpu) + self.assertTrue(mx.allclose(G, C, rtol=1e-4, atol=1e-5)) + tri = mx.tril(G, k=-1) if upper else mx.triu(G, k=1) + self.assertTrue(mx.all(tri == 0)) + + # Empty and non contiguous inputs + self.assertEqual( + mx.linalg.cholesky(mx.zeros((0, 0)), stream=mx.gpu).shape, (0, 0) + ) + At = mx.swapaxes(A, -1, -2) + self.assertTrue( + mx.allclose( + mx.linalg.cholesky(At, stream=mx.gpu), + mx.linalg.cholesky(At, stream=mx.cpu), + rtol=1e-4, + atol=1e-5, + ) + ) + def test_pseudo_inverse(self): A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32) A_plus = mx.linalg.pinv(A, stream=mx.cpu) diff --git a/Source/Cmlx/mlx/python/tests/test_load.py b/Source/Cmlx/mlx/python/tests/test_load.py index f5947de47..d9ad9751c 100644 --- a/Source/Cmlx/mlx/python/tests/test_load.py +++ b/Source/Cmlx/mlx/python/tests/test_load.py @@ -122,6 +122,24 @@ def test_async_load_npy_read_error_across_streams(self): except Exception: pass + def test_save_over_lazily_loaded_safetensors(self): + # The save must eval the lazy input before truncating the file it was + # loaded from. Shift the data so stale content cannot pass. + save_file = os.path.join(self.test_dir, "resave.safetensors") + mx.save_safetensors(save_file, {"a": mx.ones((4, 4))}) + a = mx.load(save_file, stream=mx.cpu)["a"] + mx.save_safetensors(save_file, {"a": a + 1}) + out = mx.load(save_file, stream=mx.cpu)["a"] + self.assertTrue(np.array_equal(np.array(out), 2 * np.ones((4, 4)))) + + def test_save_over_lazily_loaded_npy(self): + save_file = os.path.join(self.test_dir, "resave.npy") + mx.save(save_file, mx.ones((4, 4))) + a = mx.load(save_file, stream=mx.cpu) + mx.save(save_file, a + 2) + out = mx.load(save_file, stream=mx.cpu) + self.assertTrue(np.array_equal(np.array(out), 3 * np.ones((4, 4)))) + def test_save_and_load_safetensors(self): test_file = os.path.join(self.test_dir, "test.safetensors") with self.assertRaises(Exception): diff --git a/Source/Cmlx/mlx/python/tests/test_memory.py b/Source/Cmlx/mlx/python/tests/test_memory.py index 2a0a544b6..8436dab09 100644 --- a/Source/Cmlx/mlx/python/tests/test_memory.py +++ b/Source/Cmlx/mlx/python/tests/test_memory.py @@ -1,4 +1,4 @@ -# Copyright © 2023-2024 Apple Inc. +# Copyright © 2023-2026 Apple Inc. import unittest @@ -7,6 +7,31 @@ class TestMemory(mlx_tests.MLXTestCase): + def test_array_buffer_size(self): + a = mx.array([1.0, 2.0, 3.0, 4.0]) + view = a[:1] + lazy = a + 1 + + with self.assertRaises(ValueError): + mx.get_array_buffer_size({"lazy": lazy}) + + mx.eval(view) + a_size = mx.get_array_buffer_size(a) + self.assertGreaterEqual(a_size, a.nbytes) + self.assertLess(view.nbytes, a.nbytes) + self.assertEqual(mx.get_array_buffer_size(), 0) + self.assertEqual(mx.get_array_buffer_size(view), a_size) + self.assertEqual( + mx.get_array_buffer_size({"a": a, "nested": (a, None)}), a_size + ) + self.assertEqual(mx.get_array_buffer_size(mx.array([])), 0) + + b = mx.array([5.0, 6.0]) + self.assertEqual( + mx.get_array_buffer_size(a, b), + a_size + mx.get_array_buffer_size(b), + ) + def test_memory_info(self): old_limit = mx.set_cache_limit(0) @@ -18,6 +43,7 @@ def test_memory_info(self): self.assertEqual(mx.set_cache_limit(old_limit), old_limit) old_limit = mx.set_memory_limit(10) + self.assertEqual(mx.get_memory_limit(), 10) self.assertEqual(mx.set_memory_limit(old_limit), 10) self.assertEqual(mx.set_memory_limit(old_limit), old_limit) diff --git a/Source/Cmlx/mlx/python/tests/test_ops.py b/Source/Cmlx/mlx/python/tests/test_ops.py index 0c94989e0..151eddea8 100644 --- a/Source/Cmlx/mlx/python/tests/test_ops.py +++ b/Source/Cmlx/mlx/python/tests/test_ops.py @@ -378,6 +378,31 @@ def test_divide(self): self.assertEqual(z.dtype, mx.int32) self.assertEqual(z.item(), 2) + def test_floor_divide(self): + a = [4, 5, -1, -6] + b = [-2, 3, 2, -3] + + mx_int_float_dtypes = mx.__array_namespace_info__().dtypes( + kind=("signed integer", "real floating") + ) + np_int_float_dtypes = np.__array_namespace_info__().dtypes( + kind=("signed integer", "real floating") + ) + + for kind, dtype in mx_int_float_dtypes.items(): + with self.subTest(dtype=dtype): + result = mx.floor_divide( + mx.array(a, dtype=dtype), mx.array(b, dtype=dtype) + ) + expected = np.floor_divide( + np.array(a, dtype=np_int_float_dtypes[kind]), + np.array(b, dtype=np_int_float_dtypes[kind]), + ) + self.assertEqual(result.tolist(), expected.tolist()) + + result = mx.array(a, dtype=dtype) // mx.array(b, dtype=dtype) + self.assertEqual(result.tolist(), expected.tolist()) + def test_remainder(self): # Complex is not supported and has to say so rather than quietly # computing a componentwise remainder, which no other library defines @@ -593,6 +618,11 @@ def test_minimum(self): self.assertTrue(math.isnan(mx.minimum(a, b).item())) self.assertTrue(math.isnan(mx.minimum(b, a).item())) + # A NaN in either part of the first argument propagates. + a = mx.array([complex(3, float("nan"))]) + b = mx.array([2 + 0j]) + self.assertTrue(mx.array_equal(mx.minimum(a, b), a, equal_nan=True)) + def test_maximum(self): x = mx.array([0.0, -5, 10.0]) y = mx.array([1.0, -7.0, 3.0]) @@ -605,6 +635,11 @@ def test_maximum(self): self.assertTrue(math.isnan(mx.maximum(a, b).item())) self.assertTrue(math.isnan(mx.maximum(b, a).item())) + # A NaN in either part of the first argument propagates. + a = mx.array([complex(1, float("nan"))]) + b = mx.array([2 + 0j]) + self.assertTrue(mx.array_equal(mx.maximum(a, b), a, equal_nan=True)) + def test_floor(self): x = mx.array([-22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf]) expected = [-23, 19, -27, 9, 0, -np.inf, np.inf] @@ -977,7 +1012,7 @@ def test_median(self): with self.assertRaises(ValueError): mx.median(x, axis=0) x = mx.array([0, 1, 2, 3, 4]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.median(x, axis=(0, 1)) with self.assertRaises(ValueError): mx.median(x, axis=(0, 0)) @@ -1562,9 +1597,9 @@ def test_along_axis_invalid_axis(self): values = mx.ones(a.shape, dtype=a.dtype) for ax in [3, 4, 100, -4, -5, -100]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.take_along_axis(a, idx, axis=ax) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.put_along_axis(a, idx, values, axis=ax) # Valid negative axes still work @@ -1576,7 +1611,7 @@ def test_cross_invalid_axis(self): a = mx.array([1.0, 2.0, 3.0]) b = mx.array([4.0, 5.0, 6.0]) for ax in [1, 2, -2, -50]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.linalg.cross(a, b, axis=ax) def test_put_along_axis(self): @@ -1646,7 +1681,7 @@ def test_split(self): self.assertEqual(y.tolist(), [[3, 4]]) self.assertEqual(z.tolist(), [[5, 6]]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.split(a, 3, axis=2) a = mx.arange(8) @@ -1668,12 +1703,17 @@ def test_flip(self): b_np = np.array([1, 2, 3, 4]) self.assertTrue(np.array_equal(mx.flip(mx.array(b_np)), np.flip(b_np))) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.flip(a, axis=2) def test_unstack(self): a_np = np.arange(6).reshape(3, 2) a = mx.array(a_np) + + self.assertIsInstance(mx.unstack(a), tuple) + self.assertIsInstance(mx.unstack(a, axis=1), tuple) + self.assertEqual(mx.unstack(mx.zeros((0, 2)), axis=0), ()) + for axis in [0, 1, -1]: parts = mx.unstack(a, axis=axis) expected = np.unstack(a_np, axis=axis) @@ -1688,7 +1728,7 @@ def test_unstack(self): # stack is the inverse of unstack. self.assertTrue(mx.array_equal(mx.stack(mx.unstack(a, axis=1), axis=1), a)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.unstack(a, axis=2) def test_split_invalid_num_splits(self): @@ -2568,6 +2608,23 @@ def test_logcumsumexp(self): c_mlx = mxop(a_mlx, axis=-1) self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3)) + def test_scan_invalid_axis(self): + a = mx.arange(24).reshape(2, 3, 4) + + for op in ["cumsum", "cumprod", "cummax", "cummin", "logcumsumexp"]: + mxop = getattr(mx, op) + for ax in [3, 4, 100, -4, -5, -100]: + with self.assertRaises(IndexError): + mxop(a, axis=ax) + + # Valid negative axes still work and agree with the positive one + # logcumsumexp has no integer kernel, so use a float input + a_ = a.astype(mx.float32) if op == "logcumsumexp" else a + for ax in [-1, -2, -3]: + out_neg = mxop(a_, axis=ax) + out_pos = mxop(a_, axis=ax + a_.ndim) + self.assertTrue(mx.array_equal(out_neg, out_pos)) + def test_scans(self): a_npy = np.random.randn(32, 32, 32).astype(np.float32) a_mlx = mx.array(a_npy) @@ -2677,6 +2734,14 @@ def fn(its): mem4 = mx.get_peak_memory() self.assertEqual(mem2, mem4) + # Scanning an empty axis is a no-op + for a in (mx.array([]), mx.zeros((2, 0)), mx.zeros((0, 2))): + for op in ("cumsum", "cumprod", "cummax", "cummin", "logcumsumexp"): + for axis in range(a.ndim): + out = getattr(mx, op)(a, axis=axis) + mx.eval(out) + self.assertEqual(out.shape, a.shape) + def test_scan_size_one_axis(self): # A size one axis can carry any stride and still be row contiguous, so # the scan must not take its row count from that stride. @@ -2753,7 +2818,7 @@ def test_diff(self): self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]]) self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]]) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.diff(a, axis=1) def test_squeeze_expand(self): @@ -2777,19 +2842,64 @@ def test_squeeze_expand_invalid_axes(self): # Out of bounds negative axes must raise instead of wrapping around a = mx.zeros(()) self.assertEqual(mx.expand_dims(a, (-2, -1)).shape, (1, 1)) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.expand_dims(a, (-3, -2)) a = mx.zeros((2, 2)) for axes in [(-5, -4), (-6, 0), (0, 5)]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.expand_dims(a, axes) a = mx.zeros((1, 1, 1)) for axes in [(-4,), (-5, 0), (0, 4)]: - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.squeeze(a, axes) + def test_out_of_bounds_axis_raises_index_error(self): + # An out of bounds axis is an indexing error, not a value error, so + # every op agrees on IndexError (numpy raises AxisError, which is a + # subclass of both). + a = mx.zeros((2, 3)) + cases = [ + lambda: mx.sum(a, axis=5), + lambda: mx.mean(a, axis=5), + lambda: mx.median(a, axis=5), + lambda: mx.expand_dims(a, 5), + lambda: mx.squeeze(mx.zeros((1, 1)), 5), + lambda: mx.moveaxis(a, 5, 0), + lambda: mx.swapaxes(a, 5, 0), + lambda: mx.transpose(a, (5, 0)), + lambda: mx.roll(a, 1, 5), + lambda: mx.diagonal(a, 0, 5, 1), + lambda: mx.take(a, mx.array([0]), axis=5), + lambda: mx.vecdot(a, a, axis=5), + lambda: mx.diff(a, axis=5), + lambda: mx.take_along_axis(a, mx.zeros((2, 3), mx.int32), axis=5), + lambda: mx.cumsum(a, axis=5), + lambda: mx.argsort(a, axis=5), + lambda: mx.split(a, 2, axis=5), + lambda: mx.flip(a, axis=5), + lambda: mx.fft.fft(a, axis=5), + lambda: mx.fft.fftshift(a, axes=[5]), + lambda: mx.random.categorical(a, axis=5), + lambda: mx.unflatten(a, 5, (1, -1)), + lambda: mx.slice(a, mx.array([0]), [5], (1,)), + lambda: mx.slice_update(a, mx.zeros((1, 1)), mx.array([0]), [5]), + lambda: mx.vmap(mx.exp, in_axes=5)(a), + ] + for case in cases: + with self.assertRaises(IndexError): + case() + + # Duplicate or mismatched axes stay ValueError, they are not a + # bounds problem. + with self.assertRaises(ValueError): + mx.transpose(a, (0, 0)) + with self.assertRaises(ValueError): + mx.transpose(a, (0,)) + with self.assertRaises(ValueError): + mx.sum(a, axis=(0, 0)) + def test_sort(self): shape = (6, 4, 10) tests = product( @@ -2936,6 +3046,16 @@ def test_sort(self): b_mx = mx.partition(a_mx, 1, axis=-1) self.assertTrue(np.array_equal(b_np[:, 1], np.array(b_mx)[:, 1])) + # Sorting an empty axis is a no-op + for a in (mx.array([]), mx.zeros((2, 0)), mx.zeros((0, 2))): + for axis in range(a.ndim): + out = mx.sort(a, axis=axis) + mx.eval(out) + self.assertEqual(out.shape, a.shape) + out = mx.argsort(a, axis=axis) + mx.eval(out) + self.assertEqual(out.shape, a.shape) + def test_partition(self): shape = (3, 4, 5) for dtype in ("int32", "float32"): @@ -3348,7 +3468,7 @@ def test_vecdot(self): with self.assertRaises(ValueError): mx.vecdot(mx.array(1), mx.array([1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vecdot(mx.array([1, 2]), mx.array([1, 2]), axis=1) with self.assertRaises(ValueError): mx.vecdot(mx.array([1, 2]), mx.array([1])) @@ -4044,6 +4164,16 @@ def test_complex_ops(self): mx_op = getattr(mx, op) self.assertTrue(np.allclose(mx_op(x), np_op(x))) + if mx.metal.is_available() and mx.default_device() == mx.gpu: + x = mx.array( + [1e20 + 1e20j, 2.5e-20 - 3e-20j, 3e38 + 0j, 1e-19j], mx.complex64 + ) + for op in ["abs", "log"]: + with self.subTest(op=op): + np_op = getattr(np, op) + mx_op = getattr(mx, op) + self.assertTrue(np.allclose(mx_op(x), np_op(x), rtol=1e-5)) + x = mx.array( [ 3.0 + 4.0j, @@ -4093,6 +4223,14 @@ def test_integer_power(self): x = mx.power(mx.array([2, -2, 1], dtype), mx.array([-1, -3, -9], dtype)) self.assertEqual(x.tolist(), [0, 0, 0]) + # A single negative exponent must only zero its own element, not the + # whole SIMD vector. Use enough elements to span a vector lane. + with mx.stream(mx.cpu): + base = mx.array([2] * 16, mx.int32) + exp = mx.array([3] * 3 + [-1] + [3] * 12, mx.int32) + expected = [8] * 3 + [0] + [8] * 12 + self.assertEqual(mx.power(base, exp).tolist(), expected) + def test_depends(self): a = mx.array([1.0, 2.0, 3.0]) b = mx.exp(a) @@ -4216,6 +4354,11 @@ def test_sort_nan(self): expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, mx.nan + 2j]) self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) + # A NaN in the imaginary part also sorts last. + x = mx.array([3.0 + 1j, complex(1.0, mx.nan), 2.0 + 1j, 0.0 + 1j]) + expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, complex(1.0, mx.nan)]) + self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) + def test_argsort_nan(self): for dtype in [mx.float32, mx.float16, mx.bfloat16]: with self.subTest(dtype=dtype): @@ -4225,6 +4368,13 @@ def test_argsort_nan(self): sorted_x = mx.take(x, indices) self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True)) + # Also test complex values + for nan_val in [mx.nan + 2j, complex(1.0, mx.nan)]: + x = mx.array([3.0 + 1j, nan_val, 2.0 + 1j, 0.0 + 1j]) + expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, nan_val]) + sorted_x = mx.take(x, mx.argsort(x)) + self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True)) + def test_to_from_fp8(self): vals = mx.array( [448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625] @@ -4232,6 +4382,17 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals)) self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) + # 0x7f and 0xff are NaN in E4M3FN. Arithmetic does not keep the sign of + # a NaN, so only test for NaN. Use float32, numpy has no bfloat16. + nan_encodings = mx.array([0x7F, 0xFF], dtype=mx.uint8) + decoded_np = np.array(mx.from_fp8(nan_encodings, mx.float32)) + self.assertTrue(np.isnan(decoded_np[0]).item()) + self.assertTrue(np.isnan(decoded_np[1]).item()) + + # The carry must not fire one byte early: 0x7e/0xfe stay finite + finite = mx.from_fp8(mx.array([0x7E, 0xFE], dtype=mx.uint8), mx.float32) + self.assertTrue(mx.array_equal(finite, mx.array([448.0, -448.0]))) + def test_zeros_ones_empty_like_dtype(self): x = mx.array([1, 2, 3], dtype=mx.int32) diff --git a/Source/Cmlx/mlx/python/tests/test_optimizers.py b/Source/Cmlx/mlx/python/tests/test_optimizers.py index 3d51fedf3..c7caf5b95 100644 --- a/Source/Cmlx/mlx/python/tests/test_optimizers.py +++ b/Source/Cmlx/mlx/python/tests/test_optimizers.py @@ -364,6 +364,29 @@ def test_muon(self): optim_no_momentum = opt.Muon(learning_rate=1e-2, momentum=0.0) optim_no_momentum.apply_gradients(grads, params) + def test_weight_decay_keeps_inputs_unchanged(self): + # Weight decay must not write into the gradients or the parameters + # that the caller passed in. + params = {"w": mx.ones((3, 3))} + grads = {"w": mx.full((3, 3), 2.0)} + + optimizer = opt.SGD(learning_rate=1e-2, weight_decay=0.1) + optimizer.apply_gradients(grads, params) + self.assertTrue(mx.array_equal(grads["w"], mx.full((3, 3), 2.0))) + self.assertTrue(mx.array_equal(params["w"], mx.ones((3, 3)))) + + optimizer = opt.Adafactor(learning_rate=1e-2, weight_decay=0.1) + optimizer.apply_gradients(grads, params) + self.assertTrue(mx.array_equal(grads["w"], mx.full((3, 3), 2.0))) + self.assertTrue(mx.array_equal(params["w"], mx.ones((3, 3)))) + + # The same gradients applied twice give the same update + optimizer = opt.SGD(learning_rate=1e-2, weight_decay=0.1) + first = optimizer.apply_gradients(grads, params) + optimizer = opt.SGD(learning_rate=1e-2, weight_decay=0.1) + second = optimizer.apply_gradients(grads, params) + self.assertTrue(mx.array_equal(first["w"], second["w"])) + def test_compiled_optimizer(self): model = nn.Linear(10, 10) x = mx.random.uniform(shape=(2, 10)) diff --git a/Source/Cmlx/mlx/python/tests/test_quantized.py b/Source/Cmlx/mlx/python/tests/test_quantized.py index 461175f01..7a5547411 100644 --- a/Source/Cmlx/mlx/python/tests/test_quantized.py +++ b/Source/Cmlx/mlx/python/tests/test_quantized.py @@ -320,6 +320,52 @@ def test_qqmm(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(), "Global scale is only supported on Metal backend" + ) + def test_qqmm_global_scale_matrix_kernels(self): + mx.random.seed(0) + dtype = mx.bfloat16 + for M, N, K in product( + [32, 64], + [128, 130], + [96, 128], + ): + with self.subTest(shape=(M, N, K)): + x = mx.random.normal((M, K)).astype(dtype) + w = mx.random.normal((N, K)).astype(dtype) + global_scale_x = mx.max(mx.abs(x)).astype(mx.float32) + global_scale_w = mx.max(mx.abs(w)).astype(mx.float32) + x_hat = mx.dequantize( + *mx.quantize(x, mode="nvfp4", global_scale=global_scale_x), + mode="nvfp4", + dtype=dtype, + global_scale=global_scale_x, + ) + w_q, scales = mx.quantize(w, mode="nvfp4", global_scale=global_scale_w) + w_hat = mx.dequantize( + w_q, + scales, + mode="nvfp4", + dtype=dtype, + global_scale=global_scale_w, + ) + + expected = x_hat @ mx.swapaxes(w_hat, -1, -2) + actual = mx.qqmm( + x, + w_q, + scales, + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + ) + delta = mx.abs(actual.astype(mx.float32) - expected.astype(mx.float32)) + relative_error = delta.max() / mx.maximum( + mx.abs(expected.astype(mx.float32)).max(), 1e-20 + ) + self.assertLess(relative_error, 3e-2) + def test_qmm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -909,6 +955,167 @@ 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} + + def _fp_tail_dtypes(self): + # The half types only run on the GPU + if mx.default_device() == mx.gpu: + return [mx.float32, mx.float16, mx.bfloat16] + return [mx.float32] + + def test_fp_qmm_non_multiple_of_32(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + + def check(M, K, N, transpose, mode, dtype): + with self.subTest( + M=M, K=K, N=N, transpose=transpose, mode=mode, dtype=dtype + ): + w_shape = (N, K) if transpose else (K, N) + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + + y_q = mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode=mode + ) + # An unbounded store of the N tail spills onto the next output + # row and races the threadgroup that owns it, so two identical + # runs must agree bit for bit. + y_q2 = mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode=mode + ) + y_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat) + + self.assertEqual(y_q.shape, y_hat.shape) + self.assertTrue(mx.array_equal(y_q, y_q2)) + self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype]) + + # (M, K, N, transpose). M is kept clear of the qmv batch limit + # (get_qmv_batch_limit caps at 32) so these route to the tiled kernel. + nvfp4_shapes = [ + (50, 1040, 128, True), # K % 32 == 16 + (33, 528, 64, True), # K % 32 == 16, M just over one BM tile + (50, 512, 1040, False), # N % 32 == 16 + (50, 1040, 1040, False), # both dims carry a tail + (50, 16, 128, True), # K is a single partial tile + ] + control_shapes = [ + (50, 1056, 128, True), # K % 32 == 0 + (50, 512, 1056, False), # N % 32 == 0 + ] + + for dtype in self._fp_tail_dtypes(): + for M, K, N, transpose in nvfp4_shapes: + check(M, K, N, transpose, "nvfp4", dtype) + for mode in ["mxfp4", "mxfp8"]: + for M, K, N, transpose in control_shapes: + check(M, K, N, transpose, mode, dtype) + + def test_fp_qmm_non_multiple_of_32_vjp(self): + # The backward of quantized_matmul runs the same kernels with the + # transpose flipped, so cover both directions at a tail shape. + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + M, K, N = 50, 1040, 128 + + for dtype in self._fp_tail_dtypes(): + for transpose in [True, False]: + with self.subTest(transpose=transpose, dtype=dtype): + w_shape = (N, K) if transpose else (K, N) + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode="nvfp4") + w_hat = mx.dequantize(w_q, scales, mode="nvfp4") + cot = (mx.random.normal(shape=(M, N), key=k3) / N**0.5).astype( + dtype + ) + + def f(x): + return mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode="nvfp4" + ) + + (out,), (grad,) = mx.vjp(f, (x,), (cot,)) + tol = self.FP_TAIL_TOLS[dtype] + + # The forward runs the tail kernel for this transpose; the + # backward runs it for the other one, we might as well check both. + out_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat) + self.assertEqual(out.shape, out_hat.shape) + self.assertLess((out - out_hat).abs().max(), tol) + + grad_hat = cot @ ( + w_hat if transpose else mx.swapaxes(w_hat, -1, -2) + ) + self.assertEqual(grad.shape, grad_hat.shape) + self.assertLess((grad - grad_hat).abs().max(), tol) + + def test_fp_gather_qmm_non_multiple_of_32(self): + # Both gather_qmm kernels: the sorted-indices (MoE prefill) kernel and + # the tiled gather kernel. + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + mode = "nvfp4" + E = 4 + + def check(K, N, transpose, sorted_indices, dtype): + with self.subTest( + K=K, + N=N, + transpose=transpose, + sorted_indices=sorted_indices, + dtype=dtype, + ): + w_shape = (E, N, K) if transpose else (E, K, N) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + + if sorted_indices: + T = 64 + x = (mx.random.normal(shape=(T, 1, K), key=k1) / K**0.5).astype( + dtype + ) + rhs = mx.sort(mx.random.randint(0, E, shape=(T,), key=k3)).astype( + mx.uint32 + ) + lhs = None + else: + B, M = 8, 50 + x = (mx.random.normal(shape=(B, M, K), key=k1) / K**0.5).astype( + dtype + ) + rhs = mx.random.randint(0, E, shape=(B,), key=k3).astype(mx.uint32) + lhs = mx.arange(B, dtype=mx.uint32) + + y_q = mx.gather_qmm( + x, + w_q, + scales, + lhs_indices=lhs, + rhs_indices=rhs, + transpose=transpose, + mode=mode, + sorted_indices=sorted_indices, + ) + w_sel = mx.take(w_hat, rhs, axis=0) + y_hat = x @ (mx.swapaxes(w_sel, -1, -2) if transpose else w_sel) + + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype]) + + # (K, N, transpose): transpose puts the tail on K, else on N. + shapes = [ + (1040, 128, True), # K % 32 == 16 + (512, 1040, False), # N % 32 == 16 + ] + for dtype in self._fp_tail_dtypes(): + for K, N, transpose in shapes: + for sorted_indices in [True, False]: + check(K, N, transpose, sorted_indices, dtype) + def test_mode_error_cases(self): w = mx.random.normal(shape=(256, 256)) x = mx.random.normal(shape=(1, 256)) @@ -1394,6 +1601,73 @@ def test_gather_qqmm(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_gather_qqmm_global_scale_matrix_paths(self): + if mx.default_device() == mx.cpu: + self.skipTest("Not implemented for CPU") + + E, N, K = 3, 64, 256 + w = mx.random.normal((E, N, K), key=mx.random.key(10)).astype(mx.bfloat16) + global_scale_w = mx.max(mx.abs(w), axis=(1, 2)).astype(mx.float32) + quantized = [ + mx.quantize(w[e], mode="nvfp4", global_scale=global_scale_w[e]) + for e in range(E) + ] + w_q = mx.stack([q for q, _ in quantized]) + scales_w = mx.stack([scales for _, scales in quantized]) + w_hat = mx.stack( + [ + mx.dequantize( + q, + scales, + mode="nvfp4", + dtype=mx.bfloat16, + global_scale=global_scale_w[e], + ) + for e, (q, scales) in enumerate(quantized) + ] + ) + + def check(x, rhs, lhs=None, sorted_indices=False): + global_scale_x = mx.max(mx.abs(x)).astype(mx.float32) + x_hat = mx.dequantize( + *mx.quantize(x, mode="nvfp4", global_scale=global_scale_x), + mode="nvfp4", + dtype=mx.bfloat16, + global_scale=global_scale_x, + ) + actual = mx.gather_qqmm( + x, + w_q, + scales_w, + lhs, + rhs, + mode="nvfp4", + global_scale_x=global_scale_x, + global_scale_w=global_scale_w, + sorted_indices=sorted_indices, + ) + expected = mx.gather_mm( + x_hat, + mx.swapaxes(w_hat, -1, -2), + lhs, + rhs, + sorted_indices=sorted_indices, + ) + 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) + + # Sorted indices select the RHS-grouped matrix kernel. + x = mx.random.normal((24, 1, K), key=mx.random.key(20)).astype(mx.bfloat16) + rhs = mx.array([0] * 8 + [1] * 8 + [2] * 8) + check(x, rhs, sorted_indices=True) + + # A large M selects the gathered matrix kernel. + x = mx.random.normal((2, 64, K), key=mx.random.key(30)).astype(mx.bfloat16) + lhs = mx.array([0, 1, 0, 1]) + rhs = mx.array([0, 1, 2, 0]) + check(x, rhs, lhs) + def test_qmm_fp_type(self): indices = mx.array([[2], [0], [1]], dtype=mx.uint32) @@ -1599,6 +1873,49 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y3, atol=tol)) self.assertTrue(mx.allclose(y1, y4, atol=tol)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_gather_qmm_sorted_nax_large_m(self): + E, N, K, group_size = 16, 64, 64, 32 + dtype = mx.float16 + mx.random.seed(0) + w = (mx.random.normal((E, N, K)) * 0.1).astype(dtype) + w_q, scales, biases = mx.quantize(w, group_size=group_size, bits=4) + w_hat = mx.dequantize(w_q, scales, biases, group_size=group_size, bits=4) + + for M in (32767, 32768, 32769, 32832): + with self.subTest(M=M): + x = (mx.random.normal((M, 1, K)) * 0.1).astype(dtype) + rhs_indices = (mx.arange(M) * E // M).astype(mx.uint32) + y_hat = mx.gather_mm( + x.astype(mx.float32), + mx.swapaxes(w_hat, -1, -2).astype(mx.float32), + rhs_indices=rhs_indices, + sorted_indices=True, + ) + mx.eval(y_hat) + mx.synchronize() + + for value in (-31.0, 47.0): + poison = mx.full(y_hat.shape, value, dtype=mx.float16) + mx.eval(poison) + mx.synchronize() + del poison + + y_q = mx.gather_qmm( + x, + w_q, + scales, + biases, + rhs_indices=rhs_indices, + transpose=True, + group_size=group_size, + bits=4, + sorted_indices=True, + ) + max_error = (y_q.astype(mx.float32) - y_hat).abs().max() + self.assertLess(float(max_error.item()), 5e-2) + del y_q, max_error + @unittest.skipIf(mx.cuda.is_available(), "Not implemented for CUDA") def test_gather_qmm_sorted_sliced_weight(self): E, R, D, N = 8, 64, 256, 64 @@ -1627,6 +1944,49 @@ def test_gather_qmm_sorted_sliced_weight(self): ) ) + def test_gather_qmm_sorted_unaligned_k(self): + # Every case in test_gather_qmm_sorted uses K % 64 == 0, so the tail of + # the K loop is never exercised. When K is not a multiple of the block + # size the sorted path has to zero the partial tile correctly, otherwise + # it silently drops the tail contribution for part of the output. + L, D, E = 128, 256, 4 + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + on_gpu = mx.default_device() == mx.gpu + + dtype = mx.float16 if on_gpu else mx.float32 + + for mode, group_size in ( + ("affine", 32), + ("mxfp4", None), + ("mxfp8", None), + ("nvfp4", None), + ): + for K in (160, 288, 544): + with self.subTest(mode=mode, K=K): + indices = mx.sort( + (mx.random.uniform(shape=(L,), key=k1) * E).astype(mx.uint32) + ) + x = (mx.random.normal((L, 1, K), key=k2) / K**0.5).astype(dtype) + w = (mx.random.normal((E, D, K), key=k3) / K**0.5).astype(dtype) + + if mode == "affine": + wq = mx.quantize(w, group_size=group_size, mode=mode) + else: + wq = mx.quantize(w, mode=mode) + + kwargs = dict( + group_size=group_size, + mode=mode, + transpose=True, + rhs_indices=indices, + ) + y_sorted = mx.gather_qmm(x, *wq, sorted_indices=True, **kwargs) + y_unsorted = mx.gather_qmm(x, *wq, sorted_indices=False, **kwargs) + + tol = 1e-3 if on_gpu else 1.5e-5 + self.assertLess((y_sorted - y_unsorted).abs().max(), tol) + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: @@ -1728,6 +2088,104 @@ def 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" + ) + def test_gather_qmm_global_scale(self): + mx.random.seed(0) + N, K = 128, 256 + + def rel_err(got, expected): + # Per row + d = mx.abs(got.astype(mx.float32) - expected.astype(mx.float32)) + scale = mx.abs(expected.astype(mx.float32)) + axes = tuple(range(1, d.ndim)) + return (d.max(axis=axes) / mx.maximum(scale.max(axis=axes), 1e-20)).max() + + def quantize_experts(w): + """One tensor scale per expert.""" + E = w.shape[0] + gs = mx.stack([mx.abs(w[e]).max().astype(mx.float32) for e in range(E)]) + qs = [mx.quantize(w[e], mode="nvfp4", global_scale=gs[e]) for e in range(E)] + w_hat = mx.stack( + [ + mx.dequantize( + q, sc, mode="nvfp4", global_scale=gs[e], dtype=w.dtype + ) + for e, (q, sc) in enumerate(qs) + ] + ) + return ( + mx.stack([q for q, _ in qs]), + mx.stack([sc for _, sc in qs]), + gs, + w_hat, + ) + + tests = product( + [4, 128, 234], # E + [mx.float32, mx.float16, mx.bfloat16], # dtype + [True, False], # transpose + [(32, 2, False), (1, 2, False), (256, 4, True)], # M, B, sort + ) + for E, dtype, transpose, (M, B, sort) in tests: + with self.subTest(E=E, dtype=dtype, transpose=transpose, M=M, B=B): + # Each expert has a different scale, so we multiply by a factor + # to make the experts have different magnitudes. + factors = mx.array( + [[1, 2, 4, 8][e % 4] for e in range(E)], mx.float32 + ).reshape((E, 1, 1)) + wshape = (E, N, K) if transpose else (E, K, N) + w = (mx.random.normal(wshape) * factors).astype(dtype) + wq, s, gs, w_hat = quantize_experts(w) + self.assertEqual(gs.shape, (E,)) + + x = mx.random.normal((B, M, K)).astype(dtype) + indices = mx.random.randint(0, E, (B,)) + if sort: + indices = mx.sort(indices) + + wg = w_hat[indices] + expected = x @ (wg.swapaxes(-1, -2) if transpose else wg) + kwargs = dict( + rhs_indices=indices, + transpose=transpose, + mode="nvfp4", + sorted_indices=sort, + ) + + out = mx.gather_qmm(x, wq, s, global_scale=gs, **kwargs) + tol = 1e-5 if dtype == mx.float32 else 3e-2 + self.assertLess(rel_err(out, expected), tol) + + # Each expert uses its own scale, not a neighbour's + rotated = mx.concatenate([gs[1:], gs[:1]]) + wrong = mx.gather_qmm(x, wq, s, global_scale=rotated, **kwargs) + self.assertGreater(rel_err(wrong, expected), 0.5) + + # Repeated and out-of-order experts with explicit lhs indices on the + # M=1 vector path; each output must pick its own expert's scale. The + # long K is the decode shape that long-K vector kernels take. + E = 6 + for K in (K, 2048): + with self.subTest(K=K): + w = (mx.random.normal((E, N, K)) / K**0.5).astype(mx.bfloat16) + wq, s, gs, w_hat = quantize_experts(w) + indices = mx.array([3, 0, 5, 3, 1, 0], mx.uint32) + x = (mx.random.normal((2, 1, K)) / K**0.5).astype(mx.bfloat16) + lhs = mx.array([0, 1, 0, 1, 0, 1], mx.uint32) + expected = x[lhs] @ w_hat[indices].swapaxes(-1, -2) + out = mx.gather_qmm( + x, + wq, + s, + lhs_indices=lhs, + rhs_indices=indices, + mode="nvfp4", + global_scale=gs, + ) + self.assertLess(rel_err(out, expected), 3e-2) + def test_quantize_strided(self): N = 64 mode = "nvfp4" diff --git a/Source/Cmlx/mlx/python/tests/test_vmap.py b/Source/Cmlx/mlx/python/tests/test_vmap.py index b050a8038..7538b559d 100644 --- a/Source/Cmlx/mlx/python/tests/test_vmap.py +++ b/Source/Cmlx/mlx/python/tests/test_vmap.py @@ -9,8 +9,8 @@ class TestVmap(mlx_tests.MLXTestCase): def test_basics(self): - # Can't vmap over scalars - with self.assertRaises(ValueError): + # Can't vmap over scalars, axis 0 is out of bounds for a 0d array + with self.assertRaises(IndexError): mx.vmap(mx.exp)(mx.array(1.0)) # Invalid input @@ -21,13 +21,13 @@ def test_basics(self): with self.assertRaises(ValueError): mx.vmap(mx.exp, in_axes="hello")(mx.array([0, 1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vmap(mx.exp, in_axes=2)(mx.array([0, 1])) with self.assertRaises(ValueError): mx.vmap(mx.exp, out_axes="hello")(mx.array([0, 1])) - with self.assertRaises(ValueError): + with self.assertRaises(IndexError): mx.vmap(mx.exp, out_axes=2)(mx.array([0, 1])) def test_unary(self): diff --git a/Source/Cmlx/mlx/python/tests/test_zero_copy.py b/Source/Cmlx/mlx/python/tests/test_zero_copy.py index 231e2e45e..248e93a76 100644 --- a/Source/Cmlx/mlx/python/tests/test_zero_copy.py +++ b/Source/Cmlx/mlx/python/tests/test_zero_copy.py @@ -1,6 +1,8 @@ # Copyright © 2024 Apple Inc. +import faulthandler import gc +import threading import unittest import mlx.core as mx @@ -83,6 +85,66 @@ def test_adopt_in_loop_not_recycled(self): mx.eval(r) self.assertTrue(True) # reaching here without crashing is the assertion + def _adopted_source(self, n): + # A square source buffer that mx.asarray can adopt, else a skip. + if not mx.metal.is_available(): + self.skipTest("copy=False requires Metal") + a = np.zeros((n, n), dtype=np.float32) + if a.ctypes.data % 16384 != 0: + self.skipTest("source buffer not page-aligned; adopt path not taken") + return a + + @staticmethod + def _submit_work(a): + # The adopted array is an input of the first matmul and is dropped here, + # so only the completion handler of the command buffer that is in flight + # keeps its Python owner. The caller must keep the output until the end. + n = a.shape[0] + w = mx.ones((n, n)) + y = mx.asarray(a, copy=False) @ w + for _ in range(4): + y = y @ w + mx.async_eval(y) + return y + + # A stream callback can free an adopted buffer, which takes the GIL. A call + # that waits for a stream must release the GIL, else the two deadlock. + # faulthandler reports such a deadlock, its timer is a C thread. A watchdog + # in Python would never run. + def test_synchronize_releases_gil(self): + a = self._adopted_source(1024) + faulthandler.dump_traceback_later(30, exit=True) + try: + for _ in range(4): + y = self._submit_work(a) + mx.synchronize() + del y + finally: + faulthandler.cancel_dump_traceback_later() + + def test_clear_streams_releases_gil(self): + a = self._adopted_source(1024) + errors = [] + + # clear_streams destroys the streams of the calling thread, so it runs + # in a thread that ends right after. + def worker(): + try: + y = self._submit_work(a) + mx.clear_streams() + del y + except Exception as e: + errors.append(e) + + faulthandler.dump_traceback_later(30, exit=True) + try: + thread = threading.Thread(target=worker) + thread.start() + thread.join() + finally: + faulthandler.cancel_dump_traceback_later() + self.assertEqual(errors, []) + if __name__ == "__main__": unittest.main() diff --git a/Source/Cmlx/mlx/setup.py b/Source/Cmlx/mlx/setup.py index 52cc4f749..ad258d8af 100644 --- a/Source/Cmlx/mlx/setup.py +++ b/Source/Cmlx/mlx/setup.py @@ -339,6 +339,7 @@ def is_backend_file(file): "nvidia-cublas-cu12==12.9.*", "nvidia-cufft-cu12==11.4.*", "nvidia-cuda-nvrtc-cu12==12.9.*", + "nvidia-cusolver-cu12==11.7.*", ] elif toolkit == 13: install_requires += [ @@ -346,6 +347,7 @@ def is_backend_file(file): "nvidia-cufft==12.*", "nvidia-cuda-nvrtc==13.*", "nvidia-cuda-runtime==13.*", + "nvidia-cusolver==12.*", ] else: raise ValueError(f"Unknown toolkit {toolkit}") diff --git a/Source/Cmlx/mlx/tests/CMakeLists.txt b/Source/Cmlx/mlx/tests/CMakeLists.txt index 859215acb..aa0d90105 100644 --- a/Source/Cmlx/mlx/tests/CMakeLists.txt +++ b/Source/Cmlx/mlx/tests/CMakeLists.txt @@ -7,11 +7,16 @@ FetchContent_MakeAvailable(doctest) add_executable(tests ${PROJECT_SOURCE_DIR}/tests/tests.cpp) if(MLX_BUILD_METAL OR MLX_BUILD_CUDA) - set(METAL_TEST_SOURCES gpu_tests.cpp) + set(GPU_TEST_SOURCES gpu_tests.cpp) +endif() + +if(MLX_BUILD_CUDA) + list(APPEND GPU_TEST_SOURCES cuda_tests.cpp) + target_link_libraries(tests PRIVATE CUDA::cudart) endif() if(MLX_BUILD_METAL) - list(APPEND METAL_TEST_SOURCES residency_tests.cpp) + list(APPEND GPU_TEST_SOURCES residency_tests.cpp) endif() include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake) @@ -38,7 +43,7 @@ target_sources( utils_tests.cpp vmap_tests.cpp linalg_tests.cpp - ${METAL_TEST_SOURCES}) + ${GPU_TEST_SOURCES}) target_link_libraries(tests PRIVATE mlx doctest) target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS}) diff --git a/Source/Cmlx/mlx/tests/allocator_tests.cpp b/Source/Cmlx/mlx/tests/allocator_tests.cpp index 9658e29f0..d528f8ea2 100644 --- a/Source/Cmlx/mlx/tests/allocator_tests.cpp +++ b/Source/Cmlx/mlx/tests/allocator_tests.cpp @@ -11,8 +11,10 @@ #include "mlx/allocator.h" #include "mlx/device.h" #include "mlx/memory.h" +#include "mlx/ops.h" #include "mlx/scheduler.h" #include "mlx/stream.h" +#include "mlx/transforms.h" using namespace mlx::core; @@ -66,6 +68,24 @@ TEST_CASE("test cached allocation keeps capacity") { set_cache_limit(old_limit); } +TEST_CASE("test array buffer size") { + auto a = array({1.0f, 2.0f, 3.0f, 4.0f}); + auto view = reshape(a, {2, 2}); + + CHECK_THROWS_AS(get_array_buffer_size({view}), std::invalid_argument); + + eval(view); + auto a_size = a.buffer_size(); + CHECK_EQ(get_array_buffer_size({}), 0); + CHECK_EQ(get_array_buffer_size({a}), a_size); + CHECK_EQ(get_array_buffer_size({a, a}), a_size); + CHECK_EQ(get_array_buffer_size({a, view}), a_size); + + auto b = array({5.0f, 6.0f}); + CHECK_EQ(get_array_buffer_size({a, b}), a_size + b.buffer_size()); + CHECK_EQ(get_array_buffer_size({array({})}), 0); +} + TEST_CASE("test clear cache synchronizes cpu streams") { if (is_available(Device{Device::gpu})) { return; diff --git a/Source/Cmlx/mlx/tests/arg_reduce_tests.cpp b/Source/Cmlx/mlx/tests/arg_reduce_tests.cpp index 124fd67ed..6bcce81bd 100644 --- a/Source/Cmlx/mlx/tests/arg_reduce_tests.cpp +++ b/Source/Cmlx/mlx/tests/arg_reduce_tests.cpp @@ -1,5 +1,7 @@ // Copyright © 2023 Apple Inc. +#include + #include "doctest/doctest.h" #include "mlx/mlx.h" @@ -187,6 +189,27 @@ TEST_CASE("test arg reduce edge cases") { CHECK_THROWS(argmax(array({}))); } +TEST_CASE("test arg reduce NaN") { + auto d = default_device(); + auto nan = std::numeric_limits::quiet_NaN(); + auto x = array({3.0f, nan, 1.0f, 5.0f, nan, -1.0f}, {2, 3}); + test_arg_reduce_small(d, x, ArgReduce::ArgMin, {2}, 1, {1, 1}); + test_arg_reduce_small(d, x, ArgReduce::ArgMax, {2}, 1, {1, 1}); + + std::vector wide(1024, 0.0f); + wide[17] = nan; + wide[529] = nan; + x = array(wide.data(), {1024}); + test_arg_reduce_small(d, x, ArgReduce::ArgMin, {}, 0, {17}); + test_arg_reduce_small(d, x, ArgReduce::ArgMax, {}, 0, {17}); + + if (d == Device::cpu) { + auto z = array({complex64_t{1.0f, 0.0f}, complex64_t{1.0f, nan}}, {2}); + test_arg_reduce_small(d, z, ArgReduce::ArgMin, {}, 0, {1}); + test_arg_reduce_small(d, z, ArgReduce::ArgMax, {}, 0, {1}); + } +} + TEST_CASE("test arg reduce irregular strides") { auto x = array( {0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5, diff --git a/Source/Cmlx/mlx/tests/array_tests.cpp b/Source/Cmlx/mlx/tests/array_tests.cpp index 68a4bed3a..cc0cf5239 100644 --- a/Source/Cmlx/mlx/tests/array_tests.cpp +++ b/Source/Cmlx/mlx/tests/array_tests.cpp @@ -690,3 +690,60 @@ TEST_CASE("test siblings circular references without eval") { fun(); CHECK(tracker.expired()); } + +// https://github.com/ml-explore/mlx/pull/4453 +TEST_CASE("test siblings circular references released by assignment") { + std::weak_ptr tracker; + auto fun = [&]() { + array sink({3, 4}); + array key({1, 2}); + auto splits = split(key, 2); + { + // Set fake data as a tracker for ArrayDesc's lifetime. + splits[0].set_data(allocator::malloc(0)); + tracker = splits[0].data_shared_ptr(); + } + // Drop both outputs by assigning over them rather than destroying + // them. The siblings hold each other, so unless assignment runs the + // same check the destructor does, nothing ever releases them. + splits[0] = sink; // copy assignment + splits[1] = array({5, 6}); // move assignment + }; + fun(); + CHECK(tracker.expired()); +} + +TEST_CASE("test siblings circular references released by overwrite") { + std::weak_ptr tracker; + auto fun = [&]() { + array key({1, 2}); + auto splits = split(key, 2); + { + // Set fake data as a tracker for ArrayDesc's lifetime. + splits[0].set_data(allocator::malloc(0)); + tracker = splits[0].data_shared_ptr(); + } + // Leave one output as the last reference, then replace it in place. + // This is the path __setitem__ takes, with no compile involved. + array last = splits[0]; + splits.clear(); + last.overwrite_descriptor(array({7, 8})); + }; + fun(); + CHECK(tracker.expired()); +} + +TEST_CASE("test assigning a sibling over the last reference keeps the cycle") { + array key({1, 2}); + auto splits = split(key, 2); + // Leave one output as the last outside reference to the pair, then assign + // its sibling over it. The sibling is reachable only through the array + // being replaced, so it has to be held before the old descriptor is + // checked: checked first, the pair looks unreferenced, its cycle is broken + // and the sibling list the new value lives in is freed before it is read. + array last = splits[0]; + splits.clear(); + last = last.siblings()[0]; + REQUIRE_EQ(last.siblings().size(), 1); + CHECK(array_equal(last, array({2})).item()); +} diff --git a/Source/Cmlx/mlx/tests/autograd_tests.cpp b/Source/Cmlx/mlx/tests/autograd_tests.cpp index 5a0aca080..1ff203e1e 100644 --- a/Source/Cmlx/mlx/tests/autograd_tests.cpp +++ b/Source/Cmlx/mlx/tests/autograd_tests.cpp @@ -749,7 +749,7 @@ TEST_CASE("test op vjps") { TEST_CASE("test gather and take grads") { // Check linear takes auto linear_f = [](array indices) { - auto fun_linear = [&indices](array input) { return take(input, indices); }; + auto fun_linear = [indices](array input) { return take(input, indices); }; return fun_linear; }; @@ -1532,3 +1532,33 @@ TEST_CASE("test masked_scatter autograd") { CHECK(array_equal(grads[1], array({1.f, 1.f}, {2})).item()); } } + +TEST_CASE("test pad vjp with axes subset and negative axes") { + // Pad only the last axis, given as a negative index + { + auto x = ones({2, 3}); + auto fun = [](array in) { + return pad(in, {-1}, {1}, {2}, array(0), "constant"); + }; + auto cotan = reshape(arange(12), {2, 6}); + auto [out, dout] = vjp(fun, x, cotan); + CHECK_EQ(out.shape(), Shape{2, 6}); + CHECK_EQ(dout.shape(), Shape{2, 3}); + auto expected = slice(cotan, {0, 1}, {2, 4}); + CHECK(array_equal(dout, expected).item()); + } + + // Pad sizes must be matched with axes by position, not by axis value + { + auto x = ones({2, 3}); + auto fun = [](array in) { + return pad(in, {1, 0}, {1, 2}, {3, 0}, array(0), "constant"); + }; + auto cotan = reshape(arange(28), {4, 7}); + auto [out, dout] = vjp(fun, x, cotan); + CHECK_EQ(out.shape(), Shape{4, 7}); + CHECK_EQ(dout.shape(), Shape{2, 3}); + auto expected = slice(cotan, {2, 1}, {4, 4}); + CHECK(array_equal(dout, expected).item()); + } +} diff --git a/Source/Cmlx/mlx/tests/cuda_tests.cpp b/Source/Cmlx/mlx/tests/cuda_tests.cpp new file mode 100644 index 000000000..aca39ca88 --- /dev/null +++ b/Source/Cmlx/mlx/tests/cuda_tests.cpp @@ -0,0 +1,111 @@ +// Copyright © 2026 Apple Inc. + +#include "doctest/doctest.h" +#include "mlx/mlx.h" + +#ifdef _WIN32 +#include "mlx/backend/cuda/wddm.h" +#endif + +#include + +#include +#include + +using namespace mlx::core; + +namespace { + +bool memory_pools_supported() { + int supported = 0; + REQUIRE_EQ( + cudaDeviceGetAttribute(&supported, cudaDevAttrMemoryPoolsSupported, 0), + cudaSuccess); + return supported; +} + +struct CacheLimitGuard { + ~CacheLimitGuard() { + set_cache_limit(previous); + } + + size_t previous; +}; + +struct MemoryLimitGuard { + ~MemoryLimitGuard() { + set_memory_limit(previous); + } + + size_t previous; +}; + +} // namespace + +constexpr uint64_t gib = 1ULL << 30; + +#ifdef _WIN32 +TEST_CASE("test memory limit from budget") { + CHECK_EQ( + cu::memory_limit_from_budget(24 * gib, 8 * gib, 6 * gib), + 22 * gib - (24 * gib) / 20); + CHECK_EQ( + cu::memory_limit_from_budget(24 * gib, 4 * gib, 6 * gib), + 24 * gib - (24 * gib) / 20); + CHECK_EQ(cu::memory_limit_from_budget(24 * gib, 25 * gib, 0), 0); +} + +TEST_CASE("test wddm memory limit") { + MemoryLimitGuard memory_limit{ + set_memory_limit(std::numeric_limits::max())}; + cudaDeviceProp default_properties{}; + REQUIRE_EQ(cudaGetDeviceProperties(&default_properties, 0), cudaSuccess); + auto default_limit = get_memory_limit(); + if (default_properties.memoryPoolsSupported && + !default_properties.integrated && !default_properties.tccDriver && + default_properties.luidDeviceNodeMask != 0) { + CHECK_LT(default_limit, std::numeric_limits::max()); + } +} +#endif // _WIN32 + +TEST_CASE("test clear cache trims CUDA pool") { + if (!memory_pools_supported()) { + return; + } + + cudaMemPool_t pool = nullptr; + REQUIRE_EQ(cudaDeviceGetDefaultMemPool(&pool, 0), cudaSuccess); + + CacheLimitGuard cache_limit{set_cache_limit(gib)}; + clear_cache(); + + uint64_t initial_reserved = 0; + REQUIRE_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &initial_reserved), + cudaSuccess); + + auto s = default_stream(Device::gpu); + { + auto a = zeros({16 * 1024 * 1024}, float32, s); + eval(a); + synchronize(s); + } + CHECK_GE(get_cache_memory(), 64ULL << 20); + + uint64_t allocated_reserved = 0; + REQUIRE_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &allocated_reserved), + cudaSuccess); + CHECK_GT(allocated_reserved, initial_reserved); + + clear_cache(); + uint64_t final_reserved = 0; + REQUIRE_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &final_reserved), + cudaSuccess); + CHECK_LT(final_reserved, allocated_reserved); +} diff --git a/Source/Cmlx/mlx/tests/fft_tests.cpp b/Source/Cmlx/mlx/tests/fft_tests.cpp index 7dc1b6e66..6adc4c1d4 100644 --- a/Source/Cmlx/mlx/tests/fft_tests.cpp +++ b/Source/Cmlx/mlx/tests/fft_tests.cpp @@ -134,6 +134,9 @@ TEST_CASE("test real ffts") { TEST_CASE("test fftn") { auto x = zeros({5, 5, 5}); CHECK_THROWS_AS(fft::fftn(x, {}, {0, 3}), std::invalid_argument); + // Matching n/axes sizes so the bounds check is what trips. + CHECK_THROWS_AS(fft::fftn(x, {5, 5}, {0, 3}), std::out_of_range); + CHECK_THROWS_AS(fft::fftn(x, {5, 5}, {0, -4}), std::out_of_range); CHECK_THROWS_AS(fft::fftn(x, {}, {0, -4}), std::invalid_argument); CHECK_THROWS_AS(fft::fftn(x, {}, {0, 0}), std::invalid_argument); CHECK_THROWS_AS(fft::fftn(x, {5, 5, 5}, {0}), std::invalid_argument); @@ -388,8 +391,8 @@ TEST_CASE("test fftshift and ifftshift") { CHECK(array_equal(y, expected).item()); // Test error cases - CHECK_THROWS_AS(fft::fftshift(x, {3}), std::invalid_argument); - CHECK_THROWS_AS(fft::fftshift(x, {-5}), std::invalid_argument); - CHECK_THROWS_AS(fft::ifftshift(x, {3}), std::invalid_argument); - CHECK_THROWS_AS(fft::ifftshift(x, {-5}), std::invalid_argument); + CHECK_THROWS_AS(fft::fftshift(x, {3}), std::out_of_range); + CHECK_THROWS_AS(fft::fftshift(x, {-5}), std::out_of_range); + CHECK_THROWS_AS(fft::ifftshift(x, {3}), std::out_of_range); + CHECK_THROWS_AS(fft::ifftshift(x, {-5}), std::out_of_range); } diff --git a/Source/Cmlx/mlx/tests/gpu_tests.cpp b/Source/Cmlx/mlx/tests/gpu_tests.cpp index 8bef07a61..8f492e52f 100644 --- a/Source/Cmlx/mlx/tests/gpu_tests.cpp +++ b/Source/Cmlx/mlx/tests/gpu_tests.cpp @@ -57,6 +57,24 @@ TEST_CASE("test gpu full") { } } +TEST_CASE("test gpu strided scan grid") { + // Regression for #4419 strided_scan writes out of bounds. + std::vector values(72, 1.0f); + auto x = array(values.data(), {72}); + x = transpose(reshape(x, {2, 9, 4}, Device::gpu), {0, 2, 1}, Device::gpu); + x = cumsum(x, -1, false, true, Device::gpu); + + auto causal = tril(ones({9, 9}, float32, Device::gpu), 0, Device::gpu); + auto relative = subtract( + expand_dims(x, -1, Device::gpu), + expand_dims(x, -2, Device::gpu), + Device::gpu); + eval(relative, causal); + + auto expected = tri(9, 9, 0, float32, Device::cpu); + CHECK(array_equal(causal, expected, Device::cpu).item()); +} + TEST_CASE("test gpu astype") { array x = array({-4, -3, -2, -1, 0, 1, 2, 3}); // Check all types work diff --git a/Source/Cmlx/mlx/tests/load_tests.cpp b/Source/Cmlx/mlx/tests/load_tests.cpp index 6ef7bc276..b4b54046a 100644 --- a/Source/Cmlx/mlx/tests/load_tests.cpp +++ b/Source/Cmlx/mlx/tests/load_tests.cpp @@ -257,6 +257,104 @@ TEST_CASE("test gguf tensor data offset validation") { } } +// Writes a one-tensor GGUF whose tensor header carries `dims` and `type` +// verbatim, so a test can drive the dimension product independently of the +// amount of tensor data in the file. The name ends in ".weight" because the +// quantized loader derives the scales and biases names from it. +void write_raw_gguf_dims( + const std::string& path, + const std::vector& dims, + uint32_t type, + size_t data_bytes) { + std::ofstream out(path, std::ios::binary); + auto u32 = [&out](uint32_t v) { + out.write(reinterpret_cast(&v), 4); + }; + auto u64 = [&out](uint64_t v) { + out.write(reinterpret_cast(&v), 8); + }; + out.write("GGUF", 4); + u32(3); // version + u64(1); // tensor_count + u64(0); // metadata_kv_count + u64(8); // tensor name length + out.write("w.weight", 8); + u32(dims.size()); + for (auto dim : dims) { + u64(dim); + } + u32(type); + u64(0); // tensor data offset + while (out.tellp() % 32 != 0) { // default GGUF alignment + out.put(0); + } + std::vector data(data_bytes, 0); + out.write(data.data(), data.size()); +} + +TEST_CASE("test gguf tensor dimension validation") { + // The element count exists in two widths: gguflib keeps the 64 bit product in + // tensor.num_weights, and get_shape narrows each dimension to a 32 bit + // ShapeElem. Only the byte size is checked against the file, so dimensions + // that make the two disagree size a buffer from one and index it with the + // other. See ml-explore/mlx#4244. + const uint32_t q8_0 = 8; + + SUBCASE("valid quantized tensor loads") { + std::string file_path = get_temp_file("test_gguf_dims_ok.gguf"); + // One block of 32 weights: 2 bytes of scale plus 32 bytes of weights. + write_raw_gguf_dims(file_path, {32}, q8_0, 34); + auto [weights, metadata] = load_gguf(file_path); + CHECK_EQ(weights.at("w.weight").shape(), Shape{8}); + CHECK_EQ(weights.at("w.scales").shape(), Shape{1}); + CHECK_EQ(weights.at("w.biases").shape(), Shape{1}); + } + + SUBCASE("dimension past the ShapeElem range") { + // Narrows to a shape of 32 that no longer describes the file. + std::string file_path = get_temp_file("test_gguf_dim_narrowed.gguf"); + write_raw_gguf_dims(file_path, {(1ull << 32) + 32}, q8_0, 34); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension that narrows to a negative shape") { + std::string file_path = get_temp_file("test_gguf_dim_negative.gguf"); + write_raw_gguf_dims(file_path, {1ull << 31}, q8_0, 34); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension product that wraps to a small byte size") { + // 96 * 384307168202282326 wraps to 64, so the tensor needs only 68 bytes of + // data and passes the byte size check, while the narrowed shape + // {1431655766, 96} still describes 2^37 elements. That drives the scales + // element count past INT32_MAX, which is what truncates the allocation size + // while the extractor still loops over the full count. + std::string file_path = get_temp_file("test_gguf_dim_wrap.gguf"); + write_raw_gguf_dims(file_path, {96, 384307168202282326ull}, q8_0, 68); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension product that wraps to zero") { + // Wraps num_weights to exactly 0, so the tensor claims no data at all. + std::string file_path = get_temp_file("test_gguf_dim_wrap_zero.gguf"); + write_raw_gguf_dims( + file_path, + {(1ull << 32) + (1ull << 25), + (1ull << 32) + (1ull << 15), + (1ull << 32) + (1ull << 24)}, + q8_0, + 0); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor without dimensions") { + // gguf_load_quantized indexes the last shape element unconditionally. + std::string file_path = get_temp_file("test_gguf_no_dims.gguf"); + write_raw_gguf_dims(file_path, {}, q8_0, 0); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } +} + // Writes a metadata-only GGUF (no tensors) whose metadata KV section is // `kv_section` verbatim, so a caller can encode values whose lengths exceed the // file to exercise check_metadata_value_in_file(). `kv_count` must match the diff --git a/Source/Cmlx/mlx/tests/ops_tests.cpp b/Source/Cmlx/mlx/tests/ops_tests.cpp index 09236f1da..9c95e9b5e 100644 --- a/Source/Cmlx/mlx/tests/ops_tests.cpp +++ b/Source/Cmlx/mlx/tests/ops_tests.cpp @@ -706,7 +706,7 @@ TEST_CASE("test transpose") { CHECK_EQ(y.shape(), Shape{1}); CHECK_EQ(y.item(), 1); - CHECK_THROWS_AS(transpose(x, {1}), std::invalid_argument); + CHECK_THROWS_AS(transpose(x, {1}), std::out_of_range); CHECK_THROWS_AS(transpose(x, {0, 0}), std::invalid_argument); // Works with empty array @@ -2003,6 +2003,16 @@ TEST_CASE("test arithmetic binary ops") { y = array(false); CHECK(std::isnan(divide(x, y).item())); + // 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()); + CHECK(array_equal(remainder(num, den), num).item()); + } + } + // Check maximum and minimum x = array(1.0f); y = array(0.0f); @@ -2890,6 +2900,44 @@ TEST_CASE("test as_strided op") { CHECK(array_equal(y, expected).item()); } +TEST_CASE("test sort and scan on empty arrays") { + // These must be evaluated, not only shape checked, to reach the kernel. + auto check_empty = [](const array& r, Shape shape, Dtype dt) { + auto out = r; + eval(out); + CHECK_EQ(out.shape(), shape); + CHECK_EQ(out.dtype(), dt); + CHECK_EQ(out.size(), 0); + }; + + for (auto dt : {float32, int32, uint32, bool_}) { + auto x = zeros({0}, dt); + check_empty(sort(x), Shape{0}, dt); + check_empty(argsort(x), Shape{0}, uint32); + check_empty(cumsum(x), Shape{0}, dt == bool_ ? int32 : dt); + check_empty(cumprod(x), Shape{0}, dt); + check_empty(cummax(x), Shape{0}, dt); + check_empty(cummin(x), Shape{0}, dt); + } + + // Empty along the sorted axis, non-empty elsewhere. + auto x = zeros({3, 0}, float32); + check_empty(sort(x, 1), Shape{3, 0}, float32); + check_empty(argsort(x, 1), Shape{3, 0}, uint32); + check_empty(cumsum(x, 1), Shape{3, 0}, float32); + + // Empty on another axis, so the sorted axis itself is non-empty. + x = zeros({0, 3}, float32); + check_empty(sort(x, 1), Shape{0, 3}, float32); + check_empty(argsort(x, 1), Shape{0, 3}, uint32); + check_empty(cumsum(x, 1), Shape{0, 3}, float32); + + // Non-contiguous empty input reaches the strided path. + x = transpose(zeros({4, 0}, float32)); + check_empty(sort(x, 1), Shape{0, 4}, float32); + check_empty(cumsum(x, 1), Shape{0, 4}, float32); +} + TEST_CASE("test scan op") { auto x = array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}, {2, 3}); auto y = cumsum(x, 1, false, true); @@ -2949,6 +2997,40 @@ TEST_CASE("test scan op") { y = vmap(fun, 1, 1)(x); expected = array({1.0f, 2.0f, 4.0f, 6.0f, 9.0f, 12.0f, 16.0f, 20.0f}, {4, 2}); CHECK(array_equal(y, expected).item()); + + // Scanning an empty axis is a no-op + x = zeros({2, 0}); + y = cumsum(x, 1); + eval(y); + CHECK_EQ(y.shape(), Shape{2, 0}); + y = cummax(x, 1); + eval(y); + CHECK_EQ(y.shape(), Shape{2, 0}); + x = zeros({0, 2}); + y = cumsum(x, 0); + eval(y); + CHECK_EQ(y.shape(), Shape{0, 2}); + y = cummin(x, 0); + eval(y); + CHECK_EQ(y.shape(), Shape{0, 2}); +} + +TEST_CASE("test sort op") { + // Sorting an empty axis is a no-op + auto x = zeros({2, 0}); + auto y = sort(x, 1); + eval(y); + CHECK_EQ(y.shape(), Shape{2, 0}); + y = argsort(x, 1); + eval(y); + CHECK_EQ(y.shape(), Shape{2, 0}); + x = zeros({0, 2}); + y = sort(x, 0); + eval(y); + CHECK_EQ(y.shape(), Shape{0, 2}); + y = argsort(x, 0); + eval(y); + CHECK_EQ(y.shape(), Shape{0, 2}); } TEST_CASE("test pad") { @@ -4491,6 +4573,66 @@ TEST_CASE("test conv shape overflow") { Shape{1, 8, 8, 1}); } +TEST_CASE("test pad with an axes subset") { + // edge_pad ignored `axes` and indexed the pad sizes by array axis, so it + // read past the end of those vectors and placed the input on the wrong + // axes. reflect and symmetric indexed by position but did not normalize a + // negative axis. + auto x = reshape(arange(24.0f), {2, 3, 4}); + auto all = {"constant", "edge", "reflect", "symmetric"}; + + // Padding a subset of the axes matches the equivalent all-axes spelling. + for (auto mode : all) { + CHECK(array_equal( + pad(x, {1}, Shape{1}, Shape{2}, array(0.0f), mode), + pad(x, + {0, 1, 2}, + Shape{0, 1, 0}, + Shape{0, 2, 0}, + array(0.0f), + mode)) + .item()); + } + + // A negative axis matches its non-negative spelling. + for (auto mode : all) { + CHECK(array_equal( + pad(x, {-2}, Shape{1}, Shape{2}, array(0.0f), mode), + pad(x, {1}, Shape{1}, Shape{2}, array(0.0f), mode)) + .item()); + } + + // The order of `axes` does not change the result. + auto y = reshape(arange(25.0f), {5, 5}); + for (auto mode : all) { + CHECK(array_equal( + pad(y, {1, 0}, Shape{1, 2}, Shape{1, 2}, array(0.0f), mode), + pad(y, {0, 1}, Shape{2, 1}, Shape{2, 1}, array(0.0f), mode)) + .item()); + } + + // An ndim past SmallVector's inline capacity puts the pad sizes on the heap, + // where the old read ran off the end of the allocation. + std::vector wide_axes(11); + std::iota(wide_axes.begin(), wide_axes.end(), 0); + auto wide = + pad(zeros(Shape(34, 1)), + wide_axes, + Shape(11, 1), + Shape(11, 0), + array(0.0f), + "edge"); + CHECK_EQ(wide.size(), 2048); + + // An axis outside the array is rejected rather than indexed. + for (auto mode : all) { + CHECK_THROWS_AS( + pad(x, {5}, Shape{1}, Shape{1}, array(0.0f), mode), std::out_of_range); + CHECK_THROWS_AS( + pad(x, {-5}, Shape{1}, Shape{1}, array(0.0f), mode), std::out_of_range); + } +} + TEST_CASE("test pad shape overflow") { // A padding sum that overflows int32 is rejected, not wrapped. // https://github.com/ml-explore/mlx/issues/3611 diff --git a/Source/Cmlx/mlx/tests/scheduler_tests.cpp b/Source/Cmlx/mlx/tests/scheduler_tests.cpp index 3a8f7e86b..cc2228e0f 100644 --- a/Source/Cmlx/mlx/tests/scheduler_tests.cpp +++ b/Source/Cmlx/mlx/tests/scheduler_tests.cpp @@ -114,6 +114,35 @@ TEST_CASE("test thread unsafe stream") { CHECK_EQ(expected, actual); } +TEST_CASE("test eval does not create default stream") { + auto s = new_thread_unsafe_stream(default_device()); + size_t num_streams = get_streams().size(); + + std::thread t([&] { + async_eval(arange(10, s)); + eval(arange(10, s)); + }); + t.join(); + + CHECK_EQ(get_streams().size(), num_streams); +} + +TEST_CASE("test compile does not create default stream") { + auto s = new_thread_unsafe_stream(default_device()); + size_t num_streams = get_streams().size(); + + std::function(const std::vector&)> fun = + [s](const std::vector& inputs) { + return std::vector{abs(inputs[0], s)}; + }; + auto cfun = compile(fun); + + std::thread t([&] { eval(cfun({array({-1, 2})})); }); + t.join(); + + CHECK_EQ(get_streams().size(), num_streams); +} + TEST_CASE("test thread local stream") { auto s = new_thread_local_stream(default_device()); int result = sum(arange(10, s)).item(); @@ -249,3 +278,29 @@ TEST_CASE("test scheduler races") { } eval(a, y); } + +// The fence orders work between two streams. Check that the consumer sees the +// producer result for every combination of producer and consumer device. +TEST_CASE("test cross stream fence ordering") { + if (!gpu::is_available()) { + return; + } + std::vector devices = {Device::cpu, Device::gpu}; + + for (auto pd : devices) { + for (auto cd : devices) { + auto ps = new_stream(pd); + auto cs = new_stream(cd); + + array one = full({1}, 1.0f, float32, ps); + array x = full({64, 64}, 1.0f, float32, ps); + for (int i = 0; i < 20; ++i) { + x = add(x, one, ps); + } + + // The consumer is on the other stream, so eval builds a fence. + array y = sum(x, cs); + CHECK_EQ(y.item(), doctest::Approx(64 * 64 * 21.0f)); + } + } +} diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift index 4fd0b1f9e..427e976e5 100644 --- a/Tests/MLXTests/FP8Tests.swift +++ b/Tests/MLXTests/FP8Tests.swift @@ -10,6 +10,8 @@ class FP8Tests: XCTestCase { } 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) } let bytes = MLXArray(allBytes) let decoded = fromFP8(bytes, dtype: .float32) @@ -17,8 +19,14 @@ class FP8Tests: XCTestCase { XCTAssertEqual(decoded.dtype, .float32) XCTAssertEqual(decoded.shape, [256]) - let expected = MLXArray(allBytes.map { e4m3Value($0) }) - assertEqual(decoded, expected) + 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)") + } + } } func testToFP8EncodesE4M3Bytes() { diff --git a/UPSTREAM_MERGE_PLAN.md b/UPSTREAM_MERGE_PLAN.md index 2d8ce6484..f4466f017 100644 --- a/UPSTREAM_MERGE_PLAN.md +++ b/UPSTREAM_MERGE_PLAN.md @@ -14,8 +14,10 @@ this repository and a SwiftPM consumer needs nothing but this package. | Directory | Contents | |---|---| -| `Source/Cmlx/mlx` | upstream `ml-explore/mlx` **v0.32.2** (`1f8e74e`) + the patch series below | +| `Source/Cmlx/mlx` | upstream `ml-explore/mlx` **v0.32.3** (`64ea011`) + the patch series below | | `Source/Cmlx/mlx-c` | upstream `ml-explore/mlx-c` **`ebc88f1`** + the patch series below | + +> mlx-c has not caught up with core 0.32.3 upstream (its `main` is `a341b49`, a version bump on top of `ebc88f1`), so this fork carries one adaptation: `mlx_gather_qmm` passes `std::nullopt` for the core's new `global_scale` argument (mlx #4458). The C API does not expose it yet. | `Source/Cmlx/mlx-generated` | produced by `tools/update-mlx.sh` (never edit by hand) | `.gitmodules` is intentionally absent. When a merge reports a conflict on @@ -31,7 +33,7 @@ submodule pointer: delete it (`git rm`) and re-vendor the tree instead. and the matching hooks in `mlx/fast.{h,cpp}`. - **TurboKV**: `mlx/fast/turbo_quant.h` and the TurboQuant decompression path in `mlx/backend/metal/kernels/sdpa_vector.h`. -- **I/O fast loaders**: `mlx/io/load.{h,cpp}` (thread-pooled, pread-based loaders). +- **I/O loaders**: `mlx/io/load.{h,cpp}`. Since 0.32.3 upstream's `ParallelFileReader` has its own shared batch pool; the fork keeps that and adds only the iOS guards (`MLX_IO_THREAD_COUNT`, `MLX_IOS_SEQUENTIAL_IO`: sequential `pread()`, one reader thread). - **Metal device**: `mlx/backend/metal/device.cpp` carries only quieter Metal compile logging. The old "lazily create a command encoder when a stream is first used on a new thread" patch is **not** carried: upstream's Swift `StreamPool` creates every stream with `mlx_stream_new_thread_unsafe` From a08b272d11b3aff52531825e1db88c3062d5660e Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Mon, 5 Oct 2026 23:12:21 +0200 Subject: [PATCH 2/6] Backport mlx #4594: fix cooperative tensor build on macOS 27 mlx v0.32.3's gated delta NAX kernel does not compile with the macOS 27 SDK (get_destination_cooperative_tensor rejects decltype of the operand variables), which breaks the mlx.metallib build with -DMLX_ENABLE_NAX=1. Apply the same change as upstream mlx #4594, which landed after v0.32.3, to gated_delta_update_nax.h and regenerate mlx-generated. Drop this commit when the vendored core moves past #4594. Co-Authored-By: Claude Sonnet 5.5 --- .../mlx-generated/gated_delta_update_nax.cpp | 48 ++++++++++++++----- .../metal/gated_delta_update_nax.h | 48 ++++++++++++++----- .../metal/kernels/gated_delta_update_nax.h | 48 ++++++++++++++----- 3 files changed, 108 insertions(+), 36 deletions(-) diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp index 56afceba5..2dff85266 100644 --- a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp @@ -1240,15 +1240,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -1286,15 +1294,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -1337,6 +1353,15 @@ METAL_FUNC static constexpr void mman( 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 = @@ -1344,10 +1369,9 @@ METAL_FUNC static constexpr void mman( .template get_right_input_cooperative_tensor(); // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); // Load A in to left operand registers STEEL_PRAGMA_UNROLL 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 15b98fd21..1a576f4aa 100644 --- a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h @@ -121,15 +121,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -167,15 +175,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -218,6 +234,15 @@ METAL_FUNC static constexpr void mman( 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 = @@ -225,10 +250,9 @@ METAL_FUNC static constexpr void mman( .template get_right_input_cooperative_tensor(); // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); // Load A in to left operand registers STEEL_PRAGMA_UNROLL 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 f4ae51826..4e65fe1e5 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 @@ -121,15 +121,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -167,15 +175,23 @@ METAL_FUNC static constexpr void mma( 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< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { @@ -218,6 +234,15 @@ METAL_FUNC static constexpr void mman( 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 = @@ -225,10 +250,9 @@ METAL_FUNC static constexpr void mman( .template get_right_input_cooperative_tensor(); // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); // Load A in to left operand registers STEEL_PRAGMA_UNROLL From d22833dc84be787f03350555ea0e55509d73c34e Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Tue, 6 Oct 2026 13:28:01 +0200 Subject: [PATCH 3/6] Skip the signed floor_divide correction for unsigned integers mlx v0.32.3 (#4515) made integer floor_divide round toward negative infinity by adding remainder, less, not_equal, logical_and and subtract after the divide. For unsigned dtypes that correction is always zero. mlx-swift-lm's MoE expert sort calls order.floorDivide(topK) on uint32 indices in every layer; on this core a streamed-experts decode step died with a Metal GPU timeout (kIOGPUCommandBufferCallbackErrorTimeout) on every run, and decode was slower even when it survived. Found by bisecting 0.32.2..0.32.3: #4515 is the first bad commit, and the 0.32.2 core never fails the same request. Return the plain quotient for unsigned dtypes. Signed integers keep the upstream fix. FloorDivideTests covers both. Drop this when upstream makes the same change. Co-Authored-By: Claude Sonnet 5.5 --- Source/Cmlx/mlx/mlx/ops.cpp | 7 +++++++ Tests/MLXTests/FloorDivideTests.swift | 24 ++++++++++++++++++++++++ UPSTREAM_MERGE_PLAN.md | 7 +++++++ 3 files changed, 38 insertions(+) create mode 100644 Tests/MLXTests/FloorDivideTests.swift diff --git a/Source/Cmlx/mlx/mlx/ops.cpp b/Source/Cmlx/mlx/mlx/ops.cpp index 83094d25b..095ba9dcf 100644 --- a/Source/Cmlx/mlx/mlx/ops.cpp +++ b/Source/Cmlx/mlx/mlx/ops.cpp @@ -3167,6 +3167,13 @@ array floor_divide( 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. + if (issubdtype(dtype, unsignedinteger)) { + return quotient; + } auto rem = remainder(a, b, s); auto zero = array(0, dtype); auto step = logical_and( diff --git a/Tests/MLXTests/FloorDivideTests.swift b/Tests/MLXTests/FloorDivideTests.swift new file mode 100644 index 000000000..26a52b905 --- /dev/null +++ b/Tests/MLXTests/FloorDivideTests.swift @@ -0,0 +1,24 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +class FloorDivideTests: XCTestCase { + + /// mlx #4515: integer floor_divide rounds toward negative infinity like numpy. + func testFloorDivideSignedIntegersRoundsTowardNegativeInfinity() { + let a = MLXArray([4, 5, -1, -6] as [Int32]) + let b = MLXArray([-2, 3, 2, -3] as [Int32]) + XCTAssertEqual(a.floorDivide(b).asArray(Int32.self), [-2, 1, -1, 2]) + } + + /// The MoE expert sort divides uint32 sort indices by top-k. Unsigned values need no sign + /// correction, so the result must match plain integer division. + func testFloorDivideUnsignedMatchesIntegerDivision() { + let order = MLXArray((0 ..< 64).map { UInt32($0) }) + let q = order.floorDivide(8) + XCTAssertEqual(q.dtype, .uint32) + XCTAssertEqual(q.asArray(UInt32.self), (0 ..< 64).map { UInt32($0 / 8) }) + } +} diff --git a/UPSTREAM_MERGE_PLAN.md b/UPSTREAM_MERGE_PLAN.md index f4466f017..d4657275d 100644 --- a/UPSTREAM_MERGE_PLAN.md +++ b/UPSTREAM_MERGE_PLAN.md @@ -34,6 +34,13 @@ submodule pointer: delete it (`git rm`) and re-vendor the tree instead. - **TurboKV**: `mlx/fast/turbo_quant.h` and the TurboQuant decompression path in `mlx/backend/metal/kernels/sdpa_vector.h`. - **I/O loaders**: `mlx/io/load.{h,cpp}`. Since 0.32.3 upstream's `ParallelFileReader` has its own shared batch pool; the fork keeps that and adds only the iOS guards (`MLX_IO_THREAD_COUNT`, `MLX_IOS_SEQUENTIAL_IO`: sequential `pread()`, one reader thread). +- **`floor_divide` on unsigned integers** (`mlx/ops.cpp`): upstream #4515 (in 0.32.3) made integer + `floor_divide` correct for negative operands by adding `remainder`, `less`, `not_equal`, + `logical_and` and `subtract` after the divide. For unsigned dtypes that correction is always zero, so + the fork returns the plain quotient. mlx-swift-lm's MoE expert sort (`SwitchLayers.gatherSort`) calls + `order.floorDivide(topK)` on uint32 indices in every layer; with #4515 a streamed-experts decode step + hit a Metal GPU timeout (found by bisecting 0.32.2..0.32.3). Test: + `FloorDivideTests`. Drop this when upstream makes the same change. - **Metal device**: `mlx/backend/metal/device.cpp` carries only quieter Metal compile logging. The old "lazily create a command encoder when a stream is first used on a new thread" patch is **not** carried: upstream's Swift `StreamPool` creates every stream with `mlx_stream_new_thread_unsafe` From fa87f63d86810dfb1c5c231b6a3bba5f9a294f41 Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 6 Oct 2026 10:42:37 -0700 Subject: [PATCH 4/6] fix: bundle the gather_mm_offsets kernel into the SwiftPM metallib mlx 0.32.3 sorted gatherMM / gatherQuantizedMM (sortedIndices: true) call the new ahead-of-time kernel gather_mm_offsets. SwiftPM builds default.metallib only from the files listed in KERNEL_LIST in tools/fix-metal-includes.sh, which did not include it, so the sorted path aborts with "Unable to load kernel gather_mm_offsets" in any consumer that has no mlx.metallib next to its binary. Add the kernel to KERNEL_LIST and regenerate mlx-generated/metal. SortedGatherTests covers the sorted and unsorted gatherMM (float16, bfloat16, float32) and gatherQuantizedMM paths. Inputs are small integers, so every sum is exact and the Metal results are compared with arrayEqual against a CPU reference; there is no tolerance to tune. Without the kernel the tests abort with the error above. Co-Authored-By: Claude Sonnet 5.5 --- .../metal/gather_mm_offsets.metal | 23 +++++ Tests/MLXTests/SortedGatherTests.swift | 98 +++++++++++++++++++ tools/fix-metal-includes.sh | 1 + 3 files changed, 122 insertions(+) create mode 100644 Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal create mode 100644 Tests/MLXTests/SortedGatherTests.swift diff --git a/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal b/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal new file mode 100644 index 000000000..40b933d51 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal @@ -0,0 +1,23 @@ +// Copyright © 2026 Apple Inc. + +#include + +using namespace metal; + +[[kernel]] void gather_mm_offsets( + const device uint32_t* indices [[buffer(0)]], + device int32_t* offsets [[buffer(1)]], + const constant int& M [[buffer(2)]], + uint g [[thread_position_in_grid]]) { + int lo = 0; + int hi = M; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (indices[mid] < g) { + lo = mid + 1; + } else { + hi = mid; + } + } + offsets[g] = lo; +} diff --git a/Tests/MLXTests/SortedGatherTests.swift b/Tests/MLXTests/SortedGatherTests.swift new file mode 100644 index 000000000..f3a4bb0a4 --- /dev/null +++ b/Tests/MLXTests/SortedGatherTests.swift @@ -0,0 +1,98 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +/// `gatherMM` / `gatherQuantizedMM` with `sortedIndices: true` take the "rhs sorted" Metal paths, which +/// call the ahead-of-time kernel `gather_mm_offsets` (mlx >= 0.32.3, #4567 / #4572). SwiftPM builds +/// `default.metallib` only from the kernels listed in `tools/fix-metal-includes.sh`, so a kernel that is +/// missing from that list only fails at run time, and only when this path is hit. +/// +/// All inputs are small integers, so every sum is exact in float16, bfloat16, float32 and TF32, and the +/// Metal results can be compared with `arrayEqual` against a CPU reference. There is no tolerance to tune. +class SortedGatherTests: XCTestCase { + + /// Deterministic integers in `[-2, 2]`. + private func smallIntegers(_ count: Int, seed: Int) -> [Float] { + (0 ..< count).map { Float((($0 &* 7 &+ seed &* 13) % 5) - 2) } + } + + /// Ascending expert ids. The group sizes are uneven and expert 1 is empty, which exercises the + /// binary search in `gather_mm_offsets` for a group with no rows. + private func sortedExpertIds(_ counts: [Int]) -> [UInt32] { + counts.enumerated().flatMap { Array(repeating: UInt32($0.offset), count: $0.element) } + } + + func testSortedGatherMMMatchesCPUReference() { + let counts = [30, 0, 20, 14] + let (batch, experts, k, n) = (counts.reduce(0, +), counts.count, 32, 32) + let ids = MLXArray(sortedExpertIds(counts)) + + let a32 = MLXArray(smallIntegers(batch * k, seed: 1), [batch, 1, k]) + let b32 = MLXArray(smallIntegers(experts * k * n, seed: 2), [experts, k, n]) + + let expected = matmul(a32, b32.take(ids, axis: 0, stream: .cpu), stream: .cpu) + + for dtype in [DType.float16, .bfloat16, .float32] { + let a = a32.asType(dtype) + let b = b32.asType(dtype) + + let sorted = gatherMM(a, b, rhsIndices: ids, sortedIndices: true, stream: .gpu) + let unsorted = gatherMM(a, b, rhsIndices: ids, sortedIndices: false, stream: .gpu) + + XCTAssertEqual(sorted.shape, [batch, 1, n], "\(dtype)") + XCTAssertTrue( + arrayEqual(sorted.asType(.float32), expected, stream: .cpu).item(Bool.self), + "sorted gatherMM differs from the CPU reference (\(dtype))") + XCTAssertTrue( + arrayEqual(unsorted.asType(.float32), expected, stream: .cpu).item(Bool.self), + "unsorted gatherMM differs from the CPU reference (\(dtype))") + } + } + + func testSortedGatherQuantizedMMMatchesCPUReference() { + let counts = [30, 0, 20, 14] + let (batch, experts, k, n) = (counts.reduce(0, +), counts.count, 64, 64) + let ids = MLXArray(sortedExpertIds(counts)) + + // Each quantization group (one row of `k == groupSize` values) holds the integers 0...15 and + // contains both 0 and 15, so the affine scale is exactly 1 and the bias exactly 0. + let weightValues = (0 ..< experts * n).flatMap { row in + (0 ..< k).map { col -> Float in + switch col { + case 0: return 0 + case 1: return 15 + default: return Float((col &* 7 &+ row &* 3) % 16) + } + } + } + let w = MLXArray(weightValues, [experts, n, k]).asType(.float16) + let (wq, scales, biases) = quantized(w, groupSize: 64, bits: 4) + + // x values in [-1, 1]: |sum| <= 64 * 15, exact in float16. + let x32 = MLXArray(smallIntegers(batch * k, seed: 3).map { max(-1, min(1, $0)) }, [batch, 1, k]) + let x = x32.asType(.float16) + + let dequantizedWeights = dequantized( + wq, scales: scales, biases: biases, groupSize: 64, bits: 4, dtype: .float32) + let expected = matmul( + x32, dequantizedWeights.take(ids, axis: 0, stream: .cpu).transposed(0, 2, 1, stream: .cpu), + stream: .cpu) + + let sorted = gatherQuantizedMM( + x, wq, scales: scales, biases: biases, rhsIndices: ids, transpose: true, + groupSize: 64, bits: 4, sortedIndices: true, stream: .gpu) + let unsorted = gatherQuantizedMM( + x, wq, scales: scales, biases: biases, rhsIndices: ids, transpose: true, + groupSize: 64, bits: 4, sortedIndices: false, stream: .gpu) + + XCTAssertEqual(sorted.shape, [batch, 1, n]) + XCTAssertTrue( + arrayEqual(sorted.asType(.float32), expected, stream: .cpu).item(Bool.self), + "sorted gatherQuantizedMM differs from the CPU reference") + XCTAssertTrue( + arrayEqual(unsorted.asType(.float32), expected, stream: .cpu).item(Bool.self), + "unsorted gatherQuantizedMM differs from the CPU reference") + } +} diff --git a/tools/fix-metal-includes.sh b/tools/fix-metal-includes.sh index 6606683cf..3493d8470 100755 --- a/tools/fix-metal-includes.sh +++ b/tools/fix-metal-includes.sh @@ -23,6 +23,7 @@ KERNEL_LIST=" \ arg_reduce.metal \ conv.metal \ dot.metal \ +gather_mm_offsets.metal \ layer_norm.metal \ random.metal \ rms_norm.metal \ From 2c7c12a695ad391a97b580f9733dca9f278dc2fa Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 6 Oct 2026 10:42:37 -0700 Subject: [PATCH 5/6] build(xcode): exclude the files new in mlx 0.32.3 from the Cmlx target The Xcode project lists the files the Cmlx target must not compile (kernel sources, CUDA, tests, python and benchmarks). The 0.32.3 sync added files that are not on that list, so `xcodebuild -scheme MLX` failed: - mlx/backend/metal/kernels/{cross_entropy,gated_delta_update, gated_delta_update_nax}.{h,metal} and gather_mm_offsets.metal (the generated copies in mlx-generated/metal are the ones built) - mlx/backend/cuda/{cholesky.cu,cusolver_utils.*,wddm.*} and tests/cuda_tests.cpp (need CUDA headers / doctest) - the new python tests and benchmarks Also exclude the fork's mlx/backend/metal/kernels/moe_stream.metal. Its generated copy is already built, so linking default.metallib failed with a duplicated symbol 'streamed_moe_gemm'. Co-Authored-By: Claude Sonnet 5.5 --- xcode/MLX.xcodeproj/project.pbxproj | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 0bb55a5ab..3a62a5c99 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -502,13 +502,16 @@ benchmarks/python/distributed_bench.py, benchmarks/python/einsum_bench.py, benchmarks/python/fft_bench.py, + benchmarks/python/gated_delta_bench.py, benchmarks/python/gather_bench.py, benchmarks/python/gather_mm_bench.py, benchmarks/python/gather_qmm_bench.py, + benchmarks/python/gather_qqmm_bench.py, benchmarks/python/hadamard_bench.py, benchmarks/python/large_gemm_bench.py, benchmarks/python/layer_norm_bench.py, benchmarks/python/masked_scatter.py, + benchmarks/python/qqmm_bench.py, benchmarks/python/rms_norm_bench.py, benchmarks/python/rope_bench.py, benchmarks/python/scatter_bench.py, @@ -517,6 +520,8 @@ benchmarks/python/segmented_mm_bench.py, benchmarks/python/single_ops.py, benchmarks/python/slice_update_bench.py, + benchmarks/python/swiglu_bench.py, + benchmarks/python/swiglu_qmm_bench.py, benchmarks/python/synchronize_bench.py, benchmarks/python/time_utils.py, CITATION.cff, @@ -725,6 +730,7 @@ mlx/backend/cuda/binary/power.cu, mlx/backend/cuda/binary/remainder.cu, mlx/backend/cuda/binary/subtract.cu, + mlx/backend/cuda/cholesky.cu, mlx/backend/cuda/CMakeLists.txt, mlx/backend/cuda/compiled.cpp, mlx/backend/cuda/conv.cpp, @@ -744,6 +750,8 @@ mlx/backend/cuda/cuda.h, mlx/backend/cuda/cudnn_utils.cpp, mlx/backend/cuda/cudnn_utils.h, + mlx/backend/cuda/cusolver_utils.cpp, + mlx/backend/cuda/cusolver_utils.h, mlx/backend/cuda/custom_kernel.cpp, mlx/backend/cuda/cutlass_utils.cuh, mlx/backend/cuda/delayload.cpp, @@ -889,6 +897,8 @@ mlx/backend/cuda/utils.cpp, mlx/backend/cuda/utils.h, mlx/backend/cuda/vector_types.cuh, + mlx/backend/cuda/wddm.cpp, + mlx/backend/cuda/wddm.h, mlx/backend/cuda/worker.cpp, mlx/backend/cuda/worker.h, mlx/backend/gpu/CMakeLists.txt, @@ -922,6 +932,8 @@ mlx/backend/metal/kernels/conv.metal, mlx/backend/metal/kernels/copy.h, mlx/backend/metal/kernels/copy.metal, + mlx/backend/metal/kernels/cross_entropy.h, + mlx/backend/metal/kernels/cross_entropy.metal, mlx/backend/metal/kernels/defines.h, mlx/backend/metal/kernels/dot.h, mlx/backend/metal/kernels/dot.metal, @@ -938,6 +950,11 @@ 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_update_nax.h, + mlx/backend/metal/kernels/gated_delta_update_nax.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, mlx/backend/metal/kernels/gemv_masked.h, mlx/backend/metal/kernels/gemv_masked.metal, mlx/backend/metal/kernels/gemv.h, @@ -954,6 +971,7 @@ mlx/backend/metal/kernels/logging.h, mlx/backend/metal/kernels/logsumexp.h, mlx/backend/metal/kernels/logsumexp.metal, + mlx/backend/metal/kernels/moe_stream.metal, mlx/backend/metal/kernels/quantized_nax.h, mlx/backend/metal/kernels/quantized_nax.metal, mlx/backend/metal/kernels/quantized_utils.h, @@ -1234,6 +1252,7 @@ python/tests/test_einsum.py, python/tests/test_eval.py, python/tests/test_export_import.py, + python/tests/test_fast_gated_delta.py, python/tests/test_fast_sdpa.py, python/tests/test_fast.py, python/tests/test_fft.py, @@ -1263,6 +1282,7 @@ tests/CMakeLists.txt, tests/compile_tests.cpp, tests/creations_tests.cpp, + tests/cuda_tests.cpp, tests/custom_vjp_tests.cpp, tests/device_tests.cpp, tests/einsum_tests.cpp, From 033d4285bfcf1cb6583f0a90bf90ea3fd1d927be Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 6 Oct 2026 10:42:37 -0700 Subject: [PATCH 6/6] ci: run MLXTests against the SwiftPM metallib and build the xcodeproj pull_request.yml only runs its build/test jobs on ml-explore/mlx-swift (self-hosted runner), and the integration workflows copy a full mlx.metallib from the Python wheel. Neither exercises the Metal kernels that SwiftPM itself bundles into default.metallib, so a kernel missing from KERNEL_LIST in tools/fix-metal-includes.sh only fails at run time. Likewise nothing builds xcode/MLX.xcodeproj, whose exclusion list needs an update for every new file in Source/Cmlx/mlx. - mlx-tests: build MLXTests with Xcode and run them with no mlx.metallib on disk (a guard step fails if one shows up), so the tests load the metallib SwiftPM built from Source/Cmlx. - xcodeproj-build: xcodebuild build-for-testing -scheme MLX. Co-Authored-By: Claude Sonnet 5.5 --- .github/workflows/mlx_tests.yml | 77 +++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 .github/workflows/mlx_tests.yml diff --git a/.github/workflows/mlx_tests.yml b/.github/workflows/mlx_tests.yml new file mode 100644 index 000000000..966322142 --- /dev/null +++ b/.github/workflows/mlx_tests.yml @@ -0,0 +1,77 @@ +name: MLX Tests (SwiftPM metallib) + +# The "Build and Test" job in pull_request.yml only runs on ml-explore/mlx-swift (self-hosted runner), and the +# integration workflows copy a full mlx.metallib from the Python wheel, so none of them runs the Metal kernels +# that SwiftPM itself bundles into default.metallib. A kernel missing from KERNEL_LIST in +# tools/fix-metal-includes.sh therefore only shows up at run time in a consumer that has no mlx.metallib next +# to its binary. This job builds MLXTests with Xcode and runs them with no mlx.metallib on disk, so the tests +# load the metallib that SwiftPM built from Source/Cmlx. The second job builds xcode/MLX.xcodeproj, whose +# Cmlx target keeps its own list of excluded kernel sources that must be extended for every new .metal file. + +on: + pull_request: + branches: [ main ] + push: + branches: [ main ] + +permissions: + contents: read + +jobs: + mlx-tests: + name: MLXTests (SwiftPM bundle metallib) + runs-on: macos-15 + timeout-minutes: 60 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + + - name: Install Metal Toolchain + run: xcodebuild -downloadComponent MetalToolchain || true + + - name: Build for testing (macOS) + run: | + xcodebuild build-for-testing \ + -scheme mlx-swift-Package \ + -destination 'platform=macOS' \ + -derivedDataPath "$RUNNER_TEMP/dd" \ + -skipPackagePluginValidation -skipMacroValidation + + - name: Run MLXTests + run: | + # Fail early if a stray mlx.metallib could hide a kernel missing from default.metallib. + if find "$RUNNER_TEMP/dd/Build/Products" -name mlx.metallib | grep -q .; then + echo "mlx.metallib found next to the test products; this job must use the SwiftPM bundle metallib" + exit 1 + fi + xcrun xctest "$RUNNER_TEMP/dd/Build/Products/Debug/MLXTests.xctest" + + xcodeproj-build: + name: Build xcode/MLX.xcodeproj (macOS) + runs-on: macos-15 + timeout-minutes: 60 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + + - name: Install Metal Toolchain + run: xcodebuild -downloadComponent MetalToolchain || true + + - name: Build xcodeproj + working-directory: xcode + run: | + xcodebuild build-for-testing \ + -scheme MLX \ + -destination 'platform=macOS' \ + -derivedDataPath "$RUNNER_TEMP/dd-xcodeproj" \ + -skipPackagePluginValidation -skipMacroValidation