Checkpoint Reader

class CheckpointReader

Read individual tensors from Hugging Face safetensors or indexed PyTorch ZIP checkpoint ranges.

Shards are memory-mapped on first use. The reader is only used while model weights are prepared during runtime initialization and is destroyed before the first inference enqueue.

Public Functions

explicit CheckpointReader(std::filesystem::path const &directory)

CheckpointReader(

std::filesystem::path const &directory,

std::vector<TensorLocation> const &explicitTensorLocations

)

~CheckpointReader() noexcept

void registerTensors(std::vector<std::string> const &names)

CUDA-map the page ranges containing the requested tensors. Ranges are merged per shard and must be registered before find() returns a device pointer.

View registerTensorRange(

std::string const &name,

size_t offset,

size_t bytes

)

CUDA-map one byte range inside a checkpoint tensor.

No other ranges may be registered while the returned view is used. unregisterTensors() releases the mapping after its consumer stream has completed.

void unregisterTensors() noexcept

Release every CUDA registration created by registerTensors().

Runtime initialization calls this after each completed output binding so source checkpoint pages never accumulate beside the final weight arena.

void discardTensors(std::vector<std::string> const &names) noexcept

Evict no-longer-needed source pages from this process’s file mappings.

void discardTensorRange(

std::string const &name,

size_t offset,

size_t bytes

) noexcept

inline size_t peakRegisteredBytes() const noexcept

inline size_t registeredBytes() const noexcept

inline std::chrono::nanoseconds registrationTime() const noexcept

bool findHost(std::string const &name, View &view) const

Find a tensor without requiring CUDA page registration. The returned view has a host pointer and a null device pointer.

bool find(std::string const &name, View &view) const

struct TensorLocation

Public Members

std::string name

std::filesystem::path file

std::string dtype

Coords shape

size_t offset = {0}

size_t bytes = {0}

struct View

A checkpoint tensor has separate host and CUDA aliases into one mapped file range. This is intentionally not rt::Tensor: Tensor has one mutable address and one device type, so wrapping this range would discard either the const host view used for metadata checks or the CUDA alias used by transform kernels.

Public Members

nvinfer1::DataType dtype = {}

Coords shape

uint8_t const *data = {nullptr}

uint8_t const *deviceData = {nullptr}

size_t bytes = {0}

struct TensorLocation

Public Members

std::string name

std::filesystem::path file

std::string dtype

Coords shape

size_t offset = {0}

size_t bytes = {0}

struct View

A checkpoint tensor has separate host and CUDA aliases into one mapped file range. This is intentionally not rt::Tensor: Tensor has one mutable address and one device type, so wrapping this range would discard either the const host view used for metadata checks or the CUDA alias used by transform kernels.

Public Members

nvinfer1::DataType dtype = {}

Coords shape

uint8_t const *data = {nullptr}

uint8_t const *deviceData = {nullptr}

size_t bytes = {0}

struct Registration

Public Members

size_t offset = {0}

size_t bytes = {0}

uint8_t const *deviceData = {nullptr}