diff --git a/MLX_VERSION b/MLX_VERSION index 74940ba8f..99a1ec20c 100644 --- a/MLX_VERSION +++ b/MLX_VERSION @@ -1 +1 @@ -adf21deabd32c4fe3726ab2d373ff2165c3f27d6 \ No newline at end of file +7f062ddcb81bedb537f3f95d11a7343ba6b5f909 \ No newline at end of file diff --git a/cmake/local.cmake b/cmake/local.cmake index f16518f81..7b65e7693 100644 --- a/cmake/local.cmake +++ b/cmake/local.cmake @@ -169,6 +169,14 @@ if(OLLAMA_MLX_BACKENDS) list(APPEND _mlx_source_targets ollama-mlx-source) endif() + # Temporary MLX-C carry patch: regenerated bindings for force_fused and the + # thread-local compile cache, carried until they merge upstream into + # ml-explore/mlx-c. Then bump MLX_C_VERSION and delete mlx/compat/. + find_package(Git REQUIRED) + set(OLLAMA_MLX_C_COMPAT_PATCH_COMMAND + ${GIT_EXECUTABLE} apply ${CMAKE_SOURCE_DIR}/mlx/compat/0001-mlx-c-regen-0.32.1.patch + CACHE INTERNAL "MLX-C carry patch") + if(DEFINED "FETCHCONTENT_SOURCE_DIR_MLX-C" AND NOT "${FETCHCONTENT_SOURCE_DIR_MLX-C}" STREQUAL "") get_filename_component(OLLAMA_MLX_C_SOURCE_DIR "${FETCHCONTENT_SOURCE_DIR_MLX-C}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}") @@ -188,7 +196,9 @@ if(OLLAMA_MLX_BACKENDS) CONFIGURE_COMMAND "" BUILD_COMMAND "" INSTALL_COMMAND "" - USES_TERMINAL_DOWNLOAD TRUE) + PATCH_COMMAND ${OLLAMA_MLX_C_COMPAT_PATCH_COMMAND} + USES_TERMINAL_DOWNLOAD TRUE + USES_TERMINAL_PATCH TRUE) list(APPEND _mlx_source_targets ollama-mlx-c-source) endif() # Refresh the vendored MLX-C headers once the sources are present. Every MLX diff --git a/mlx/compat/0001-mlx-c-regen-0.32.1.patch b/mlx/compat/0001-mlx-c-regen-0.32.1.patch new file mode 100644 index 000000000..d97813685 --- /dev/null +++ b/mlx/compat/0001-mlx-c-regen-0.32.1.patch @@ -0,0 +1,88 @@ +diff --git a/mlx/c/compile.cpp b/mlx/c/compile.cpp +index bc7b864..260ee86 100644 +--- a/mlx/c/compile.cpp ++++ b/mlx/c/compile.cpp +@@ -42,22 +42,25 @@ extern "C" int mlx_detail_compile( + } + extern "C" int mlx_detail_compile_clear_cache(void) { + try { +- mlx::core::detail::compile_clear_cache(); ++ mlx::core::detail::compile_clear_cache(mlx::core::detail::compile_cache()); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; + } ++ + extern "C" int mlx_detail_compile_erase(uintptr_t fun_id) { + try { +- mlx::core::detail::compile_erase(fun_id); ++ mlx::core::detail::compile_erase( ++ mlx::core::detail::compile_cache(), fun_id); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; + } ++ + extern "C" int mlx_disable_compile(void) { + try { + mlx::core::disable_compile(); +diff --git a/mlx/c/compile.h b/mlx/c/compile.h +index 2892a1f..ff3abc4 100644 +--- a/mlx/c/compile.h ++++ b/mlx/c/compile.h +@@ -44,7 +44,9 @@ int mlx_detail_compile( + const uint64_t* constants, + size_t constants_num); + int mlx_detail_compile_clear_cache(void); ++ + int mlx_detail_compile_erase(uintptr_t fun_id); ++ + int mlx_disable_compile(void); + int mlx_enable_compile(void); + int mlx_set_compile_mode(mlx_compile_mode mode); +diff --git a/mlx/c/fast.cpp b/mlx/c/fast.cpp +index ac17d54..eb35d16 100644 +--- a/mlx/c/fast.cpp ++++ b/mlx/c/fast.cpp +@@ -3,8 +3,8 @@ + /* This file is auto-generated. Do not edit manually. */ + /* */ + +-#include "mlx/c/fast.h" + #include "mlx/c/error.h" ++#include "mlx/c/fast.h" + #include "mlx/c/private/mlx.h" + #include "mlx/fast.h" + +@@ -610,6 +610,7 @@ extern "C" int mlx_fast_scaled_dot_product_attention( + const char* mask_mode, + const mlx_array mask_arr /* may be null */, + const mlx_array sinks /* may be null */, ++ bool force_fused, + const mlx_stream s) { + try { + mlx_array_set_( +@@ -624,6 +625,7 @@ extern "C" int mlx_fast_scaled_dot_product_attention( + : std::nullopt), + (sinks.ctx ? std::make_optional(mlx_array_get_(sinks)) + : std::nullopt), ++ force_fused, + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); +diff --git a/mlx/c/fast.h b/mlx/c/fast.h +index c825d00..e6c5bdf 100644 +--- a/mlx/c/fast.h ++++ b/mlx/c/fast.h +@@ -195,6 +195,7 @@ int mlx_fast_scaled_dot_product_attention( + const char* mask_mode, + const mlx_array mask_arr /* may be null */, + const mlx_array sinks /* may be null */, ++ bool force_fused, + const mlx_stream s); + + /**@}*/ diff --git a/x/mlxrunner/mlx/fast.go b/x/mlxrunner/mlx/fast.go index f38a670df..27d5724ed 100644 --- a/x/mlxrunner/mlx/fast.go +++ b/x/mlxrunner/mlx/fast.go @@ -21,7 +21,7 @@ func FastScaledDotProductAttention(q, k, v *Array, scale float32, mode string, m } out := New("FAST_SDPA") - C.mlx_fast_scaled_dot_product_attention(&out.ctx, q.ctx, k.ctx, v.ctx, C.float(scale), cMode, maskCtx, sinks.ctx, DefaultStream().ctx) + C.mlx_fast_scaled_dot_product_attention(&out.ctx, q.ctx, k.ctx, v.ctx, C.float(scale), cMode, maskCtx, sinks.ctx, C.bool(false), DefaultStream().ctx) return out } diff --git a/x/mlxrunner/mlx/generated.c b/x/mlxrunner/mlx/generated.c index 8147fe5b4..a775149a9 100644 --- a/x/mlxrunner/mlx/generated.c +++ b/x/mlxrunner/mlx/generated.c @@ -518,6 +518,7 @@ int (*mlx_fast_scaled_dot_product_attention_)( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s) = NULL; int (*mlx_fft_fft_)( mlx_array* res, diff --git a/x/mlxrunner/mlx/generated.h b/x/mlxrunner/mlx/generated.h index 02beed9d0..592aede4e 100644 --- a/x/mlxrunner/mlx/generated.h +++ b/x/mlxrunner/mlx/generated.h @@ -1761,6 +1761,7 @@ extern int (*mlx_fast_scaled_dot_product_attention_)( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s); extern int (*mlx_fft_fft_)( mlx_array* res, @@ -4518,8 +4519,9 @@ static inline int mlx_fast_scaled_dot_product_attention( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s) { - return mlx_fast_scaled_dot_product_attention_(res, queries, keys, values, scale, mask_mode, mask_arr, sinks, s); + return mlx_fast_scaled_dot_product_attention_(res, queries, keys, values, scale, mask_mode, mask_arr, sinks, force_fused, s); } static inline int mlx_fft_fft( mlx_array* res, diff --git a/x/mlxrunner/mlx/include/mlx/c/compile.h b/x/mlxrunner/mlx/include/mlx/c/compile.h index 2892a1f0b..ff3abc4ee 100644 --- a/x/mlxrunner/mlx/include/mlx/c/compile.h +++ b/x/mlxrunner/mlx/include/mlx/c/compile.h @@ -44,7 +44,9 @@ int mlx_detail_compile( const uint64_t* constants, size_t constants_num); int mlx_detail_compile_clear_cache(void); + int mlx_detail_compile_erase(uintptr_t fun_id); + int mlx_disable_compile(void); int mlx_enable_compile(void); int mlx_set_compile_mode(mlx_compile_mode mode); diff --git a/x/mlxrunner/mlx/include/mlx/c/fast.h b/x/mlxrunner/mlx/include/mlx/c/fast.h index c825d00e5..e6c5bdf62 100644 --- a/x/mlxrunner/mlx/include/mlx/c/fast.h +++ b/x/mlxrunner/mlx/include/mlx/c/fast.h @@ -195,6 +195,7 @@ int mlx_fast_scaled_dot_product_attention( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s); /**@}*/