Describe the bug
Two mx.fast.metal_kernel instances sharing name but with different source, both dispatched before a single mx.eval, silently execute the first kernel's compiled code for the second call. No error or warning. Return wrong results.
To Reproduce (affected: MLX 0.31.1, 0.31.2, 0.32.0 from PyPI, and git main 0.32.1.dev20260710+4367c73b built from source; M4 Pro, macOS 26):
import mlx.core as mx
n = 64
a = mx.arange(n, dtype=mx.float32)
mk = lambda src: mx.fast.metal_kernel(
name="probe", input_names=["a"], output_names=["out"], source=src)
body = "uint i = thread_position_in_grid.x; if (i < %du) out[i] = a[i] %s;"
kA = mk(body % (n, "* 2.0f"))
kB = mk(body % (n, "+ 100.0f"))
run = lambda k, x: k(inputs=[x], output_shapes=[(n,)], output_dtypes=[mx.float32],
grid=(n, 1, 1), threadgroup=(min(n, 256), 1, 1))[0]
yA, yB = run(kA, a), run(kB, a)
mx.eval(yA, yB) # ONE eval batch
print(yB[:3]) # observed: [0, 2, 4] (kernel A's code!)
# expected: [100, 101, 102]
# With mx.eval(yA) before building yB, the result is correct.
Expected behavior
Correct dispatch (cache keyed by name + source identity), or a loud error on name collision. Shouldn't perform silent execution of a different kernel's code.
Desktop (please complete the following information):
- OS Version: macOS 26 (Darwin 25.5.0)
- M4 Pro, MLX 0.31.1 (PyPI), Python 3.13
Additional context
The custom-kernel library cache is keyed by kernel name (Device::get_library(name_, …) in backend/metal/custom_kernel.cpp); the stale-source invalidation works across eval boundaries but not within one batch. Frameworks that execute a whole program as one batch (e.g. PJRT plugins lowering to fast::metal_kernel) hit this deterministically for any name reuse.
Potential workaround: suffix the name with a content hash of the source.
Describe the bug
Two
mx.fast.metal_kernelinstances sharingnamebut with differentsource, both dispatched before a singlemx.eval, silently execute the first kernel's compiled code for the second call. No error or warning. Return wrong results.To Reproduce (affected: MLX 0.31.1, 0.31.2, 0.32.0 from PyPI, and git main
0.32.1.dev20260710+4367c73bbuilt from source; M4 Pro, macOS 26):Expected behavior
Correct dispatch (cache keyed by name + source identity), or a loud error on name collision. Shouldn't perform silent execution of a different kernel's code.
Desktop (please complete the following information):
Additional context
The custom-kernel library cache is keyed by kernel name (
Device::get_library(name_, …)inbackend/metal/custom_kernel.cpp); the stale-source invalidation works across eval boundaries but not within one batch. Frameworks that execute a whole program as one batch (e.g. PJRT plugins lowering tofast::metal_kernel) hit this deterministically for any name reuse.Potential workaround: suffix the name with a content hash of the source.