include/brotensor/safetensors.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // safetensors reader + writer for brotensor. | ||
| 4 | // | ||
| 5 | // Format (from huggingface/safetensors): | ||
| 6 | // bytes 0..7 : little-endian u64 header_size N | ||
| 7 | // bytes 8..8+N : UTF-8 JSON header | ||
| 8 | // bytes 8+N.. : raw tensor payload | ||
| 9 | // | ||
| 10 | // Header JSON is an object mapping tensor_name -> { | ||
| 11 | // "dtype": "F32" | "F16" | "BF16" | "I32" | "I64" | "U8" | "BOOL", | ||
| 12 | // "shape": [N1, N2, ...], | ||
| 13 | // "data_offsets": [start, end] // bytes relative to start of payload | ||
| 14 | // } | ||
| 15 | // plus an optional "__metadata__" entry (string -> string) that we skip. | ||
| 16 | // | ||
| 17 | // The file is mmap'd; TensorView::data points directly into the mapping, so | ||
| 18 | // it's valid only while the owning File is alive. Zero-copy on the read path. | ||
| 19 | // | ||
| 20 | // safetensors is a tensor-container format — a named collection of typed N-D | ||
| 21 | // arrays. It lives in brotensor because its natural output type is | ||
| 22 | // brotensor::Tensor; the upload* helpers below depend only on brotensor. | ||
| 23 | |||
| 24 | #include <cstddef> | ||
| 25 | #include <cstdint> | ||
| 26 | #include <functional> | ||
| 27 | #include <string> | ||
| 28 | #include <string_view> | ||
| 29 | #include <unordered_map> | ||
| 30 | #include <vector> | ||
| 31 | |||
| 32 | #include <brotensor/detail/string_hash.h> | ||
| 33 | |||
| 34 | namespace brotensor { | ||
| 35 | struct Tensor; | ||
| 36 | enum class Dtype; // brotensor/tensor.h — target dtype for upload_as | ||
| 37 | } | ||
| 38 | |||
| 39 | namespace brotensor::safetensors { | ||
| 40 | |||
| 41 | enum class Dtype { | ||
| 42 | Unknown = 0, | ||
| 43 | F32, | ||
| 44 | F16, | ||
| 45 | BF16, | ||
| 46 | I32, | ||
| 47 | I64, | ||
| 48 | U8, | ||
| 49 | BOOL, | ||
| 50 | }; | ||
| 51 | |||
| 52 | const char* dtype_name(Dtype d); | ||
| 53 | int dtype_size_bytes(Dtype d); // 0 for Unknown | ||
| 54 | |||
| 55 | 82 | struct TensorView { | |
| 56 | std::string name; | ||
| 57 | 41 | Dtype dtype = Dtype::Unknown; | |
| 58 | std::vector<int64_t> shape; | ||
| 59 | 41 | const uint8_t* data = nullptr; // into the mmap; valid while File lives | |
| 60 | 41 | std::size_t nbytes = 0; | |
| 61 | |||
| 62 | int64_t numel() const; // product of shape | ||
| 63 | }; | ||
| 64 | |||
| 65 | class File { | ||
| 66 | public: | ||
| 67 | 72 | File() = default; | |
| 68 | ~File(); | ||
| 69 | |||
| 70 | File(File&& other) noexcept; | ||
| 71 | File& operator=(File&& other) noexcept; | ||
| 72 | File(const File&) = delete; | ||
| 73 | File& operator=(const File&) = delete; | ||
| 74 | |||
| 75 | // Open a .safetensors file. Throws std::runtime_error on I/O or parse | ||
| 76 | // failure. The returned File owns the mmap and tensor index. | ||
| 77 | static File open(const std::string& path); | ||
| 78 | |||
| 79 | // Lookup by name. find() returns nullptr if missing; get() throws. | ||
| 80 | const TensorView* find(std::string_view name) const; | ||
| 81 | const TensorView& get(std::string_view name) const; | ||
| 82 | |||
| 83 | 7 | const std::vector<TensorView>& tensors() const { return tensors_; } | |
| 84 | 14 | std::size_t size() const { return tensors_.size(); } | |
| 85 | |||
| 86 | private: | ||
| 87 | void release_() noexcept; | ||
| 88 | |||
| 89 | // Platform-specific mmap handles. On Windows: file_ + mapping_ + base_. | ||
| 90 | // On POSIX: file_ holds the fd, mapping_ is unused. | ||
| 91 | 24 | void* file_ = nullptr; | |
| 92 | 24 | void* mapping_ = nullptr; | |
| 93 | 24 | void* base_ = nullptr; | |
| 94 | 24 | std::size_t file_size_ = 0; | |
| 95 | |||
| 96 | std::vector<TensorView> tensors_; | ||
| 97 | std::unordered_map<std::string, std::size_t, detail::StringHash, std::equal_to<>> index_; | ||
| 98 | }; | ||
| 99 | |||
| 100 | // Upload a tensor view as a 2D (rows, cols) Tensor. | ||
| 101 | // | ||
| 102 | // brotensor is 2D-only by design; higher-rank tensors (e.g. conv weights | ||
| 103 | // [Cout, Cin, kH, kW]) are flattened by the caller to whatever 2D layout the | ||
| 104 | // consuming op expects. | ||
| 105 | // | ||
| 106 | // Requirements: | ||
| 107 | // - view.dtype must be F32, F16, or BF16 | ||
| 108 | // - rows * cols * dtype_size_bytes(view.dtype) == view.nbytes | ||
| 109 | // - dst is resized to (rows, cols) at the matching brotensor::Dtype | ||
| 110 | // (a BF16 view produces a brotensor::Dtype::BF16 tensor, zero-converted) | ||
| 111 | // | ||
| 112 | // Throws std::runtime_error if the byte count or dtype is wrong. | ||
| 113 | void upload(const TensorView& view, int rows, int cols, brotensor::Tensor& dst); | ||
| 114 | |||
| 115 | // Like upload(), but always produces an FP16 Tensor. If the source view is | ||
| 116 | // F32, it is converted host-side via brotensor::fp32_to_fp16_bits before | ||
| 117 | // upload. If the source is F16, it's uploaded as-is. | ||
| 118 | void upload_fp16(const TensorView& view, int rows, int cols, brotensor::Tensor& dst); | ||
| 119 | |||
| 120 | // Upload a weight at brotensor's compute dtype — FP32 on the CPU backend, | ||
| 121 | // FP16 on a GPU backend (see brotensor::compute_dtype()). The source view | ||
| 122 | // (F16, F32, or BF16) is converted host-side as needed, so a single | ||
| 123 | // checkpoint serves either backend. BF16 is the on-disk dtype of Flux-family | ||
| 124 | // weights; it is widened to FP32 on the CPU backend (which has no BF16 | ||
| 125 | // arithmetic) and narrowed to FP16 on a GPU backend. | ||
| 126 | // Upload `view` (F16 / F32 / BF16 source) as a (rows, cols) tensor at an | ||
| 127 | // EXPLICIT arithmetic dtype (FP32 / FP16 / BF16), converting host-side as | ||
| 128 | // needed, on the current default device. Lets a module pick a compute dtype | ||
| 129 | // different from the global one — e.g. Flux runs BF16 on a GPU backend | ||
| 130 | // whose pipeline dtype is FP16, because its activations overflow FP16. | ||
| 131 | void upload_as(const TensorView& view, int rows, int cols, | ||
| 132 | brotensor::Dtype want, brotensor::Tensor& dst); | ||
| 133 | |||
| 134 | void upload_compute(const TensorView& view, int rows, int cols, | ||
| 135 | brotensor::Tensor& dst); | ||
| 136 | |||
| 137 | // Like upload_compute(), but first validates the view: its dtype must be | ||
| 138 | // F16, F32, or BF16 and its element count must equal rows*cols. Throws | ||
| 139 | // std::runtime_error tagged with `name` (a caller-supplied label) and the | ||
| 140 | // safetensors key on mismatch — the loader entry point a model component | ||
| 141 | // uses so a malformed checkpoint fails with a clear message. | ||
| 142 | void upload_compute_checked(const TensorView& view, int rows, int cols, | ||
| 143 | brotensor::Tensor& dst, const std::string& name); | ||
| 144 | |||
| 145 | // ─── Writer ──────────────────────────────────────────────────────────────── | ||
| 146 | // | ||
| 147 | // Minimal safetensors writer. Builds the JSON header in insertion order and | ||
| 148 | // concatenates raw tensor bytes. All entries' `host_data` must point to | ||
| 149 | // `bytes` bytes of valid host memory matching `dtype`/`shape`. Tensor names | ||
| 150 | // must be unique. Throws std::runtime_error on I/O failure or invalid input. | ||
| 151 | |||
| 152 | 70 | struct WriteEntry { | |
| 153 | std::string name; | ||
| 154 | 28 | Dtype dtype = Dtype::F16; | |
| 155 | std::vector<int64_t> shape; | ||
| 156 | 28 | const void* host_data = nullptr; | |
| 157 | 28 | std::size_t bytes = 0; | |
| 158 | }; | ||
| 159 | |||
| 160 | void write_file(const std::string& path, const std::vector<WriteEntry>& entries); | ||
| 161 | |||
| 162 | } // namespace brotensor::safetensors | ||
| 163 |