Skip to content

[BUG] fast.metal_kernel: same-name kernels with different sources silently run stale code within one eval batch #3832

Description

@Jinge-Wang

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions