Bump MLX pin to v32.2 (pytorch#22495)

As titled

Author: metascroy

.github/workflows/mlx.yml

@@ -71,6 +71,37 @@ jobs:
         ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 ))
         echo "::endgroup::"
 
+        echo "::group::Check MLX artifact sizes"
+        size_check_failed=0
+        check_artifact_size() {
+          local path="$1"
+          local max_size="$2"
+          local baseline_size="$3"
+
+          if [[ ! -f "${path}" ]]; then
+            echo "::error file=${path}::Expected MLX artifact not found"
+            size_check_failed=1
+            return
+          fi
+
+          local actual_size
+          actual_size=$(wc -c < "${path}" | tr -d '[:space:]')
+          echo "${path}: ${actual_size} bytes (baseline ${baseline_size}, limit ${max_size})"
+          if (( actual_size > max_size )); then
+            echo "::error file=${path}::MLX artifact is ${actual_size} bytes, exceeding the ${max_size}-byte limit (baseline ${baseline_size})"
+            size_check_failed=1
+          fi
+        }
+
+        # Baselines recorded from the MLX v0.32.2 Release build on 2026-09-02.
+        check_artifact_size cmake-out/backends/mlx/libmlxdelegate.a 5000000 4159552
+        check_artifact_size cmake-out/backends/mlx/mlx/libmlx.a 14000000 11458544
+        check_artifact_size cmake-out/backends/mlx/mlx/mlx/backend/metal/kernels/mlx.metallib 2000000 1460840
+        echo "::endgroup::"
+        if (( size_check_failed )); then
+          exit 1
+        fi
+
         echo "::group::Run mutable-state (multi-session) unit test"
         ./cmake-out/backends/mlx/test/mlx_mutable_state_test
         echo "::endgroup::"

backends/mlx/CMakeLists.txt

@@ -167,8 +167,9 @@ endif()
 # with add_subdirectory would drop its whole project into ExecuTorch's
 # target/option namespace, which collides with shared deps MLX fetches (e.g.
 # nlohmann_json) and leaks MLX's MLX_BUILD_* options into our cache. The
-# isolated scope runs MLX's FetchContent in its own namespace, so no collision
-# and no submodule patching are needed.
+# isolated scope runs MLX's FetchContent in its own namespace, so dependency
+# collisions do not require patching. The platform packaging patches below are
+# still applied before the external configure step.
 include(ExternalProject)
 
 set(_mlx_binary_dir ${CMAKE_CURRENT_BINARY_DIR}/mlx)
@@ -185,11 +186,9 @@ message(
 # mismatch after an MLX bump fails loudly rather than silently no-op'ing. See
 # each patch file under patches/ for its rationale.
 set(_mlx_patches
-    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_jit_sdk_gate.patch
-    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
-    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch
     ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_sdk_per_platform.patch
     ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_swiftpm_metallib_name.patch
+    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_remove_addrspace_compat.patch
 )
 # In a framework build the delegate is static, so MLX cannot find a colocated
 # metallib and instead loads one from a SwiftPM resource bundle. SWIFTPM_BUNDLE

backends/mlx/patches/mlx_gather_mm_rhs_lda.patch

@@ -0,0 +1,39 @@
+Keep Steel integral constants compatible with older Metal frontends.
+
+MLX supports macOS 14 and selects Metal language 3.1 there. The Metal 4.1
+address-space fix uses metal::remove_addrspace_t, but that trait is unavailable
+in the Xcode 15.4 Metal standard library. Because integral_constant.h is embedded
+in runtime-JIT sources, this prevents kernels from compiling on that supported
+toolchain.
+
+Use remove_addrspace_t on Metal 4.1+, where explicit thread qualifiers make it
+necessary and the trait is available. For older language versions, preserve the
+previous decltype behavior with an identity alias.
+
+Upstream candidate; carried locally until MLX includes the compatibility fix.
+
+diff --git a/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/mlx/backend/metal/kernels/steel/utils/integral_constant.h
+index 2f153f48..97022416 100644
+--- a/mlx/backend/metal/kernels/steel/utils/integral_constant.h
++++ b/mlx/backend/metal/kernels/steel/utils/integral_constant.h
+@@ -47,11 +47,19 @@ using Int = integral_constant<int, val>;
+ // Binary Operators on Integral constants
+ ///////////////////////////////////////////////////////////////////////////////
+ 
++#if __METAL_VERSION__ >= 410
++template <typename T>
++using remove_addrspace_t = metal::remove_addrspace_t<T>;
++#else
++template <typename T>
++using remove_addrspace_t = T;
++#endif
++
+ #define integral_const_binop(__op__, __operator__)          \
+   template <typename T, T tv, typename U, U uv>             \
+   METAL_FUNC constexpr auto __operator__(                   \
+       integral_constant<T, tv>, integral_constant<U, uv>) { \
+     constexpr auto res = tv __op__ uv;                      \
+-    using res_t = metal::remove_addrspace_t<decltype(res)>; \
++    using res_t = remove_addrspace_t<decltype(res)>;        \
+     return integral_constant<res_t, res>{};                 \
+   }

backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch

@@ -293,7 +293,15 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) {
   }
 
   array out = fast::scaled_dot_product_attention(
-      Q, K, V, static_cast<float>(n.scale), mask_mode, mask_arr, sinks, s);
+      Q,
+      K,
+      V,
+      static_cast<float>(n.scale),
+      mask_mode,
+      mask_arr,
+      sinks,
+      false,
+      s);
   st.set_tensor(n.out, std::move(out));
 }
 
@@ -382,6 +390,7 @@ inline void exec_update_and_attend(
       mask_mode,
       spec.mask,
       std::nullopt,
+      false,
       s);
   // Honor the op's output-dtype contract (unset -> SDPA's native output).
   if (n.out_dtype) {