Skip to content

Commit c2b273e

Browse files
metascroymetascroy
authored
[executorch][mlx] bump MLX version to v0.32 (#20970)
As titled
1 parent 9435118 commit c2b273e

5 files changed

Lines changed: 125 additions & 1 deletion

File tree

‎backends/mlx/CMakeLists.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,20 @@ set(_mlx_metallib ${_mlx_binary_dir}/mlx/backend/metal/kernels/mlx.metallib)
180180
message(
181181
STATUS "Building MLX from submodule (ExternalProject): ${MLX_SOURCE_DIR}"
182182
)
183+
# Local patches applied to the MLX submodule before its build. Kept minimal and
184+
# idempotent (patches/apply.sh reverse-checks each before applying); a context
185+
# mismatch after an MLX bump fails loudly rather than silently no-op'ing. See
186+
# each patch file under patches/ for its rationale.
187+
set(_mlx_patches
188+
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch
189+
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
190+
)
183191
ExternalProject_Add(
184192
mlx_external
185193
SOURCE_DIR ${MLX_SOURCE_DIR}
186194
BINARY_DIR ${_mlx_binary_dir}
195+
PATCH_COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh
196+
${MLX_SOURCE_DIR} ${_mlx_patches}
187197
CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
188198
-DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD}
189199
-DCMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}

‎backends/mlx/patches/apply.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
#
3+
# Copyright (c) Meta Platforms, Inc. and affiliates.
4+
# All rights reserved.
5+
#
6+
# This source code is licensed under the BSD-style license found in the
7+
# LICENSE file in the root directory of this source tree.
8+
#
9+
# Idempotently apply local patches to the MLX submodule before its build.
10+
#
11+
# Usage: apply.sh <mlx_source_dir> <patch> [<patch> ...]
12+
#
13+
# For each patch: skip it if it is already applied (reverse-check succeeds),
14+
# otherwise apply it. A patch that neither reverse-applies nor forward-applies
15+
# (e.g. context drift after an MLX bump) fails loudly via `set -e`.
16+
set -euo pipefail
17+
18+
mlx_dir="$1"
19+
shift
20+
21+
for patch in "$@"; do
22+
if git -C "$mlx_dir" apply --reverse --check "$patch" 2>/dev/null; then
23+
echo "MLX patch already applied, skipping: $patch"
24+
else
25+
echo "Applying MLX patch: $patch"
26+
git -C "$mlx_dir" apply --verbose "$patch"
27+
fi
28+
done

‎backends/mlx/patches/mlx_nax_has_include.patch‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
Guard the NAX MetalPerformancePrimitives includes behind __has_include.
2+
3+
MLX's NAX kernels (GEMM and attention) unconditionally include
4+
<MetalPerformancePrimitives/MetalPerformancePrimitives.h>, a framework that
5+
only ships in the macOS 26 / Xcode 26 SDK. On older SDKs the JIT preamble
6+
generator (make_compiled_preamble.sh) runs `metal -E` over these headers, which
7+
fatals on the missing include and fails the build. MLX already gates NAX on the
8+
non-JIT metallib path (kernels/CMakeLists.txt), but the JIT path
9+
(MLX_METAL_JIT=ON, which ExecuTorch uses) is ungated.
10+
11+
This wraps the includes in `#if __has_include(...)` so preprocessing succeeds on
12+
pre-26 SDKs (NAX kernels are never JIT-compiled at runtime there anyway, since
13+
device.cpp:is_nax_available() returns false), while macOS 26 still gets NAX.
14+
15+
Upstream candidate; carried locally until an MLX release gates the JIT path.
16+
17+
diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kernels/steel/gemm/nax.h
18+
--- a/mlx/backend/metal/kernels/steel/gemm/nax.h
19+
+++ b/mlx/backend/metal/kernels/steel/gemm/nax.h
20+
@@ -9,7 +9,9 @@
21+
#include "mlx/backend/metal/kernels/steel/defines.h"
22+
#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
23+
24+
-#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
25+
+#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
26+
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
27+
+#endif
28+
29+
using namespace metal;
30+
31+
diff --git a/mlx/backend/metal/kernels/steel/attn/nax.h b/mlx/backend/metal/kernels/steel/attn/nax.h
32+
--- a/mlx/backend/metal/kernels/steel/attn/nax.h
33+
+++ b/mlx/backend/metal/kernels/steel/attn/nax.h
34+
@@ -9,7 +9,9 @@
35+
#include "mlx/backend/metal/kernels/steel/defines.h"
36+
#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
37+
38+
-#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
39+
+#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
40+
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
41+
+#endif
42+
43+
using namespace metal;
44+

‎backends/mlx/patches/mlx_qmm_splitk_bk_align.patch‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
Align split-K partitions to the qmm K-tile (BK=32), fixing nvfp4.
2+
3+
MLX v0.32.0's qmm_splitk caps split_k only by the quantization group count
4+
(K / group_size), not by the kernel's K-tile width BK (=32). For nvfp4
5+
(group_size=16) this yields a per-partition K of 16 < BK, and fp_qmm_t_splitk's
6+
tile load reads a full BK-wide K-tile with no K bound -- spilling 16 columns
7+
past the partition into the next group's packed weights and fp8 scales. That
8+
corrupts every partial (non-uniform ~2x error) and reads past the buffer on the
9+
last partition (NaN/inf). Affine (group_size >= 32) is unaffected because its
10+
partitions are already >= BK.
11+
12+
Fix in the dispatch (MLX's pattern for tile-alignment constraints): require each
13+
K partition to be a whole number of BK-wide tiles as well as whole groups, i.e.
14+
align split_k to max(group_size, BK). Only changes behavior for group_size < 32.
15+
16+
Upstream candidate.
17+
18+
diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp
19+
index 62d48714..94c56307 100644
20+
--- a/mlx/backend/metal/quantized.cpp
21+
+++ b/mlx/backend/metal/quantized.cpp
22+
@@ -884,11 +884,15 @@ void qmm_splitk(
23+
int current_tgs = n_tiles * m_tiles;
24+
int split_k = std::max(1, 512 / current_tgs);
25+
26+
- // Cap split_k by the number of quantization groups
27+
- split_k = std::min(split_k, K / group_size);
28+
-
29+
- // Ensure K divides evenly by split_k * group_size
30+
- while (split_k > 1 && (K % (split_k * group_size) != 0)) {
31+
+ // Each K partition must be a whole number of BK-wide (32) K-tiles as well as
32+
+ // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do
33+
+ // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's
34+
+ // group_size=16) would over-read into the next group's weights/scales.
35+
+ int k_align = group_size > 32 ? group_size : 32;
36+
+ split_k = std::min(split_k, K / k_align);
37+
+
38+
+ // Ensure K divides evenly by split_k * k_align
39+
+ while (split_k > 1 && (K % (split_k * k_align) != 0)) {
40+
split_k--;
41+
}
42+
if (split_k <= 1) {

‎backends/mlx/third-party/mlx‎

Submodule mlx updated 381 files

0 commit comments

Comments
 (0)