From 6ee5fd9136d2e3b20a19df9c79584522d3a5b9ca Mon Sep 17 00:00:00 2001 From: HaoXuAI Date: Sat, 15 Aug 2026 01:21:11 -0700 Subject: [PATCH 1/2] [CUDA] Fix custom kernel cache collision for same name, different source 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 #3833 did for Metal. --- mlx/backend/cuda/custom_kernel.cpp | 7 ++++++- python/tests/test_fast.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cuda/custom_kernel.cpp b/mlx/backend/cuda/custom_kernel.cpp index 9b5bd38b7f..383f5cb331 100644 --- a/mlx/backend/cuda/custom_kernel.cpp +++ b/mlx/backend/cuda/custom_kernel.cpp @@ -310,9 +310,14 @@ void CustomKernel::eval_gpu( // Compile the custom kernel std::string kernel_name = (is_precompiled_) ? name_ : "mlx::core::cu::" + name_; + // The module cache is keyed on this name, so it has to include the source: + // two kernels sharing a name but not a body would otherwise both run + // whichever was compiled first. Same fix as #3833 on the Metal side. + std::string module_name = + fmt::format("{}_{:x}", name_, std::hash{}(source_)); cu::JitModule& mod = cu::get_jit_module( encoder.device(), - name_, + module_name, [&]() { return std::make_tuple( is_precompiled_, source_, std::vector{kernel_name}); diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 5dacaa605c..ba5b8f3138 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1058,6 +1058,35 @@ def call_kernel(a, source): self.assertTrue(mx.array_equal(out_a, a * 2.0)) self.assertTrue(mx.array_equal(out_b, a + 100.0)) + @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") + def test_cuda_kernel_same_name_different_source(self): + # The CUDA module cache was keyed on the kernel name alone, so the + # second kernel here silently ran the first one's code. Metal had the + # same bug, fixed in #3833. + def call_kernel(a, source): + kernel = mx.fast.cuda_kernel( + name="dup_name", + input_names=["inp"], + output_names=["out"], + source=source, + ) + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.arange(32, dtype=mx.float32) + elem = "auto e = cooperative_groups::this_grid().thread_rank();" + out_a = call_kernel(a, f"{elem} out[e] = inp[e] * 2.0f;") + out_b = call_kernel(a, f"{elem} out[e] = inp[e] + 100.0f;") + mx.eval(out_a, out_b) + self.assertTrue(mx.array_equal(out_a, a * 2.0)) + self.assertTrue(mx.array_equal(out_b, a + 100.0)) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") def test_custom_metal_kernel_math_mode(self): with self.assertRaises(ValueError): From 8d351fbe121be27c860e5e3e73e74a89783c699a Mon Sep 17 00:00:00 2001 From: Cheng Date: Sun, 16 Aug 2026 19:34:45 +0900 Subject: [PATCH 2/2] nit --- mlx/backend/cuda/custom_kernel.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/mlx/backend/cuda/custom_kernel.cpp b/mlx/backend/cuda/custom_kernel.cpp index 383f5cb331..c230656d80 100644 --- a/mlx/backend/cuda/custom_kernel.cpp +++ b/mlx/backend/cuda/custom_kernel.cpp @@ -310,9 +310,6 @@ void CustomKernel::eval_gpu( // Compile the custom kernel std::string kernel_name = (is_precompiled_) ? name_ : "mlx::core::cu::" + name_; - // The module cache is keyed on this name, so it has to include the source: - // two kernels sharing a name but not a body would otherwise both run - // whichever was compiled first. Same fix as #3833 on the Metal side. std::string module_name = fmt::format("{}_{:x}", name_, std::hash{}(source_)); cu::JitModule& mod = cu::get_jit_module(