Module mem

Expand description

Memory-related helpers. Thin safe wrappers around the CUDA driver memory API.

Functions§

alloc: Allocate byte_size bytes of device memory in the current CUDA context.
alloc_host: Allocate n_elems elements of T in page-locked (pinned) host memory.
copy_d_to_h^⚠: Copy count elements of type T from device memory to host memory.
copy_h_to_d^⚠: Copy count elements of type T from host memory to device memory.
copy_rows_d_to_d: Copy num_rows rows of row_bytes bytes from a device buffer with src_stride_bytes row stride to a device buffer with dst_stride_bytes row stride. Used to depad TMA-aligned output tensors back to tight NCHW.
free: Free a device allocation previously returned by alloc.
free_host^⚠: Free a pinned host allocation returned by alloc_host.