From b5f33b0e3e81970e4b2f18d3f801f3881f3b7a6f Mon Sep 17 00:00:00 2001 From: Googler Date: Wed, 29 Jul 2026 11:00:59 -0700 Subject: [PATCH] Fix xla::Future mapped callback completion drop-off across pip-wheel .so boundary Capture copy_future by value in the .Map(...) closure in weight_synchronizer_base.cc (D2h) to preserve AsyncValue reference counting across dynamic library (.so) boundaries without dropping completion signals or hanging execution. PiperOrigin-RevId: 955977043 --- tpu_raiden/weight_sync/weight_synchronizer_base.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tpu_raiden/weight_sync/weight_synchronizer_base.cc b/tpu_raiden/weight_sync/weight_synchronizer_base.cc index c9dde6fd..d17a751b 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_base.cc +++ b/tpu_raiden/weight_sync/weight_synchronizer_base.cc @@ -325,7 +325,7 @@ absl::StatusOr WeightSynchronizerBase::D2h() { shard_hold.CopyRawDeviceToHost(temp_buffer_ptr, 0, layer_size); xla::Future<> detile_future = - copy_future.Map([temp_buffer, dst_host_ptr, + copy_future.Map([copy_future, temp_buffer, dst_host_ptr, shape = shard_hold.buffer->on_device_shape(), layout = *xla_layout]() -> absl::Status { return tpu_raiden::weight_sync::DetileBuffer(