aitune.torch.libs.cuda.memory | NVIDIA AITune Documentation

CUDA memory utilities for GPU-to-GPU transfers.

Module Contents

Functions

Name	Description
`memcpy_to_torch`	Copy a CUDA device buffer into a new torch tensor via D2D memcpy (no CPU round-trip).

Data

_ORT_DTYPE_TO_TORCH

API

aitune.torch.libs.cuda.memory.memcpy_to_torch(
    ptr: int,
    shape: list[int],
    ort_dtype: str,
    device: torch.device
) -> torch.Tensor

Copy a CUDA device buffer into a new torch tensor via D2D memcpy (no CPU round-trip).

Parameters:

ptr

int

Raw CUDA pointer to the source buffer (e.g. OrtValue.data_ptr()).

shape

list[int]

Shape of the buffer.

ort_dtype

str

ORT data type string (e.g. "tensor(float)").

device

torch.device

Target CUDA device for the output tensor.

Returns: torch.Tensor

A new CUDA torch tensor containing a copy of the source buffer.

Raises:

KeyError: If ort_dtype is not a recognised ORT type string.
RuntimeError: If the CUDA memcpy fails.

aitune.torch.libs.cuda.memory._ORT_DTYPE_TO_TORCH: dict[str, dtype] = {'tensor(float)': torch.float32, 'tensor(float16)': torch.float16, 'tensor(doubl...