Fix custom metal kernel cache collision for same name, different source (#3832)#3833
Open
katlun-lgtm wants to merge 1 commit into
Open
Fix custom metal kernel cache collision for same name, different source (#3832)#3833katlun-lgtm wants to merge 1 commit into
katlun-lgtm wants to merge 1 commit into
Conversation
…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).
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
1 participant
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_ + 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 library map key is never used as a Metal identifier (only as a cache key), so the full source string is a collision-free key. A content hash — as suggested in the issue — would have a (tiny) collision chance that could alias two distinct sources and silently reintroduce exactly this bug, so the full string is used instead.
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.