Repository navigation
Fix custom metal kernel cache collision for same name, different source - #3833
Merged
Merged
Conversation
3 of 4 tasks
zcbenz
reviewed
Aug 2, 2026
katlun-lgtm
added a commit
to katlun-lgtm/mlx
that referenced
this pull request
Aug 2, 2026
Keeps the library_map_ key short instead of embedding the full kernel source, per review feedback on ml-explore#3833.
zcbenz
approved these changes
Aug 3, 2026
zcbenz
left a comment
Member
There was a problem hiding this comment.
Thanks for the PR, the original cache design looked odd to me but I was worried about breaking things, let's just make the change and it should be safe.
…ce (ml-explore#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).
Keeps the library_map_ key short instead of embedding the full kernel source, per review feedback on ml-explore#3833.
zcbenz
force-pushed
the
fix/metal-kernel-cache-3832
branch
from
August 3, 2026 23:38
6e757ce to
75d994c
Compare
haoxu0
added a commit
to haoxu0/mlx
that referenced
this pull request
Aug 15, 2026
get_jit_module keys the module cache on the kernel name alone, so a second kernel sharing a name silently runs the first one's compiled code. Include a hash of the source in the module name, as ml-explore#3833 did for Metal.
This was referenced Aug 15, 2026
1 task done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #3832. Two
fast.metal_kernelinstances that share anamebut have differentsource, dispatched before a singlemx.eval, silently execute the first kernel's compiled code for the second — wrong results, no error.Cause
The custom-kernel library cache is keyed by
name_alone (backend/metal/custom_kernel.cpp), with a hand-rolled "if the source changed,clear_library(name_)and rebuild" invalidation. That is safe acrossevalboundaries (the command buffer has completed before the clear) but not within one batch: the second kernel clears and rebuilds the sharedname_slot while the first kernel's dispatch is still pending in the same command buffer, and the second ends up bound to the first's code.Fix
Key the library by
name_ + hash(source_) + compile_options_, so kernels that differ in any of those map to distinct libraries — and therefore distinct pipelines (the pipeline cache is already keyed by(library, function)). They now coexist for the whole batch instead of thrashing one slot, and the fragile clear-and-rebuild (with itsCustomKernelCachebookkeeping) is removed.compile_options_is still passed toget_library, so a kernel's math mode is honored.The source is hashed with
std::hash<std::string>rather than embedded whole, so the key stays short — the same treatmentbackend/cpu/compiled.cppgives over-long kernel names. A wrong cache hit would need two kernels to share a name, share compile options, and collide on the 64-bit hash.Test
Adds
test_custom_kernel_same_name_different_source_one_eval: two same-name / different-source kernels evaluated in a single batch must each return their own result. It fails before this change and passes after.Validation
Built against
main;python/tests/test_fast.pypasses (27/27, including the new test and the compile-options / caching tests), and the issue's repro now returns the correct[100, 101, 102]instead of[0, 2, 4].Thanks to @Jinge-Wang for the precise report and minimal repro.