[executorch][mlx] bump MLX version to v0.32

backends/mlx/CMakeLists.txt

@@ -180,10 +180,20 @@ set(_mlx_metallib ${_mlx_binary_dir}/mlx/backend/metal/kernels/mlx.metallib)
 message(
   STATUS "Building MLX from submodule (ExternalProject): ${MLX_SOURCE_DIR}"
 )
+# Local patches applied to the MLX submodule before its build. Kept minimal and
+# idempotent (patches/apply.sh reverse-checks each before applying); a context
+# 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_has_include.patch
+    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
+)
 ExternalProject_Add(
   mlx_external
   SOURCE_DIR ${MLX_SOURCE_DIR}
   BINARY_DIR ${_mlx_binary_dir}
+  PATCH_COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh
+                ${MLX_SOURCE_DIR} ${_mlx_patches}
   CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
              -DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD}
              -DCMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}

backends/mlx/patches/apply.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Idempotently apply local patches to the MLX submodule before its build.
+#
+# Usage: apply.sh <mlx_source_dir> <patch> [<patch> ...]
+#
+# For each patch: skip it if it is already applied (reverse-check succeeds),
+# otherwise apply it. A patch that neither reverse-applies nor forward-applies
+# (e.g. context drift after an MLX bump) fails loudly via `set -e`.
+set -euo pipefail
+
+mlx_dir="$1"
+shift
+
+for patch in "$@"; do
+  if git -C "$mlx_dir" apply --reverse --check "$patch" 2>/dev/null; then
+    echo "MLX patch already applied, skipping: $patch"
+  else
+    echo "Applying MLX patch: $patch"
+    git -C "$mlx_dir" apply --verbose "$patch"
+  fi
+done

backends/mlx/patches/mlx_nax_has_include.patch

@@ -0,0 +1,44 @@
+Guard the NAX MetalPerformancePrimitives includes behind __has_include.
+
+MLX's NAX kernels (GEMM and attention) unconditionally include
+<MetalPerformancePrimitives/MetalPerformancePrimitives.h>, a framework that
+only ships in the macOS 26 / Xcode 26 SDK. On older SDKs the JIT preamble
+generator (make_compiled_preamble.sh) runs `metal -E` over these headers, which
+fatals on the missing include and fails the build. MLX already gates NAX on the
+non-JIT metallib path (kernels/CMakeLists.txt), but the JIT path
+(MLX_METAL_JIT=ON, which ExecuTorch uses) is ungated.
+
+This wraps the includes in `#if __has_include(...)` so preprocessing succeeds on
+pre-26 SDKs (NAX kernels are never JIT-compiled at runtime there anyway, since
+device.cpp:is_nax_available() returns false), while macOS 26 still gets NAX.
+
+Upstream candidate; carried locally until an MLX release gates the JIT path.
+
+diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kernels/steel/gemm/nax.h
+--- a/mlx/backend/metal/kernels/steel/gemm/nax.h
++++ b/mlx/backend/metal/kernels/steel/gemm/nax.h
+@@ -9,7 +9,9 @@
+ #include "mlx/backend/metal/kernels/steel/defines.h"
+ #include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+-#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
++#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
++#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
++#endif
+
+ using namespace metal;
+
+diff --git a/mlx/backend/metal/kernels/steel/attn/nax.h b/mlx/backend/metal/kernels/steel/attn/nax.h
+--- a/mlx/backend/metal/kernels/steel/attn/nax.h
++++ b/mlx/backend/metal/kernels/steel/attn/nax.h
+@@ -9,7 +9,9 @@
+ #include "mlx/backend/metal/kernels/steel/defines.h"
+ #include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+-#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
++#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
++#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
++#endif
+
+ using namespace metal;
+

backends/mlx/patches/mlx_qmm_splitk_bk_align.patch

@@ -0,0 +1,42 @@
+Align split-K partitions to the qmm K-tile (BK=32), fixing nvfp4.
+
+MLX v0.32.0's qmm_splitk caps split_k only by the quantization group count
+(K / group_size), not by the kernel's K-tile width BK (=32). For nvfp4
+(group_size=16) this yields a per-partition K of 16 < BK, and fp_qmm_t_splitk's
+tile load reads a full BK-wide K-tile with no K bound -- spilling 16 columns
+past the partition into the next group's packed weights and fp8 scales. That
+corrupts every partial (non-uniform ~2x error) and reads past the buffer on the
+last partition (NaN/inf). Affine (group_size >= 32) is unaffected because its
+partitions are already >= BK.
+
+Fix in the dispatch (MLX's pattern for tile-alignment constraints): require each
+K partition to be a whole number of BK-wide tiles as well as whole groups, i.e.
+align split_k to max(group_size, BK). Only changes behavior for group_size < 32.
+
+Upstream candidate.
+
+diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp
+index 62d48714..94c56307 100644
+--- a/mlx/backend/metal/quantized.cpp
++++ b/mlx/backend/metal/quantized.cpp
+@@ -884,11 +884,15 @@ void qmm_splitk(
+   int current_tgs = n_tiles * m_tiles;
+   int split_k = std::max(1, 512 / current_tgs);
+
+-  // Cap split_k by the number of quantization groups
+-  split_k = std::min(split_k, K / group_size);
+-
+-  // Ensure K divides evenly by split_k * group_size
+-  while (split_k > 1 && (K % (split_k * group_size) != 0)) {
++  // Each K partition must be a whole number of BK-wide (32) K-tiles as well as
++  // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do
++  // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's
++  // group_size=16) would over-read into the next group's weights/scales.
++  int k_align = group_size > 32 ? group_size : 32;
++  split_k = std::min(split_k, K / k_align);
++
++  // Ensure K divides evenly by split_k * k_align
++  while (split_k > 1 && (K % (split_k * k_align) != 0)) {
+     split_k--;
+   }
+   if (split_k <= 1) {