Metal fill_gpu fails to keep fill value

[ronaldmannak]

The Metal fill_gpu fails to keep the fill value alive, and the GPU can read overwritten values, depending on timing.

Callers of fill_gpu currently use different patterns. There are 12 call sites in conv.cpp and gated_delta_update.cpp that do not keep the fill value alive and potentially expose the bug (confirmed with 3D Conv). There are 8 calls from attention, matmul and more that correctly keep the value alive using add_temporary, and there's one call from Pad::eval_gpu where the eval loop retains the value

This PR adds a two line fix to fill_gpu in metal/copy.cpp and a unit test. An alternative solution would be to add add_temporary(zero) to the 12 call sites that fail to keep the fill value alive, but that seems like a brittle fix. A second alternative would be adding add_temporary(val) to fill_gpu, but that may break Pad since its fill value is a graph input and not temporary.

To replicate the issue:

import mlx.core as mx

# A conv3d over all-ones data. Every output value is the sum of one
# 2 x 2 x 2 window over 3 channels: 2 * 2 * 2 * 3 = 24.
x = mx.ones((1, 4, 8, 8, 3))
w = mx.ones((4, 2, 2, 2, 3))
y = mx.conv3d(x, w)

# The Metal conv pads the 3 channels of x and w to 16 with a temporary zero
# scalar. async_eval queues that work and returns before the GPU runs it.
mx.async_eval(y)

# Before the fix, the zero scalar is already freed here. This new scalar gets
# the same buffer and writes 1.0 to it, so the GPU pads with 1.0 instead of 0.
one = mx.array(1.0)

# Expected:       24.0
# Before the fix: 128.0 (all 16 channels count: 2 * 2 * 2 * 16)
print(y.max().item())

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: I ran into this bug while working on MLX-Swift-LM and Opus 5.5 traced it back to mlx.

[zcbenz]

This is usually caused by some code forgetting to call add_temporary(val) before using fill_gpu. There is no need to change fill_gpu.

[ronaldmannak]

@zcbenz Sounds good, I can remove the code in fill_gpu and focus on the callers instead. I found the issue in pad_and_slice_conv_3D_gpu in conv.cpp and can add add_temporary there. If you want, I can also apply the same update to pad_in_channels_conv_2D_gpu and the 10 call sites in gated_delta_update.cpp so the calls conform your suggestion, though I haven't tested the issue there

[zcbenz]

Yeah please update all the callers to use add_temporary. For a more decent solution I think we should have a CommandEncoder.memset shortcut that fills array with scalar since it is a common and error-prone op, and CUDA has a builtin op for it, but we can do it in future.

Regarding the test, I think it is likely going to be flaky, I'm good with no test since the change would be trivial to verify.