From ca83aade0465bf5a30248e22dc3a742d6eda55f4 Mon Sep 17 00:00:00 2001
From: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com>
Date: Sat, 11 Jul 2026 04:33:08 -0400
Subject: [PATCH] Fix custom metal kernel cache collision for same name,
 different source (#3832)

Two fast.metal_kernel instances sharing a name but with different sources,
dispatched before a single mx.eval, silently ran the first kernel's compiled
code for the second (wrong results, no error). The library cache was keyed by
name alone with a hand-rolled clear-and-rebuild of that one slot, which is not
safe within a single command batch.

Key the library by name + source + compile options so kernels that differ in
any of those get distinct libraries (and distinct pipelines) that coexist for
the whole batch, instead of thrashing one slot. compile_options_ is still
passed to get_library so per-kernel math mode is honored. The map key is never
used as a Metal identifier, so the full source is a collision-free key.

Adds a regression test (two same-name/different-source kernels in one eval).
---
 mlx/backend/metal/custom_kernel.cpp | 36 +++++++----------------------
 python/tests/test_fast.py           | 32 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp
index 0648ed221e..9b92f6ec3c 100644
--- a/mlx/backend/metal/custom_kernel.cpp
+++ b/mlx/backend/metal/custom_kernel.cpp
@@ -8,16 +8,6 @@
 
 namespace mlx::core::fast {
 
-struct CustomKernelCache {
-  std::unordered_map<std::string, std::pair<std::string, CompileOptions::Data>>
-      libraries;
-};
-
-static CustomKernelCache& cache() {
-  static CustomKernelCache cache_;
-  return cache_;
-};
-
 void CustomKernel::eval_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
@@ -55,25 +45,15 @@ void CustomKernel::eval_gpu(
 
   auto& d = metal::device(s.device);
 
-  {
-    // Clear kernels from the device library cache if needed
-    auto& kernel_cache = cache();
-    if (auto it = kernel_cache.libraries.find(name_);
-        it != kernel_cache.libraries.end()) {
-      if (it->second.first != source_ ||
-          it->second.second != compile_options_) {
-        auto& d = metal::device(s.device);
-        d.clear_library(name_);
-        it->second = std::make_tuple(source_, compile_options_);
-      }
-    } else {
-      kernel_cache.libraries.emplace(
-          name_, std::make_tuple(source_, compile_options_));
-    }
-  }
-
+  // Key the library by (name, source, compile options) so kernels that differ
+  // in any of those coexist instead of thrashing a single cache slot within one
+  // eval batch (fixes #3832). The map key is never used as a Metal identifier,
+  // so the full source is a collision-free key (a content hash could alias
+  // distinct sources and silently reintroduce the bug).
+  std::string lib_name =
+      name_ + "\n" + source_ + "\n" + std::to_string(compile_options_);
   auto lib = d.get_library(
-      name_, compile_options_, [this] { return metal::utils() + source_; });
+      lib_name, compile_options_, [this] { return metal::utils() + source_; });
   auto kernel = d.get_kernel(name_, lib);
   auto& compute_encoder = metal::get_command_encoder(s);
   compute_encoder.set_compute_pipeline_state(kernel);
diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py
index 8343c45941..5dacaa605c 100644
--- a/python/tests/test_fast.py
+++ b/python/tests/test_fast.py
@@ -1026,6 +1026,38 @@ def call_kernel(a: mx.array, source):
         out = call_kernel(a, source)
         self.assertTrue(mx.array_equal(out, mx.ones_like(out)))
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
+    def test_custom_kernel_same_name_different_source_one_eval(self):
+        # Regression test for #3832: two kernels sharing a name but with
+        # different sources, dispatched in a SINGLE eval batch, must each run
+        # their own compiled code instead of silently reusing the first's.
+        def call_kernel(a, source):
+            kernel = mx.fast.metal_kernel(
+                name="dup_name",
+                input_names=["inp"],
+                output_names=["out"],
+                source=source,
+            )
+            return kernel(
+                inputs=[a],
+                grid=(a.size, 1, 1),
+                threadgroup=(a.size, 1, 1),
+                output_shapes=[a.shape],
+                output_dtypes=[a.dtype],
+                stream=mx.gpu,
+            )[0]
+
+        a = mx.arange(32, dtype=mx.float32)
+        out_a = call_kernel(
+            a, "uint e = thread_position_in_grid.x; out[e] = inp[e] * 2.0f;"
+        )
+        out_b = call_kernel(
+            a, "uint e = thread_position_in_grid.x; out[e] = inp[e] + 100.0f;"
+        )
+        mx.eval(out_a, out_b)  # one batch — the reported failure case
+        self.assertTrue(mx.array_equal(out_a, a * 2.0))
+        self.assertTrue(mx.array_equal(out_b, a + 100.0))
+
     @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
     def test_custom_metal_kernel_math_mode(self):
         with self.assertRaises(ValueError):