WeightsReader.ixx File Reference

Reader for a Mila weights file. More...

#include <cstdio>
#include <filesystem>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <format>
#include <algorithm>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <exception>
#include <type_traits>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
import Compute.CpuMemoryResource;
import Compute.Device;
import Dnn.TensorTypes;
import Dnn.TensorDataTypeTraits;
import Dnn.TensorDataType;
import Dnn.ITensor;
import Dnn.TensorBuffer;
import Dnn.Tensor;
import Serialization.SafeTensors;
import Serialization.Tensor;
import Serialization.OpenMode;
import Serialization.Serializer;
import nlohmann.json;

Classes
struct	Mila::Dnn::Serialization::TensorBlobMetadata
	Metadata for a tensor blob in a weights file. More...
struct	Mila::Dnn::Serialization::WeightsMetadata
	The model configuration a weights file carries. More...
class	Mila::Dnn::Serialization::WeightsReader
	Reader for a Mila weights file. More...

Namespaces
namespace	Mila
	Mila main API namespace.

Enumerations
enum class	DType : uint32_t {   Float32 = 0 , Float16 = 1 , BFloat16 = 2 , Int32 = 3 ,   UInt8 = 4 , Float8E4M3 = 5 , Float8E5M2 = 6 , Int8 = 7 }

Functions
TensorDataType	Mila::Dnn::Serialization::dtypeToTensorDataType (uint32_t dtype)
uint32_t	Mila::Dnn::Serialization::tensorDataTypeToWireCode (TensorDataType type)
std::string	Mila::Dnn::Serialization::toMetadataJSON (const WeightsMetadata &metadata)
	Serialize WeightsMetadata to the JSON the reader parses back.

Detailed Description

Reader for a Mila weights file.

Provides direct access to model weights stored in either Mila container. Used by the models' load() factories.

The whole file is memory-mapped at construction (CreateFileMapping/MapViewOfFile on Windows, mmap on POSIX). streamTensorBlobs() consumes blobs in ascending file offset order, turning the former 224+ random per-tensor reads into a single sequential scan. On the CUDA path a background producer thread stages each blob into a pinned host buffer (double-buffered) so disk I/O overlaps the H2D copy; all CUDA work stays on the consuming thread. Staging reads each blob directly from the file handle (positioned ReadFile/pread), not by faulting through the mapped view – the map's 4 KB on-demand faults throttle a large model well below disk bandwidth once the file no longer fits resident alongside the model's own GPU staging. The mapping is kept for oversized blobs (consumed straight from the view) and the legacy random-access readTensorBlob<MR>() fallback.

Function Documentation

toMetadataJSON()

std::string Mila::Dnn::Serialization::toMetadataJSON ( const WeightsMetadata & metadata )

inlineexport

Serialize WeightsMetadata to the JSON the reader parses back.

The inverse of parseMetadataJSON, which until now had none – an artifact could be inspectable without being loadable. Every field the parser extracts is emitted, so a written model carries the same architecture description a converted one does.

Keys are quoted on both sides by the parser, so no key can match inside a longer one ("rope_theta" does not match within "rope_theta_local"). Do not introduce a key that is a prefix of another up to its closing quote.