fix(rocm): key custom-kernel JIT modules by their generated source #2149

Part of #1801

Problem / Background

Found in #2100's security review. The vendored ROCm fast::hip_kernel caches compiled modules under a key that does not cover everything the generated source depends on, so a second call with different input dtypes can silently reuse a module compiled for the first call and read its buffers through the wrong pointer type. Nothing throws; the output is just wrong.

Current Behavior

In src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp:

• :265-266 names the kernel "custom_kernel_" + name + template_arguments_hash(template_args); template_arguments_hash (:108) hashes only the template args.
• build_kernel (:129-202) generates a source that also depends on: each input's dtype (dtype_to_hip_type(arr.dtype()), :150), each output's dtype (:171), and whether each input has ndim() > 0, which decides whether its _shape / _strides / _ndim parameters are declared at all (:153-164). The ndim value itself does not matter: since LOCAL_FIXES.md item 30, <input>_shape is a fixed-width KernelShape passed by value.
• CustomKernel::eval_gpu (:346-354) calls rocm::get_jit_module(s.device, name_, builder, false), and get_jit_module (jit_module.cpp:502-517) memoises by device.index + ":" + name and runs the builder only on a miss.

Upstream CUDA at the MLX pin 81ba1c6a already fixed this: mlx/backend/cuda/custom_kernel.cpp keys the module by fmt::format("{}_{:x}", name_, std::hash<std::string>{}(source_)) (around line 313), and Metal appends every input dtype to the name. The ROCm fork keeps the old name-only key.

mlxcel's own launchers are mostly shielded by scripts/ci/check_kernel_dtype_keys.py (make verify-kernel-dtype-keys), which requires at least one input dtype in each launch's template_args. It cannot cover a launch whose inputs have different dtypes that vary independently, an output dtype that differs at one input dtype, or a scalar versus non-scalar input, and its docstring (:15-19) still describes CUDA as name-only, which is stale at the pin.

Proposed Solution

Mirror upstream CUDA in the overlay: in CustomKernel::eval_gpu, pass module_name = name_ + "_" + hex(std::hash<std::string>{}(source_)) to rocm::get_jit_module while keeping kernel_name (the symbol looked up in the module) unchanged. Because the source is the complete input to compilation, this covers dtypes, the scalar case and any future source-affecting input at once. Add a LOCAL_FIXES.md entry citing the CUDA line it copies. Leave get_jit_module's device-index prefix as is.

Rejected: appending input and output dtypes to the name as Metal does. It fixes dtypes but not the scalar-input case, and diverges from the backend the overlay tracks.

Also correct the CUDA paragraph of check_kernel_dtype_keys.py's docstring (the check itself stays as defense in depth for a future backend or a re-vendored fork that regresses the key).

Implementation Notes

• Reuse: the rocm_fault_probe_array bridge fixture (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5163, wrapper mlxcel_core::rocm_faults) is the pattern for a ROCm-only test kernel reachable from Rust.
• Edge cases: the hash is computed per eval_gpu; source_ is built once per call in hip_kernel's closure, so the extra cost is one string hash per launch, the same as CUDA. Two sources that collide in std::hash would still share a module; CUDA accepts that and so does this fix.
• Constraints: no change for Metal or CUDA builds; no change to kernel symbol names, so rocprofv3 traces keep their custom_kernel_* names.

Acceptance Criteria

• A new ROCm bridge fixture rocm_jit_key_probe(input) -> Result<UniquePtr<MlxArray>> launches one hip_kernel with a fixed name and no template args that writes out[i] = (float)inp[i] into an f32 output.
• tests/rocm_custom_kernel_jit_key.rs calls it with an f32 array, then an f16 array, then a bf16 array of known values, and asserts each output equals the input. It fails on current main (second and third calls reuse the f32 module) and passes with the fix.
• The same test covers a 0-d input after a 1-d input with the probe reading inp_shape, run in a child process as tests/rocm_gpu_faults.rs does (on current main the argument lists disagree and can fault the queue), and passes with the fix.
• make verify-rocm passes with no other test change; LOCAL_FIXES.md and make verify-rocm-overlay agree.

Verification

cargo test --release --features rocm --test rocm_custom_kernel_jit_key -- --test-threads=1
make verify-rocm verify-kernel-dtype-keys