From cf6c9de620caac8c3a1167ed476ec4773617ce40 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 5 Oct 2026 16:49:32 -0700 Subject: [PATCH] mlx: mitigate high latency after GPU idle (#18807) Carry the MLX residency-refresh patch and enable it every second in the runner, preserving explicit environment overrides. Fixes #18744. --- Dockerfile | 4 +- cmake/local.cmake | 37 +++++++-- .../mlx/0001-metal-residency-refresh.patch | 83 +++++++++++++++++++ mlxrunner/client.go | 7 ++ 4 files changed, 120 insertions(+), 11 deletions(-) create mode 100644 mlx/compat/mlx/0001-metal-residency-refresh.patch diff --git a/Dockerfile b/Dockerfile index 7099e1c2a..118082ba8 100644 --- a/Dockerfile +++ b/Dockerfile @@ -235,8 +235,8 @@ RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux- ENV PATH=/usr/local/go/bin:$PATH RUN go mod download RUN --mount=type=cache,target=/root/.ccache \ - --mount=type=bind,from=local-mlx,target=/tmp/local-mlx \ - --mount=type=bind,from=local-mlx-c,target=/tmp/local-mlx-c \ + --mount=type=bind,from=local-mlx,target=/tmp/local-mlx,rw \ + --mount=type=bind,from=local-mlx-c,target=/tmp/local-mlx-c,rw \ if [ -f /tmp/local-mlx/CMakeLists.txt ]; then \ export OLLAMA_MLX_SOURCE=/tmp/local-mlx; \ fi \ diff --git a/cmake/local.cmake b/cmake/local.cmake index 2fb9b8429..b2dad3adb 100644 --- a/cmake/local.cmake +++ b/cmake/local.cmake @@ -146,14 +146,26 @@ if(OLLAMA_MLX_BACKENDS) file(READ "${CMAKE_SOURCE_DIR}/MLX_C_VERSION" OLLAMA_MLX_C_GIT_TAG) string(STRIP "${OLLAMA_MLX_C_GIT_TAG}" OLLAMA_MLX_C_GIT_TAG) + # Apply carried MLX and MLX-C patches to fetched sources and local overrides. + find_package(Git REQUIRED) + set(OLLAMA_MLX_COMPAT_PATCH_COMMAND + ${CMAKE_COMMAND} + -DPATCH_DIR=${CMAKE_SOURCE_DIR}/mlx/compat/mlx + -DPATCH_LABEL=mlx/compat/mlx + -P ${CMAKE_SOURCE_DIR}/cmake/apply-git-patches.cmake + CACHE INTERNAL "MLX carry patches") + + set(_mlx_source_override FALSE) if(DEFINED FETCHCONTENT_SOURCE_DIR_MLX AND NOT "${FETCHCONTENT_SOURCE_DIR_MLX}" STREQUAL "") get_filename_component(OLLAMA_MLX_SOURCE_DIR "${FETCHCONTENT_SOURCE_DIR_MLX}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}") message(STATUS "Using MLX source override: ${OLLAMA_MLX_SOURCE_DIR}") + set(_mlx_source_override TRUE) elseif(DEFINED ENV{OLLAMA_MLX_SOURCE}) get_filename_component(OLLAMA_MLX_SOURCE_DIR "$ENV{OLLAMA_MLX_SOURCE}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}") message(STATUS "Using local MLX source: ${OLLAMA_MLX_SOURCE_DIR}") + set(_mlx_source_override TRUE) else() set(OLLAMA_MLX_SOURCE_DIR "${CMAKE_BINARY_DIR}/_deps/mlx-src") ExternalProject_Add(ollama-mlx-source @@ -165,20 +177,27 @@ if(OLLAMA_MLX_BACKENDS) CONFIGURE_COMMAND "" BUILD_COMMAND "" INSTALL_COMMAND "" - USES_TERMINAL_DOWNLOAD TRUE) + PATCH_COMMAND ${OLLAMA_MLX_COMPAT_PATCH_COMMAND} + USES_TERMINAL_DOWNLOAD TRUE + USES_TERMINAL_PATCH TRUE) list(APPEND _mlx_source_targets ollama-mlx-source) endif() - # Temporary MLX-C carry patch for gather_qmm global-scale support, carried - # until it merges upstream into ml-explore/mlx-c. Then bump MLX_C_VERSION - # and delete mlx/compat/. - find_package(Git REQUIRED) + if(_mlx_source_override) + add_custom_target(ollama-mlx-override-patch + COMMAND ${OLLAMA_MLX_COMPAT_PATCH_COMMAND} + WORKING_DIRECTORY ${OLLAMA_MLX_SOURCE_DIR} + COMMENT "Applying carried MLX patches to ${OLLAMA_MLX_SOURCE_DIR}" + VERBATIM) + list(APPEND _mlx_source_targets ollama-mlx-override-patch) + endif() + set(OLLAMA_MLX_C_COMPAT_PATCH_COMMAND ${CMAKE_COMMAND} -DPATCH_DIR=${CMAKE_SOURCE_DIR}/mlx/compat/mlx-c -DPATCH_LABEL=mlx/compat/mlx-c -P ${CMAKE_SOURCE_DIR}/cmake/apply-git-patches.cmake - CACHE INTERNAL "MLX-C carry patch") + CACHE INTERNAL "MLX-C carry patches") set(_mlx_c_source_override FALSE) if(DEFINED "FETCHCONTENT_SOURCE_DIR_MLX-C" AND NOT "${FETCHCONTENT_SOURCE_DIR_MLX-C}" STREQUAL "") @@ -208,13 +227,13 @@ if(OLLAMA_MLX_BACKENDS) list(APPEND _mlx_source_targets ollama-mlx-c-source) endif() if(_mlx_c_source_override) - # Source overrides bypass the ExternalProject patch step, so the carry - # patch has to be applied to the override checkout as well. The + # Source overrides bypass the ExternalProject patch step, so patches + # have to be applied to the override checkout as well. The # applier is idempotent, which keeps repeated builds safe. add_custom_target(ollama-mlx-c-override-patch COMMAND ${OLLAMA_MLX_C_COMPAT_PATCH_COMMAND} WORKING_DIRECTORY ${OLLAMA_MLX_C_SOURCE_DIR} - COMMENT "Applying MLX-C compat patches to ${OLLAMA_MLX_C_SOURCE_DIR}" + COMMENT "Applying carried MLX-C patches to ${OLLAMA_MLX_C_SOURCE_DIR}" VERBATIM) list(APPEND _mlx_source_targets ollama-mlx-c-override-patch) endif() diff --git a/mlx/compat/mlx/0001-metal-residency-refresh.patch b/mlx/compat/mlx/0001-metal-residency-refresh.patch new file mode 100644 index 000000000..31e45b388 --- /dev/null +++ b/mlx/compat/mlx/0001-metal-residency-refresh.patch @@ -0,0 +1,83 @@ +diff --git a/mlx/backend/metal/resident.cpp b/mlx/backend/metal/resident.cpp +--- a/mlx/backend/metal/resident.cpp ++++ b/mlx/backend/metal/resident.cpp +@@ -1,7 +1,8 @@ +-// Copyright © 2024 Apple Inc. ++// Copyright © 2024-2026 Apple Inc. + + #include + #include ++#include + #include + #include + #include +@@ -42,10 +43,38 @@ + } + throw std::runtime_error(msg.str()); + } +- } +-} +- +-ResidencySets::~ResidencySets() = default; ++ if (int interval = ++ env::get_var("MLX_METAL_RESIDENCY_REFRESH_INTERVAL_MS", 0); ++ interval > 0) { ++ refresh_thread_ = std::thread([this, interval] { ++ std::unique_lock lock(mtx_); ++ while (!refresh_cv_.wait_for( ++ lock, std::chrono::milliseconds(interval), [this] { ++ return stop_refresh_; ++ })) { ++ auto pool = new_scoped_memory_pool(); ++ // macOS can drop residency after GPU idle despite a standing request. ++ for (auto& s : sets_) { ++ if (s.size != 0) { ++ s.set->requestResidency(); ++ } ++ } ++ } ++ }); ++ } ++ } ++} ++ ++ResidencySets::~ResidencySets() { ++ if (refresh_thread_.joinable()) { ++ { ++ std::lock_guard lock(mtx_); ++ stop_refresh_ = true; ++ } ++ refresh_cv_.notify_one(); ++ refresh_thread_.join(); ++ } ++} + + bool ResidencySets::add_set_locked(NS::Error** error_out) { + NS::SharedPtr set; +diff --git a/mlx/backend/metal/resident.h b/mlx/backend/metal/resident.h +--- a/mlx/backend/metal/resident.h ++++ b/mlx/backend/metal/resident.h +@@ -1,10 +1,12 @@ +-// Copyright © 2024 Apple Inc. ++// Copyright © 2024-2026 Apple Inc. + + #pragma once + + #include ++#include + #include + #include ++#include + #include + #include + +@@ -118,6 +120,9 @@ + // taking the lock. + std::atomic num_sets_{0}; + mutable std::mutex mtx_; ++ std::condition_variable refresh_cv_; ++ bool stop_refresh_{false}; ++ std::thread refresh_thread_; + }; + + } // namespace mlx::core::metal diff --git a/mlxrunner/client.go b/mlxrunner/client.go index aa03ccd34..9cfb84775 100644 --- a/mlxrunner/client.go +++ b/mlxrunner/client.go @@ -370,6 +370,13 @@ func (c *Client) Load(ctx context.Context, systemInfo ml.SystemInfo, gpus []ml.D cmd := exec.Command(exe, "runner", "--model", c.modelName, "--port", strconv.Itoa(port)) cmd.Env = os.Environ() + // Keep Metal weights resident between requests until MLX fixes idle eviction. + if runtime.GOOS == "darwin" { + if _, ok := os.LookupEnv("MLX_METAL_RESIDENCY_REFRESH_INTERVAL_MS"); !ok { + setEnv(cmd, "MLX_METAL_RESIDENCY_REFRESH_INTERVAL_MS", "1000") + } + } + // Set library path environment variable for MLX libraries // Linux: LD_LIBRARY_PATH, Windows: PATH var libPathEnvVar string