Skip to content

Fix missing add_temporary in pad_gpu/fill_gpu #4660

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zcbenz merged 6 commits into from
Oct 13, 2026
Merged

Fix missing add_temporary in pad_gpu/fill_gpu #4660

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7a7aedf
retain value in fill_gpu
Oct 8, 2026
aa775c0
Add unit test
Oct 8, 2026
b9ff671
Remove unit test (as suggested by Cheng)
Oct 12, 2026
5d66e3f
remove fill_gpu edit
Oct 12, 2026
af65e4d
Merge branch 'main' into fill-gpu
Oct 12, 2026
fe497bd
Add add_temporary to fill_gpu callers
ronaldmannak Oct 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions mlx/backend/metal/conv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -790,6 +790,7 @@ void pad_and_slice_conv_3D_gpu(
array x_copy(xshape, x.dtype(), nullptr, {});
array zero(0, x.dtype());
pad_gpu(x, zero, x_copy, {0, -1}, {0, 0}, s);
metal::get_command_encoder(s).add_temporary(zero);
metal::get_command_encoder(s).add_temporary(x_copy);

return x_copy;
Expand Down Expand Up @@ -1393,6 +1394,7 @@ void pad_in_channels_conv_2D_gpu(
array x_copy(xshape, x.dtype(), nullptr, {});
array zero(0, x.dtype());
pad_gpu(x, zero, x_copy, {-1}, {0}, s);
metal::get_command_encoder(s).add_temporary(zero);
metal::get_command_encoder(s).add_temporary(x_copy);

return x_copy;
Expand Down
26 changes: 16 additions & 10 deletions mlx/backend/metal/gated_delta_update.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,9 @@ void GatedDeltaUpdate::eval_gpu(
out.set_data(allocator::malloc(out.nbytes()));
hf.set_data(allocator::malloc(hf.nbytes()));

fill_gpu(array(0, out.dtype()), out, s);
array zero(0, out.dtype());
compute_encoder.add_temporary(zero);
fill_gpu(zero, out, s);

// The forward is inference-only here: the vjp runs its own save pass. The
// kernel declares buffers 9, 10 and 11 unconditionally, so bind a placeholder
Expand Down Expand Up @@ -289,12 +291,16 @@ void GatedDeltaUpdateVJP::eval_gpu(
auto dg_acc = stage_fp32 ? scratch_alloc(dg.shape(), float32, s) : dg;
auto db_acc = stage_fp32 ? scratch_alloc(db.shape(), float32, s) : db;

fill_gpu(array(0, float32), dq_acc, s);
fill_gpu(array(0, float32), dk_acc, s);
fill_gpu(array(0, float32), dv_acc, s);
fill_gpu(array(0, float32), dg_acc, s);
fill_gpu(array(0, float32), db_acc, s);
fill_gpu(array(0, dh.dtype()), dh, s);
array zero(0, float32);
array zero_h(0, dh.dtype());
compute_encoder.add_temporary(zero);
compute_encoder.add_temporary(zero_h);
fill_gpu(zero, dq_acc, s);
fill_gpu(zero, dk_acc, s);
fill_gpu(zero, dv_acc, s);
fill_gpu(zero, dg_acc, s);
fill_gpu(zero, db_acc, s);
fill_gpu(zero_h, dh, s);

array state_cache = scratch_alloc({B, Hv, n_states, Dv, Dk}, float32, s);
array seg_states = scratch_alloc(
Expand All @@ -305,9 +311,9 @@ void GatedDeltaUpdateVJP::eval_gpu(
chunked ? Shape({B, Hv, n_chunks, 16, Dv}) : Shape({1}), float32, s);

if (chunked) {
fill_gpu(array(0, float32), state_cache, s);
fill_gpu(array(0, float32), chunk_mats, s);
fill_gpu(array(0, float32), chunk_delta, s);
fill_gpu(zero, state_cache, s);
fill_gpu(zero, chunk_mats, s);
fill_gpu(zero, chunk_delta, s);
}

array y_scratch = scratch_alloc({B, T, Hv, Dv}, q.dtype(), s);
Expand Down
Loading