spdl.io.transfer_tensor_d2h

transfer_tensor_d2h(batch: T, /, *, device: TDevice | str | None = None, stream: torch.cuda.Stream | None = None, producer_event: torch.cuda.Event | None = None) → T[source]

Transfer PyTorch CUDA tensors to CPU through a dedicated stream.

Added in version 0.8.0.

This function transfers tensors directly to CPU with torch.Tensor.to() on a dedicated CUDA stream. Non-blocking D2H transfers produce page-locked (pinned) CPU tensors.

The transfer process: 1. Traverses the batch recursively. 2. Asynchronously transfers each matching tensor directly to CPU. 3. Synchronizes the transfer stream before returning the rebuilt batch.

The returned CPU tensors remain page-locked. Retaining many transferred batches therefore retains the corresponding amount of pinned host memory.

By default, the function records an event on the calling thread’s current stream and makes the copy stream wait for that event. When dispatching the function to another CPU thread, record the event on the producer thread before dispatch and pass it as producer_event. This captures the exact readiness boundary without waiting for later foreground work. The function waits for the copy stream before returning.

Example

import torch
from spdl.io import transfer_tensor_d2h

batch = {"scores": torch.randn(32, device="cuda:0")}
cpu_batch = transfer_tensor_d2h(batch, device="cuda:0")
Parameters:

  • batch – A torch.Tensor or a composition of tensors with container types such as list, tuple, dict and dataclass.

  • device –

    Optional CUDA device to transfer data from.

    If None the source device is determined by the LOCAL_RANK environment variable. If not set, cuda:0 is used.

  • stream –

    Optional Custom CUDA stream to use for the transfer. If None, a stream is created from the device argument, and cached to a thread-local storage for future reuse.

    When stream is not None, the device argument must be provided. The stream must be on the same device.

  • producer_event – Optional Event marking when the input tensors are ready. If None, an event is recorded on the calling thread’s current stream. Callers that dispatch this function to another CPU thread must record and pass an event from the producer thread.

Returns:

An object of the same type as the input, but the PyTorch CUDA tensors on the specified device are transferred to CPU.

If there is no PyTorch tensor in the input, the input is returned as-is.

If there is no CUDA device available, the input is returned as-is.

Raises:

  • ValueError – If device is not a CUDA device, or a custom stream does not match the requested device.

  • RuntimeError – If the resolved CUDA device index is unavailable.