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

@@ -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 \

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()

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 <algorithm>
+ #include <cassert>
++#include <chrono>
+ #include <cstdio>
+ #include <sstream>
+ #include <stdexcept>
+@@ -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<std::mutex> 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<std::mutex> lock(mtx_);
++      stop_refresh_ = true;
++    }
++    refresh_cv_.notify_one();
++    refresh_thread_.join();
++  }
++}
+ 
+ bool ResidencySets::add_set_locked(NS::Error** error_out) {
+   NS::SharedPtr<MTL::ResidencySet> 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 <atomic>
++#include <condition_variable>
+ #include <cstdint>
+ #include <mutex>
++#include <thread>
+ #include <unordered_map>
+ #include <vector>
+ 
+@@ -118,6 +120,9 @@
+   // taking the lock.
+   std::atomic<uint64_t> num_sets_{0};
+   mutable std::mutex mtx_;
++  std::condition_variable refresh_cv_;
++  bool stop_refresh_{false};
++  std::thread refresh_thread_;
+ };
+ 
+ } // namespace mlx::core::metal

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