diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 15db7ca0cb0..c6eb2ac3156 100644 --- a/backends/mlx/CMakeLists.txt +++ b/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} diff --git a/backends/mlx/patches/apply.sh b/backends/mlx/patches/apply.sh new file mode 100644 index 00000000000..f834ce59377 --- /dev/null +++ b/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 [ ...] +# +# 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 diff --git a/backends/mlx/patches/mlx_nax_has_include.patch b/backends/mlx/patches/mlx_nax_has_include.patch new file mode 100644 index 00000000000..2c917353146 --- /dev/null +++ b/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 +, 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 ++#if __has_include() ++#include ++#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 ++#if __has_include() ++#include ++#endif + + using namespace metal; + diff --git a/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch b/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch new file mode 100644 index 00000000000..f4405b5ae9f --- /dev/null +++ b/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) { diff --git a/backends/mlx/third-party/mlx b/backends/mlx/third-party/mlx index ce45c52505c..7a1d4f5c12a 160000 --- a/backends/mlx/third-party/mlx +++ b/backends/mlx/third-party/mlx @@ -1 +1 @@ -Subproject commit ce45c52505c8158ea48d2a54e8caae05efd86bfe +Subproject commit 7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247