Loading Safetensors in C++

First, use open() to open the file and check the file size.

fd = open(path.c_str(), O_RDONLY);
file_size = lseek(fd, 0, SEEK_END);

To avoid loading the entire file into memory, I use mmap() to map the file into virtual memory.

mapped = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0);

The first uint64_t in a safetensors file stores the header size.

uint64_t header_size;
memcpy(&header_size, mapped, sizeof(uint64_t));

char *header_ptr = static_cast<char *>(mapped) + sizeof(uint64_t);

std::string header(header_ptr, header_size);

The header contains the metadata information of each tensor. To parse the metadata, I use nlohmann::json.

metadata = nlohmann::json::parse(header);

The actual model parameters start after the header:

data_offset = sizeof(uint64_t) + header_size;

To access each layer's weights easily, I created a Tensor structure:

struct Tensor
{
    std::string name;
    std::string dtype;

    std::vector<int64_t> shape;

    uint64_t start;
    uint64_t end;

    char *data;
    size_t size;
};

A get_tensor() function is used to retrieve each tensor from the safetensors file. The important part is calculating the correct memory position using the tensor's data offsets.

Tensor get_tensor(const std::string &name)
{
    if (!metadata.contains(name))
        throw std::runtime_error("Tensor not found: " + name);

    auto json = metadata[name];

    Tensor tensor;

    tensor.name = name;
    tensor.dtype = json["dtype"];
    tensor.shape = json["shape"].get<std::vector<int64_t>>();

    tensor.start = json["data_offsets"][0].get<uint64_t>();
    tensor.end = json["data_offsets"][1].get<uint64_t>();

    tensor.size = tensor.end - tensor.start;

    tensor.data =
        static_cast<char *>(mapped) + data_offset + tensor.start;

    return tensor;
}

Since my PC has limited memory, I avoid loading the entire model weights. Instead, the weights remain inside the memory-mapped file, and the values are converted only when needed during inference. Note that this does not copy the actual parameters into memory; it only creates a pointer to the corresponding location inside the memory-mapped file.
During inference, the tensor data can be accessed by casting the raw memory pointer into the correct data type.

For BF16 weights:

embedding_weights = reinterpret_cast<uint16_t *>(tensor.data);
output.resize(embedding_dim);
size_t offset = token_id * embedding_dim;
for (int i = 0; i < embedding_dim; i++)
{
    output[i] = bf16_to_float(embedding_weights[offset + i]);
}

Alternatively, for faster inference, the tensor can be copied into allocated memory instead of accessing the memory-mapped file directly.

tensor.data = static_cast<char *>(malloc(tensor.size));
memcpy(
    tensor.data,
    static_cast<char *>(mapped) + data_offset + start,
    tensor.size
);

uint16_t *weights = reinterpret_cast<uint16_t *>(tensor.data);
for (int i = 0; i < size; i++)
{
    float value = bf16_to_float(weights[i]);
}

This loads the tensor data into normal allocated memory, allowing direct access during inference.