fix(rocm): make scatter apply every update, check kernel arg sizes

src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md

@@ -22,6 +22,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code
 11. **`arg_reduce` launch geometry.** For short reduction axes `ArgReduce::eval_gpu` launched one 1024-thread block per output over a 1-D grid. AMD limits each grid dimension to 2^32 - 1 threads, so the launch failed with "invalid configuration argument" past about 4.2M outputs; mxfp4 `quantize` of a 4096x4096 matrix (an argmin over 16.7M elements) hit it. Axes of at most 64 elements now use a 64-thread block, and the grid spills into y. GPU `quantize` in mxfp4 is then bit-identical to the CPU at 4096x4096. Applies to the fork; to be proposed there.
 12. **f16 in the fast `gather_qmm` path.** `GatherQMM::eval_gpu` sent only bf16 activations to `gather_qmv_warp_shared_kernel`; f16 fell through to the generic `gather_qmv_kernel`, which took about 205 ms per call on `gfx1151`. f16 checkpoints such as `mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit` decoded at about 0.1 tok/s, with 99.7% of GPU time in that kernel (rocprofv3). f16 now uses the same warp-shared kernel (the kernel was already templated on the element type), which takes 1.44 ms for a decode-shaped call at Mixtral's expert shape and matches a dequantized f32 reference. Applies to the fork; to be proposed there.
 13. **f32 activations in the tiled qmv path.** `QuantizedMatmul::eval_gpu` entered the tiled qmv branch for any activation dtype, launched `qmv_tiled_kernel` only for bf16 and f16, and then returned, so f32 activations produced no kernel launch and the output was whatever the fresh allocation held (relative error from 1.0 up to 1e35 at every tested shape). The branch is now entered only for bf16/f16 and group sizes 32/64/128, so f32 takes the generic qmv dispatch, which matches a dequantized f32 reference (relative error about 1e-7). mlxcel-core's split/fused projection tests use f32 inputs and exposed it. Applies to the fork; to be proposed there.
+14. **Scatter wrote every update to the first index.** `Scatter::eval_gpu` passed `upd_post_idx_size` as `int32_t` while `scatter_general_kernel` declares it `int64_t`. `CommandEncoder::add_kernel_node` hands HIP the address of each argument, so the kernel read eight bytes from a four-byte value, got a huge divisor, and computed index element 0 for every update: `z[idx] = v`, `.at[idx].add(...)` and everything built on scatter (`eye`, `identity`, `tri`, `diag`) applied only one update. The size is now `int64_t`, and `add_kernel_node` statically asserts that every argument has the size of its kernel parameter, so this class of mismatch fails to compile (checked by reintroducing the `int32_t`). Applies to the fork; to be proposed there.
 
 ## Dropped from the vendored set
 

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device.h

@@ -45,6 +45,25 @@ extern std::atomic<long> g_inline_launches_;
 void set_current_prim(const char* name);
 void record_inline_launch();
 
+// mlxcelverse: add_kernel_node hands HIP the address of every argument, so
+// each argument must have the same size as the kernel parameter it lands in.
+// A narrower argument is read with garbage high bytes (Scatter once passed an
+// int32_t for an int64_t parameter and every update collapsed onto the first
+// index), so the sizes are checked at compile time.
+template <typename Func>
+struct kernel_param_types;
+template <typename... Args>
+struct kernel_param_types<void(Args...)> {
+  using type = std::tuple<Args...>;
+};
+template <typename KernelArgs, typename Passed, std::size_t... I>
+constexpr bool kernel_arg_sizes_match(std::index_sequence<I...>) {
+  return (
+      (sizeof(std::tuple_element_t<I, KernelArgs>) ==
+       sizeof(std::decay_t<std::tuple_element_t<I, Passed>>)) &&
+      ...);
+}
+
 class CommandEncoder {
  public:
   explicit CommandEncoder(Device& d);
@@ -77,6 +96,15 @@ class CommandEncoder {
       uint32_t smem_bytes,
       Params&&... params) {
     constexpr size_t num = sizeof...(Params);
+    using KernelArgs = typename kernel_param_types<Func>::type;
+    static_assert(
+        std::tuple_size_v<KernelArgs> == num,
+        "add_kernel_node: argument count differs from the kernel's parameters");
+    static_assert(
+        kernel_arg_sizes_match<KernelArgs, std::tuple<Params...>>(
+            std::index_sequence_for<Params...>{}),
+        "add_kernel_node: an argument's size differs from its kernel "
+        "parameter; pass the parameter's exact type");
     if (!use_hip_graphs()) {
       // Immediate launch: kernelParams are consumed synchronously, so
       // addresses of the caller's locals are fine.

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/indexing.hip

@@ -793,11 +793,15 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   int nidx = axes_.size();
   int32_t idx_ndim = nidx > 0 ? inputs[1].ndim() : 0;
 
-  int32_t upd_post_idx_size = std::accumulate(
+  // mlxcelverse: int64_t to match scatter_general_kernel's parameter.
+  // add_kernel_node passes argument addresses straight to HIP, so an int32_t
+  // here was read as an 8-byte value with garbage high bytes; every index
+  // element then collapsed onto the first index (only one update landed).
+  int64_t upd_post_idx_size = std::accumulate(
       upd.shape().begin() + idx_ndim,
       upd.shape().end(),
-      1,
-      std::multiplies<int32_t>());
+      int64_t{1},
+      std::multiplies<int64_t>());
 
   // Prepare host data for parameters
   std::vector<int32_t> h_upd_shape(upd.shape().begin(), upd.shape().end());