torch_tpu: Who's holding the buffer?
How base[0:50] gets a second TPU buffer that only a C++ local owns, and three async copies that forgot to keep theirs alive.
How it started
I've been curious about TPUs for a while, so I've been poking around torch_tpu, the PyTorch backend for them. When I pulled, the top commit was one making the transfer utilities in csrc/pjrt/pjrt_transfer.cc always non-blocking, so that's where I started reading.
The raw DMA copy functions in pjrt_transfer.h return this, and their only caller, WindowTpu, hands it straight to a WorkTpu:
struct AsyncDmaResult {
xla::Future<> future;
std::shared_ptr<void> buffer_hold;
};
// window_tpu.h
explicit WorkTpu(xla::Future<> future, std::shared_ptr<void> hold = nullptr)
So an async copy gives you a future, plus a type-erased object whose entire job description is "exist until the future is done". WorkTpu parks it in a hold_ member that defaults to nullptr, so if a caller forgets it, nothing complains.
The hold is a RawBufferHolder, and there are two structs by that name in the repo. The one in pjrt_transfer.cc keeps the TPU buffer alive itself. The one in batch_transfer.cc leaves that to a separate vector of shared_ptr<void>.
The duplication was meh. What actually bugged me was that torch_tpu's own rule for non-blocking user copies is the same as CUDA's: keep the source tensor alive until a later torch.Event says the copy is done. The tensor is enough there. So why does the internal copy need a babysitter??
CUDA muscle memory vs PJRT (a primer)
Quick primer, because my CUDA brain kept giving me the wrong answers here.
On CUDA, device memory is a flat void*. base = torch.randn(100, device="cuda") is one allocation, s = base[0:50] is just base_ptr + 0 with a length of 50, and cudaMemcpyAsync gets a pointer and a size. One buffer, and holding s keeps it alive.
On TPU, the void* inside a tensor's DataPtr isn't a device address at all. It points at a C++ handle, and the actual memory is a few hops further down:
at::Tensor → StorageImpl → DataPtr (void*) → DeviceBufferRef*
→ shared_ptr<DeviceBufferList> → unique_ptr<xla::PjRtBuffer>
A PjRtBuffer is opaque and whole. ToLiteral (device-to-host) and CopyToRemoteDevice (cross-host sends) take the entire buffer, and even the raw DMA path works in physical, tiled byte offsets rather than element offsets. So before you can copy a view, the view has to become its own buffer.
The comments around these holds all say the same thing. Batch transfer has // Keep the DeviceBufferRef alive until the transfer completes., and the holder promises that the DeviceBufferRef and its PjRtBuffer "remain alive and referenced until the DMA completes". So some buffer can apparently lose its last owner mid-copy while s is still alive. I was ready to write that off as overly cautious, since s shares storage with base and that's the end of it. The part I was missing is in GetBuffer (tensor_to_buffer.cc:349-367):
TT_ASSIGN_OR_RETURN(DeviceBufferRef base_buffer_ref, GetBaseBuffer(tensor));
TT_ASSIGN_OR_RETURN(const TpuLayout tpu_layout, ResolveTpuLayout(tensor));
if (!tensor.is_conj() && tensor.is_contiguous() &&
tensor.storage_offset() == 0 &&
tpu_layout.sizes == base_buffer_ref.dimensions() &&
tpu_layout.element_type == base_buffer_ref.element_type()) {
// ...
return base_buffer_ref; // Buffer #1
}
return CreateViewDeviceBufferRef(base_buffer_ref, tpu_layout,
tensor.is_conj(), donate_base); // Buffer #2
base[0:50] is contiguous with a storage offset of 0, but its sizes are (50) and the base is (100), so it takes the slow path. So does a.view(10, 10) of a 100-element tensor. The slow path builds a short-lived DeferredOp (kAsStrided) that runs the view on the base and produces a brand-new DeviceBufferList. That's Buffer #2. MaterializeAndReturn hands it to whoever asked, nothing caches it, and the caller's local variable is its only owner.
It can't go into s.storage() either. Storage always holds the base, and s shares that storage with base and every other view, so swapping in a 50-element buffer would break all of them.
CUDA TPU (torch_tpu + PJRT)
base ──► [ 100 floats ] base ─┐
s ──► base_ptr + 0, len 50 s ──┴► storage ─► Buffer #1 (100 elems)
one allocation, s keeps it alive MaterializeAndReturn(s) ─► Buffer #2 (50 elems)
owned only by a C++ stack local
So the comments weren't being paranoid. I was just thinking about the wrong buffer.
There's a second route that doesn't need views at all. In-place ops in torch_tpu "mutate" a tensor by swapping which DeviceBufferRef its storage points at (OverwriteDataPtr in tensor_to_buffer.cc). Start an async copy from t, call t.zero_() before waiting, and t and its storage are still perfectly alive, but the buffer the copy was reading from has been swapped out from under it.
Why not just capture the at::Tensor in the PJRT completion callback? torch_tpu did exactly that once, and a commit titled "Remove backing_tensor from TpuMallocAndMemcpyHtoD" took it back out:
It is unsafe to free an at::Tensor from a callback on a PjRt transfer. This callback is invoked on a background thread, and then tries to acquire the GIL to free the tensor. If that tensor is blocking on a subsequent transfer, this creates an A-before-B/B-before-A deadlock.
(That's also where the torch.Event rule for user copies comes from. For internal copies, the same message says "a backing vector should be used instead".)
So the only thing left to pin is the C++ DeviceBufferRef itself, which is exactly what those comments were asking for. That's the hold. On main that pin travels next to the future instead of inside it, so every caller has to carry it by hand.
Three ways to forget a hold
Once I knew what the hold was for, I went looking for callers that dropped it. I found these by reading, not by crashing (PJRT may well be forgiving about freeing buffers with work in flight, and I can't tell from torch_tpu), but the holder comment is pretty clear about what's required. The last one doesn't need an unforgiving runtime at all. It's a plain dangling pointer, and my recently rust-ified brain couldn't let that go.
StartD2H remembered the destination
StartD2H (pjrt_transfer.cc:212-256) copies a TPU buffer into host memory with ToLiteral, then sets a promise when it's done:
xla::Future<> future = (*buffer_or)->ToLiteral(literal.get());
// ...
// `literal` is kept alive until the transfer completes.
future.OnReady([literal = std::move(literal), promise = std::move(promise),
fp4_receiver = std::move(fp4_receiver),
dst_ptr](absl::Status s) mutable {
// ... pack FP4 if needed, then promise.Set(...)
});
literal is a MutableBorrowingLiteral over the caller's host memory, i.e. the destination. The comment is accurate as far as it goes. The source buffer is right there on the first line, though, and it isn't in the capture list.
The caller that makes this matter is CopyTpuToCpuWithTpuConversion, which handles non-blocking TPU-to-CPU copies that can't be a straight storage copy. Our s = base[0:50] is one of those: its storage holds 100 elements and the destination wants 50, so s.to("cpu", non_blocking=True) ends up there. It calls GetBuffer on the view, gets Buffer #2 into a local converted_src_buf, starts the copy, hands the future to RecordAsyncDeviceToHost and returns. That local was Buffer #2's only owner. (Blocking copies wait before returning, so they're fine.)
The fix captures device_buffer_list = buffer_ref.device_buffer_list() in the lambda and resets it right before promise.Set. Small diff (+9/-4), no test.
A vector that gets moved but never filled
batch_transfer.cc has two siblings, BatchTransferD2H and BatchTransferH2D. The D2H side keeps its hold, right under that "Keep the DeviceBufferRef alive" comment from earlier, with batch_holds.push_back(std::make_shared<DeviceBufferRef>(std::move(src_buffer_ref))).
Here's the H2D side (batch_transfer.cc:504-587):
std::vector<std::shared_ptr<RawBufferHolder>> batch_c_api_holds;
std::vector<std::shared_ptr<void>> batch_holds;
// ...
for (size_t layer_idx = 0; layer_idx < n; ++layer_idx) {
// ...
TT_ASSIGN_OR_RETURN(DeviceBufferRef dst_buffer_ref,
GetPjrtBuffer(dst, dst_buffer));
// ... raw alias + CopyRawHostToDevice
}
acc.AndThen(std::move(batch_futures), std::move(batch_c_api_holds),
std::move(batch_holds));
batch_holds is declared, carefully std::moved into AndThen, and never has a single thing put in it. It might be the best-handled empty vector in the repo.
For a plain base tensor destination it mostly doesn't bite, since TransferFuture also keeps every src and dst at::Tensor and the storage still holds Buffer #1. zero_() the destination before wait() though, and the buffer the DMA is writing into gets swapped out, with that empty vector as its only protection.
View destinations are worse. They just lose the data. GetPjrtBuffer calls MaterializeAndReturn(dst), so a view gets Buffer #2, the DMA writes the host bytes into Buffer #2, and nothing ever copies them back into dst.storage(). No error, no data, and Buffer #2 is gone at the end of the loop body anyway.
The fix adds the missing push_back, which hands the ref to TransferFuture until the caller's Await(), plus a ValidateContiguousBase check that rejects views with "destination TPU tensor must be a contiguous base tensor (non-base views are not supported by BatchTransferH2D)". While I was in there I also stopped it from materializing the first tensor twice. There are two new tests: one for the overwrite-before-wait path (which doesn't actually check the data yet, oops), and one that checks a view gets that error.
.contiguous(), the no-op
Send a boring 32 MB contiguous tensor to a peer with WindowTpu::PutPeerP2P. Nothing weird about it.
PutPeerP2P stripes large sends: total size divided by an 8 MB chunk, clamped to 1-4 stripes, so 16 MB gets 2 and 32 MB gets 4. The chunks come from tensor.contiguous().view(-1).chunk(num_stripes, 0), and each one goes to backend_->send as {chunks[i].contiguous()}.
In my head, .contiguous() always copied. It does not, and I'm a little embarrassed about how long I believed that. A 1D chunk with stride 1 is already contiguous, so .contiguous() hands back the exact same tensor. Each stripe is still a view of the flattened tensor, its sizes don't match the base, and it becomes Buffer #2. (The .contiguous() at the top of PutPeerP2P does the same nothing for a user's own contiguous view, so even the single-stripe path can get here.)
backend_->send goes through ProcessGroupTpu::send to experimental_send (process_group_tpu.cc:1910-1942):
for (size_t i = 0; i < tensors.size(); ++i) {
// ...
TT_ASSIGN_OR_THROW(
const DeviceBufferRef device_buffer,
MaterializeAndReturn(tensor, MaterializationReason::kDistributedOp));
TT_ASSIGN_OR_THROW(xla::PjRtBuffer * pjrt_buffer,
device_buffer.AwaitBuffer());
pjrt_buffers.push_back(pjrt_buffer); // raw pointer
// ...
}
TT_ASSIGN_OR_THROW(auto send_future, // ... six lines after the '}'
CrossHostSendBuffers(pjrt_buffers, transfer_keys));
The loop's closing brace is the last thing that owned Buffer #2. And CrossHostSendBuffers doesn't even use the pointer right away. It first does a blocking store_->multiGet, which (per the comment above it) "will block until descriptors from the receiver are available in the store", and only then calls CopyToRemoteDevice through the raw pointer. So the window between losing the owner and using the pointer is however long the other host takes to show up.
That's a plain dangling pointer. A raw PjRtBuffer* can get dereferenced after the unique_ptr that owned it is gone, no runtime internals required.
The fix changes CrossHostSendBuffers to take absl::Span<const DeviceBufferRef>, and each send callback captures its buffer's device_buffer_list and resets it when the send finishes. It also deletes the per-chunk .contiguous() here and the identical one in WindowTpu::signal instead of turning them into real copies, since sending a view is fine once the send holds on to its buffer.
Two more turned up next door in process_group_tpu.cc. experimental_recv built its TpuWork with no future, so wait() returned while the receive could still be in flight, and TpuWork's constructor captured [this] in a callback that dangles if you drop the work handle first. The same branch fixes both.
Moving the hold into the future
All three fixes come down to making something carry the buffer until the copy finishes. Most of the hold plumbing (AsyncDmaResult, both RawBufferHolders, batch_holds) lived on the raw DMA path, so that's where the API changed. StartD2H and the P2P send aren't raw DMA, so they keep their capture-and-reset fixes for now.
The new type is PreparedRawDmaBuffer. Prepare(DeviceBufferRef) waits for the buffer to be ready and creates the raw PJRT alias once. CopyDtoH and CopyHtoD then return an xla::Future<> (wrapped in absl::StatusOr) and nothing else. Internally, a non-copyable Holder owns the DeviceBufferRef, the C API pointers and the alias, and PreparedRawDmaBuffer shares it through a shared_ptr. Each copy hands that shared_ptr to its own completion callback (pjrt_transfer.cc on the refactor branch):
auto [promise, future] = xla::MakePromise<void>();
pjrt::ConvertCEventToCppFuture(args.event, holder_->c_api)
.OnReady([holder = holder_,
promise = std::move(promise)](absl::Status status) mutable {
holder.reset();
promise.Set(std::move(status));
});
return future;
The future keeps its own buffer alive, so there's nothing extra to hand back. Several partial copies from one Prepare share one Holder, and the alias is destroyed after the last of them finishes. On the caller side, WindowTpu becomes c10::make_intrusive<WorkTpu>(std::move(dma_future)), and the batch H2D loop is a Prepare followed by CopyHtoD.
That callback runs on a PJRT thread, so the GIL question comes back. The only thing it drops is the DeviceBufferRef, and DeviceBufferList holds no Python objects, so dropping it there doesn't need the GIL. TransferFuture still keeps the user's at::Tensors, but those get released in Await() on the caller's side, not in a callback.
Stuff that's gone:
struct AsyncDmaResult.- Both
RawBufferHolders, plusRawBufferContext,PrepareRawBufferand the forward declaration inbatch_transfer.h. TransferFuture::holds_,TransferFuture::c_api_holds_, and the twoAndThenoverloads that took them.batch_holds,batch_c_api_holds, and all seven rawpjrt::C API calls inbatch_transfer.cc, which went from 617 lines to 513.WorkTpu::hold_, along with its= nullptrdefault.
git grep AsyncDmaResult on the refactor branch comes back empty, and on top of the H2D fix the refactor is net -119 lines.
What's next
StartD2Hand the cross-host send still pin their buffers in their own lambdas. It's the same capture-and-reset pattern each time, and I'd like those paths to share one abstraction the way raw DMA now does.- Tests. The
StartD2Hand P2P fixes have none, and the batch-transfer tests need real hardware ("Requires TPU v5lite hardware"). - Upstream TODOs worth watching: replacing the two-sided send with one-sided RDMA over ICI "once the PJRT C API exports direct cross-chip DMA", folding
experimental_send/experimental_recvintosend/recv, and the commit that forces blocking for non-contiguous host-to-TPU copies, which calls itself a hotfix and promises a better solution soon. - Grep the rest of the repo for futures that travel with a
shared_ptr<void>next to them. I doubtAsyncDmaResultwas the only one ;)
Until the next one, may all your hold vectors have something in them!