PUBLIC:

Release server_init_mu_ before executing SyncPull in KvCacheManagerBase to avoid holding it during the entire network transfer in SyncPull.

### Test

Before the fix, cross-node transfer takes about
With the fix, cross-node transfer takes about

PiperOrigin-RevId: 969568450

Author: benliangw

tpu_sync/kv_cache/kv_cache_manager_base.cc

@@ -1195,15 +1195,24 @@ absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::H2hReadExplicit(
     const std::vector<uint8_t*>& explicit_dst_ptrs, int parallelism,
     tpu_raiden::transport::MajorOrder major_order,
     tpu_raiden::transport::BlockReceivedCallback on_block_received) {
-  absl::MutexLock lock(server_init_mu_);
-  if (!server_) {
+  // Take the transport pointer under the lock but run the pull outside it
+  // (same pattern as StartPushInternal): SyncPull blocks for the whole
+  // network transfer, and holding server_init_mu_ across it serializes every
+  // concurrent pull in the process -- including the pull_pool_ pipeline in
+  // H2dRead, which degenerates to one round trip at a time.
+  tpu_raiden::transport::BlockTransport* transport_server = nullptr;
+  {
+    absl::MutexLock lock(server_init_mu_);
+    transport_server = server_.get();
+  }
+  if (!transport_server) {
     return absl::FailedPreconditionError("Transport server is not running");
   }
   ASSIGN_OR_RETURN(
       std::vector<int> allocated_ids,
-      server_->SyncPull({peer}, src_block_ids, local_block_ids,
-                        explicit_dst_ptrs, parallelism, major_order,
-                        on_block_received, kLeaseAuthorizedPullUuid));
+      transport_server->SyncPull({peer}, src_block_ids, local_block_ids,
+                                 explicit_dst_ptrs, parallelism, major_order,
+                                 on_block_received, kLeaseAuthorizedPullUuid));
   return raiden::PjRtCopyFuture(std::vector<raiden::BufferHolder>{});
 }