From 9c2f0f974ffb5e7c24be58eee93ddf58532fd069 Mon Sep 17 00:00:00 2001
From: Corentin Kerisit <corentin.kerisit@gmail.com>
Date: Fri, 16 Oct 2026 21:11:46 +0200
Subject: [PATCH 1/2] Allow configuring the CPU JIT cache directory

---
 docs/src/usage/environment_variables.rst |  9 +++++++++
 mlx/backend/cpu/compiled.cpp             | 13 ++++++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst
index 95ee3bdaa1..1aabecf00d 100644
--- a/docs/src/usage/environment_variables.rst
+++ b/docs/src/usage/environment_variables.rst
@@ -26,6 +26,15 @@ General
    on supported hardware. The default is ``1``. Set it to ``0`` to keep these
    operations in full ``float32`` precision. See :doc:`precision`.
 
+CPU
+---
+
+.. envvar:: MLX_CPU_CACHE_DIR
+
+   Override the directory used to cache runtime-compiled CPU kernels. An unset
+   or empty value uses ``mlx/<version>/cpu`` under the system temporary
+   directory, which is the default.
+
 Distributed
 -----------
 
diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp
index d6a53a5287..5ef78cba30 100644
--- a/mlx/backend/cpu/compiled.cpp
+++ b/mlx/backend/cpu/compiled.cpp
@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include <dlfcn.h>
+#include <cstdlib>
 #include <filesystem>
 #include <fstream>
 #include <list>
@@ -98,11 +99,13 @@ void* compile(
     kernel_file_name = kernel_name;
   }
 
-  auto output_dir =
-      std::filesystem::temp_directory_path() / "mlx" / version() / "cpu";
-  if (!std::filesystem::exists(output_dir)) {
-    std::filesystem::create_directories(output_dir);
-  }
+  static const std::filesystem::path output_dir = [] {
+    if (auto dir = std::getenv("MLX_CPU_CACHE_DIR"); dir && *dir) {
+      return std::filesystem::absolute(dir);
+    }
+    return std::filesystem::temp_directory_path() / "mlx" / version() / "cpu";
+  }();
+  std::filesystem::create_directories(output_dir);
 
 #ifdef _WIN32
   std::string shared_lib_name = kernel_file_name + ".dll";

From 9b5df7587a4857ded20192173aa4f595f3f1bbb5 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Mon, 19 Oct 2026 09:56:38 +0900
Subject: [PATCH 2/2] nit

---
 docs/src/usage/environment_variables.rst | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst
index 1aabecf00d..b43c053d52 100644
--- a/docs/src/usage/environment_variables.rst
+++ b/docs/src/usage/environment_variables.rst
@@ -26,15 +26,6 @@ General
    on supported hardware. The default is ``1``. Set it to ``0`` to keep these
    operations in full ``float32`` precision. See :doc:`precision`.
 
-CPU
----
-
-.. envvar:: MLX_CPU_CACHE_DIR
-
-   Override the directory used to cache runtime-compiled CPU kernels. An unset
-   or empty value uses ``mlx/<version>/cpu`` under the system temporary
-   directory, which is the default.
-
 Distributed
 -----------
 
@@ -208,3 +199,12 @@ The MLX-prefixed variables in this section are advanced CUDA backend controls.
 
 MLX also uses ``CUDA_HOME`` or ``CUDA_PATH`` to locate CUDA headers for runtime
 kernel compilation when they cannot be found in the Python environment.
+
+CPU
+---
+
+.. envvar:: MLX_CPU_CACHE_DIR
+
+   Override the directory used to cache runtime-compiled CPU kernels. An unset
+   or empty value uses ``mlx/<version>/cpu`` under the system temporary
+   directory, which is the default.