File size: 6,810 Bytes
9e551a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | # Tensor Roll — Recursive CUDA-Q Model Quantizer
# Copyright (C) 2026 SnapKitty Collective
# SPDX-License-Identifier: AGPL-3.0-or-later
"""The .trq container format (spec section 10) — a real binary codec.
Layout (all integers little-endian):
magic 4 bytes b'TRQ1'
header_len u32
header JSON {version, arch, n_tensors, provenance, roll_meta, created}
per tensor:
name_len u16
name bytes
ndim u8
shape ndim x u64
dtype_code u8 0=fp32 1=fp16 2=bf16 3=int8 4=int4-packed
quant_code u8 0=none 1=int8-asym 2=int4-group32 3=fp16-cast 4=bf16-cast
meta_len u32
meta JSON {quantization metadata: scales/zeros as base64 f64,
roll metadata, codebooks}
data_len u64
data bytes
checksum 32 bytes sha256(data)
footer 32 bytes sha256(all preceding bytes)
read_trq() verifies the file checksum and every per-tensor checksum and
raises on any mismatch. Corrupt files never load silently.
"""
from __future__ import annotations
import base64
import hashlib
import json
import struct
import numpy as np
from .util import utc_now_iso
MAGIC = b"TRQ1"
VERSION = 1
_DTYPE_CODES = {"fp32": 0, "fp16": 1, "bf16": 2, "int8": 3, "int4": 4}
_DTYPE_NAMES = {v: k for k, v in _DTYPE_CODES.items()}
_NP_DTYPES = {"fp32": np.float32, "fp16": np.float16, "bf16": np.uint16,
"int8": np.int8, "int4": np.uint8}
_QUANT_CODES = {"none": 0, "int8-asym": 1, "int4-group32": 2,
"fp16-cast": 3, "bf16-cast": 4}
_QUANT_NAMES = {v: k for k, v in _QUANT_CODES.items()}
def _b64(arr) -> str:
return base64.b64encode(np.asarray(arr, dtype=np.float64).tobytes()).decode()
def _unb64(s: str) -> np.ndarray:
return np.frombuffer(base64.b64decode(s.encode()), dtype=np.float64)
def _meta_to_json(meta: dict) -> dict:
m = dict(meta)
for k, v in list(m.items()):
if isinstance(v, np.ndarray):
# codebooks, scales, zero-points, or any ndarray metadata
m[k] = {"__b64_f64__": _b64(np.asarray(v, dtype=np.float64)),
"__shape__": list(v.shape), "__key__": k}
return m
def _meta_from_json(m: dict) -> dict:
m = dict(m)
for k, v in list(m.items()):
if isinstance(v, dict) and "__b64_f64__" in v:
arr = _unb64(v["__b64_f64__"])
if "__shape__" in v:
arr = arr.reshape(v["__shape__"])
m[k] = arr
return m
class TrqError(ValueError):
pass
def write_trq(path: str, *, arch: dict, tensors: dict, provenance: dict,
roll_meta: dict = None) -> dict:
"""tensors: {name: (payload_array, fmt_name, quant_meta_dict)}.
Returns {"path":..., "bytes":..., "sha256":...}."""
buf = bytearray()
buf += MAGIC
header = {
"version": VERSION,
"arch": arch,
"n_tensors": len(tensors),
"provenance": provenance,
"roll_meta": roll_meta or {},
"created": utc_now_iso(),
}
hbytes = json.dumps(header).encode()
buf += struct.pack("<I", len(hbytes)) + hbytes
for name, (payload, fmt, qmeta) in tensors.items():
payload = np.ascontiguousarray(payload)
nbytes = payload.tobytes()
name_b = name.encode()
buf += struct.pack("<H", len(name_b)) + name_b
buf += struct.pack("<B", payload.ndim)
buf += struct.pack(f"<{payload.ndim}Q", *payload.shape)
buf += struct.pack("<B", _DTYPE_CODES[fmt])
quant_name = {"fp32": "none", "fp16": "fp16-cast", "bf16": "bf16-cast",
"int8": "int8-asym", "int4": "int4-group32"}[fmt]
buf += struct.pack("<B", _QUANT_CODES[quant_name])
meta = _meta_to_json({"fmt": fmt, **qmeta,
"roll": (roll_meta or {}).get(name, {})})
mbytes = json.dumps(meta).encode()
buf += struct.pack("<I", len(mbytes)) + mbytes
buf += struct.pack("<Q", len(nbytes)) + nbytes
buf += hashlib.sha256(nbytes).digest()
digest = hashlib.sha256(bytes(buf)).digest()
buf += digest
with open(path, "wb") as f:
f.write(bytes(buf))
return {"path": path, "bytes": len(buf), "sha256": digest.hex()}
def read_trq(path: str):
"""Returns (header, tensors) with all checksums verified.
tensors: {name: (payload_array, meta_dict)}. Raises TrqError on corruption."""
with open(path, "rb") as f:
blob = f.read()
if len(blob) < 4 + 32 or blob[:4] != MAGIC:
raise TrqError("not a .trq file (bad magic)")
file_sum, body = blob[-32:], blob[:-32]
if hashlib.sha256(body).digest() != file_sum:
raise TrqError("file checksum mismatch — refusing to load")
off = 4
hlen = struct.unpack_from("<I", body, off)[0]
off += 4
header = json.loads(body[off:off + hlen])
off += hlen
tensors = {}
for _ in range(header["n_tensors"]):
nlen = struct.unpack_from("<H", body, off)[0]
off += 2
name = body[off:off + nlen].decode()
off += nlen
ndim = struct.unpack_from("<B", body, off)[0]
off += 1
shape = struct.unpack_from(f"<{ndim}Q", body, off)
off += 8 * ndim
dtype_code = struct.unpack_from("<B", body, off)[0]
off += 1
quant_code = struct.unpack_from("<B", body, off)[0]
off += 1
mlen = struct.unpack_from("<I", body, off)[0]
off += 4
meta = _meta_from_json(json.loads(body[off:off + mlen]))
off += mlen
dlen = struct.unpack_from("<Q", body, off)[0]
off += 8
data = body[off:off + dlen]
off += dlen
cksum = body[off:off + 32]
off += 32
if hashlib.sha256(data).digest() != cksum:
raise TrqError(f"tensor {name!r} checksum mismatch — refusing to load")
fmt = _DTYPE_NAMES[dtype_code]
arr = np.frombuffer(data, dtype=_NP_DTYPES[fmt])
if fmt == "int4":
arr = arr.reshape(-1) # packed nibbles stay flat; dequantize reshapes
else:
arr = arr.reshape(shape)
tensors[name] = (arr, meta)
return header, tensors
def total_params(tensors: dict) -> int:
"""Total logical parameters across tensors (uses recorded shapes, so
packed INT4 counts its true parameters, not its packed bytes)."""
total = 0
for _name, (arr, meta) in tensors.items():
shape = meta.get("shape") or list(arr.shape)
total += int(np.prod(shape))
return total
def param_budget_ok(tensors: dict, cap: float = 8_500_000_000) -> tuple:
"""Verify the exported artifact obeys the hard parameter budget.
Returns (total_params, ok)."""
total = total_params(tensors)
return total, total <= cap
|