From 8c7c29ac9eac2d5f26b6d140dcc250e7641a1d7c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 15 Jul 2026 17:52:39 -0700 Subject: [PATCH 1/5] up --- backends/mlx/third-party/mlx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 2c4cb9f021c217319db0ccbbe3266d74332d4c30 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 16 Jul 2026 11:17:04 -0700 Subject: [PATCH 2/5] up --- backends/mlx/CMakeLists.txt | 10 +++++++ .../mlx/patches/mlx_nax_has_include.patch | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 backends/mlx/patches/mlx_nax_has_include.patch diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 15db7ca0cb0..e8b9c083214 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 (reverse-check skips an already-applied patch); 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_nax_patch + ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch +) ExternalProject_Add( mlx_external SOURCE_DIR ${MLX_SOURCE_DIR} BINARY_DIR ${_mlx_binary_dir} + PATCH_COMMAND + bash -c + "git -C '${MLX_SOURCE_DIR}' apply --reverse --check '${_mlx_nax_patch}' 2>/dev/null || git -C '${MLX_SOURCE_DIR}' apply --verbose '${_mlx_nax_patch}'" 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/mlx_nax_has_include.patch b/backends/mlx/patches/mlx_nax_has_include.patch new file mode 100644 index 00000000000..1556fbf368d --- /dev/null +++ b/backends/mlx/patches/mlx_nax_has_include.patch @@ -0,0 +1,30 @@ +Guard the NAX MetalPerformancePrimitives include behind __has_include. + +MLX's NAX GEMM header unconditionally includes +, 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 this header, 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 include 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; + From 7e0717eb68666e5e7682499de867c6f160b76580 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 16 Jul 2026 11:43:29 -0700 Subject: [PATCH 3/5] up --- .../mlx/patches/mlx_nax_has_include.patch | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/backends/mlx/patches/mlx_nax_has_include.patch b/backends/mlx/patches/mlx_nax_has_include.patch index 1556fbf368d..2c917353146 100644 --- a/backends/mlx/patches/mlx_nax_has_include.patch +++ b/backends/mlx/patches/mlx_nax_has_include.patch @@ -1,14 +1,14 @@ -Guard the NAX MetalPerformancePrimitives include behind __has_include. +Guard the NAX MetalPerformancePrimitives includes behind __has_include. -MLX's NAX GEMM header unconditionally includes +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 this header, which +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 include in `#if __has_include(...)` so preprocessing succeeds on +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. @@ -28,3 +28,17 @@ diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kern 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; + From e4b22b3833fbe27cb3c3f0feb307b2909ef9a5df Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 16 Jul 2026 12:27:32 -0700 Subject: [PATCH 4/5] up --- backends/mlx/CMakeLists.txt | 15 +++++----- backends/mlx/patches/apply.sh | 28 +++++++++++++++++++ .../mlx/patches/mlx_nvfp4_no_splitk.patch | 25 +++++++++++++++++ 3 files changed, 60 insertions(+), 8 deletions(-) create mode 100644 backends/mlx/patches/apply.sh create mode 100644 backends/mlx/patches/mlx_nvfp4_no_splitk.patch diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index e8b9c083214..786fee0c199 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -181,19 +181,18 @@ message( STATUS "Building MLX from submodule (ExternalProject): ${MLX_SOURCE_DIR}" ) # Local patches applied to the MLX submodule before its build. Kept minimal and -# idempotent (reverse-check skips an already-applied patch); 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_nax_patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch +# 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_nvfp4_no_splitk.patch ) ExternalProject_Add( mlx_external SOURCE_DIR ${MLX_SOURCE_DIR} BINARY_DIR ${_mlx_binary_dir} - PATCH_COMMAND - bash -c - "git -C '${MLX_SOURCE_DIR}' apply --reverse --check '${_mlx_nax_patch}' 2>/dev/null || git -C '${MLX_SOURCE_DIR}' apply --verbose '${_mlx_nax_patch}'" + 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_nvfp4_no_splitk.patch b/backends/mlx/patches/mlx_nvfp4_no_splitk.patch new file mode 100644 index 00000000000..b06a765e3b4 --- /dev/null +++ b/backends/mlx/patches/mlx_nvfp4_no_splitk.patch @@ -0,0 +1,25 @@ +Do not use the split-K quantized matmul for the nvfp4 mode. + +MLX v0.32.0 added a split-K route (qmm_splitk / fp_qmm_t_splitk) selected in +QuantizedMatmul::eval_gpu for M >= vector_limit with transposed weights and a +single batch. The nvfp4 (fp4 + fp8-e4m3 block scale) variant of that path +produces incorrect results (non-uniform ~2x error plus NaN/inf), while the +non-split fp_qmm_t path is correct. Gate the split-K selection so nvfp4 falls +back to the proven qmm path. + +Upstream candidate; carried locally until MLX fixes fp4 split-K. + +diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp +--- a/mlx/backend/metal/quantized.cpp ++++ b/mlx/backend/metal/quantized.cpp +@@ -1511,7 +1511,9 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { + if (M >= vector_limit) { + // Use split-K qmm for small M with transposed weights (non-batched only) + int B = out.size() / M / N; +- if (transpose_ && B == 1) { ++ // nvfp4 split-K (fp_qmm_t_splitk) produces incorrect results (NaN + ~2x ++ // error); route the fp4 tensor-scale mode through the non-split qmm path. ++ if (transpose_ && B == 1 && mode_ != QuantizationMode::Nvfp4) { + qmm_splitk( + x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + return; From 88ab9438370ad8d15579b22aaa6fb788c8c4ceb5 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 16 Jul 2026 14:21:38 -0700 Subject: [PATCH 5/5] up --- backends/mlx/CMakeLists.txt | 5 ++- .../mlx/patches/mlx_nvfp4_no_splitk.patch | 25 ----------- .../mlx/patches/mlx_qmm_splitk_bk_align.patch | 42 +++++++++++++++++++ 3 files changed, 45 insertions(+), 27 deletions(-) delete mode 100644 backends/mlx/patches/mlx_nvfp4_no_splitk.patch create mode 100644 backends/mlx/patches/mlx_qmm_splitk_bk_align.patch diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 786fee0c199..c6eb2ac3156 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -184,8 +184,9 @@ message( # 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_nvfp4_no_splitk.patch +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 diff --git a/backends/mlx/patches/mlx_nvfp4_no_splitk.patch b/backends/mlx/patches/mlx_nvfp4_no_splitk.patch deleted file mode 100644 index b06a765e3b4..00000000000 --- a/backends/mlx/patches/mlx_nvfp4_no_splitk.patch +++ /dev/null @@ -1,25 +0,0 @@ -Do not use the split-K quantized matmul for the nvfp4 mode. - -MLX v0.32.0 added a split-K route (qmm_splitk / fp_qmm_t_splitk) selected in -QuantizedMatmul::eval_gpu for M >= vector_limit with transposed weights and a -single batch. The nvfp4 (fp4 + fp8-e4m3 block scale) variant of that path -produces incorrect results (non-uniform ~2x error plus NaN/inf), while the -non-split fp_qmm_t path is correct. Gate the split-K selection so nvfp4 falls -back to the proven qmm path. - -Upstream candidate; carried locally until MLX fixes fp4 split-K. - -diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp ---- a/mlx/backend/metal/quantized.cpp -+++ b/mlx/backend/metal/quantized.cpp -@@ -1511,7 +1511,9 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { - if (M >= vector_limit) { - // Use split-K qmm for small M with transposed weights (non-batched only) - int B = out.size() / M / N; -- if (transpose_ && B == 1) { -+ // nvfp4 split-K (fp_qmm_t_splitk) produces incorrect results (NaN + ~2x -+ // error); route the fp4 tensor-scale mode through the non-split qmm path. -+ if (transpose_ && B == 1 && mode_ != QuantizationMode::Nvfp4) { - qmm_splitk( - x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); - return; 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) {