Skip to content

[CUDA] Fix custom kernel cache collision for same name, different source - #4273

Open
HaoXuAI wants to merge 1 commit into
ml-explore:mainfrom
HaoXuAI:cuda-kernel-cache-fix
Open

[CUDA] Fix custom kernel cache collision for same name, different source#4273
HaoXuAI wants to merge 1 commit into
ml-explore:mainfrom
HaoXuAI:cuda-kernel-cache-fix

Conversation

@HaoXuAI

@HaoXuAI HaoXuAI commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Fixes #4275.

get_jit_module keys the module cache on the kernel name alone (jit_module.cpp:471-485), so a second custom kernel sharing a name silently runs the first one's compiled code — no error, wrong results.

This is the CUDA counterpart of #3832, fixed for Metal in #3833. It is worse here: on Metal the stale-source invalidation worked across eval boundaries and only failed within a single batch, whereas get_jit_module has no invalidation at all, so the collision persists for the lifetime of the process.

The fix hashes the source into the module name, the same way #3833 did. The kernel name passed to get_kernel() is unchanged, so lookup inside the module still works.

It also covers a second path with the same root cause: read_cached_ptx() runs unconditionally (jit_module.cpp:407) even when use_disk_cache is false, so a stale on-disk PTX under the same name could be picked up too.

Testing

test_cuda_kernel_same_name_different_source mirrors the Metal regression test from #3833. Verified on an L40S (sm_89) with CUDA 12.6 — two kernels sharing an entry name but differing in body, with a separate mx.eval between them:

before after
first source (x * 2) [0, 2, 4] correct correct
second source (x + 100) [0, 2, 4] — kernel A's code [100, 101, 102] correct

The rest of test_fast.py is unaffected: 24 passed / 4 skipped on CUDA, 27 passed / 1 skipped on Metal.

I ran into this through mx.fast.precompiled_cuda_kernel with Triton output, where it is easy to hit by accident since Triton names kernels after the Python function.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] [CUDA] fast.cuda_kernel: same-name kernels with different sources silently run the first kernel's code

1 participant