Safetensors Utils#

struct TensorMetadata#

Metadata for one tensor stored in a safetensors payload.

Public Members

std::string name#
nvinfer1::DataType dataType = {}#
Coords shape#
size_t offset = {0}#

Byte offset relative to FileMetadata::dataOffset.

size_t bytes = {0}#
struct FileMetadata#

Parsed safetensors header. Parsing touches only the fixed prefix and JSON header; tensor payload pages remain demand-paged.

Public Members

size_t dataOffset = {0}#

Byte offset of the tensor payload in the file.

std::vector<TensorMetadata> tensors#
nvinfer1::DataType trt_edgellm::rt::safetensors::dataTypeFromString(
std::string_view dtype
)#

Convert a safetensors dtype name to the corresponding TensorRT type.

FileMetadata trt_edgellm::rt::safetensors::parseMetadata(
void const *data,
size_t bytes,
std::string_view source
)#

Parse and validate a safetensors header from an existing file mapping.

bool trt_edgellm::rt::safetensors::saveSafetensors(
std::filesystem::path const &filePath,
std::vector<Tensor> const &tensors,
cudaStream_t stream
)#

Save tensors to a safetensors file.

Each tensor in the vector must have a unique name associated with it.

Parameters:

  • filePath – Path to output safetensors file

  • tensors – Vector of tensors to save

  • stream – CUDA stream for async operations

Throws:

std::runtime_error – If CUDA operations fail or data type is unsupported

Returns:

True on success, false on failure

bool trt_edgellm::rt::safetensors::loadSafetensors(
std::filesystem::path const &filePath,
std::vector<Tensor> &tensors,
cudaStream_t stream
)#

Load tensors from a safetensors file.

Tensors are loaded into the provided vector. Each tensor owns its memory.

Parameters:

  • filePath – Path to input safetensors file

  • tensors – Output vector to store loaded tensors

  • stream – CUDA stream for async operations

Throws:

std::runtime_error – If CUDA operations fail, data type is unsupported, or tensor allocation fails

Returns:

True on success, false on failure