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

Describe the bug

Two mx.fast.cuda_kernel instances sharing name but with different source silently execute the first kernel's compiled code for the second call. No error or warning, wrong results.

This is the CUDA counterpart of #3832, which was 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 on CUDA the collision persists for the lifetime of the process.

mx.fast.precompiled_cuda_kernel is affected the same way, and it is easier to hit there by accident because Triton names its kernels after the Python function.

To Reproduce (mlx 0.32.1.dev20260809, built from source; NVIDIA L40S, CUDA 12.6, driver 580.95.05)

import mlx.core as mx

n = 64
a = mx.arange(n, dtype=mx.float32)
elem = "auto i = cooperative_groups::this_grid().thread_rank();"

def run(src):
    k = mx.fast.cuda_kernel(
        name="probe", input_names=["a"], output_names=["out"], source=src
    )
    return k(inputs=[a], output_shapes=[(n,)], output_dtypes=[mx.float32],
             grid=(n, 1, 1), threadgroup=(min(n, 256), 1, 1))[0]

yA = run(f"{elem} out[i] = a[i] * 2.0f;")
mx.eval(yA)                                  # separate eval, unlike #3832
yB = run(f"{elem} out[i] = a[i] + 100.0f;")
mx.eval(yB)
print(yA[:3].tolist())   # [0.0, 2.0, 4.0]   correct
print(yB[:3].tolist())   # [0.0, 2.0, 4.0]   kernel A's code — expected [100, 101, 102]

Expected behavior

The second kernel runs its own source, or a loud error on name collision. Not silent execution of different code.

Desktop

• Linux, NVIDIA L40S (sm_89), CUDA 12.6, driver 580.95.05
• mlx 0.32.1.dev20260809 built from source, Python 3.11

Additional context

get_jit_module keys its cache on the kernel name alone:

// mlx/backend/cuda/jit_module.cpp:471-485
static auto* cache = new std::unordered_map<std::string, JitModule>;
if (auto it = cache->find(name); it != cache->end()) {
  return it->second;
}

and CustomKernel::eval_gpu passes name_ as that key, so the source never participates.

There is a second path with the same root cause: read_cached_ptx(ptx_cache_dir(), module_name, ...) runs unconditionally at jit_module.cpp:407 even when use_disk_cache is false, so a stale on-disk PTX stored under the same name can also be picked up.

Suggested fix, mirroring #3833: hash the source into the module cache key. PR #4273.