#!/usr/bin/env python3 """Pack pre-quantized Ling-3 assets into the mmap-friendly .l3r format.""" from __future__ import annotations import argparse import hashlib import json import os import struct import zlib from pathlib import Path ALIGNMENT = 4096 HEADER = struct.Struct("<8s6I6Q20sI28I4f24s") ENTRY = struct.Struct("<12I4Q2f4I32s24s") DTYPES = { "unknown": 0, "bf16": 1, "fp16": 2, "fp32": 3, "f32": 3, "int8": 4, "int4": 5, "int16": 6, "int32": 7, "rknn": 8, } ROLES = { "unknown": 0, "embedding": 1, "norm": 2, "linear_weight": 3, "linear_scale": 4, "bias": 5, "rknn_island": 6, "tokenizer": 7, } QUANTS = {"none": 0, "per_tensor": 1, "per_output_channel": 2, "fixed_activation": 3} LAYOUTS = {"row_major": 0, "rknn_native_b": 1, "packed_int4": 2, "opaque": 3} MODEL_FIELDS = ( "vocab_size", "hidden_size", "layer_count", "attention_heads", "head_dim", "kv_lora_rank", "q_lora_rank", "qk_nope_dim", "qk_rope_dim", "value_head_dim", "dense_ffn_dim", "expert_ffn_dim", "shared_ffn_dim", "expert_count", "experts_per_token", "expert_group_count", "selected_group_count", "layer_group_size", "leading_dense_layers", "convolution_kernel", "max_context", "eos_token", "pad_token", "bos_token", "mla_layer_count", "kda_layer_count", "default_weight_bits", "default_activation_bits", ) def align(value: int) -> int: return (value + ALIGNMENT - 1) & ~(ALIGNMENT - 1) def enum_value(mapping: dict[str, int], value: str, label: str) -> int: try: return mapping[value] except KeyError as error: raise ValueError(f"unknown {label} {value!r}") from error def resolve(base: Path, value: str) -> Path: path = Path(value) return path if path.is_absolute() else base / path def copy_file(source: Path, output, hasher: hashlib._Hash) -> int: total = 0 with source.open("rb") as stream: while chunk := stream.read(8 * 1024 * 1024): output.write(chunk) hasher.update(chunk) total += len(chunk) return total def pad_to(output, offset: int) -> None: current = output.tell() if current > offset: raise ValueError("package layout overlapped while writing") output.write(b"\0" * (offset - current)) def make_header(manifest: dict, tensor_count: int, string_offset: int, string_bytes: int, payload_offset: int, payload_bytes: int, file_bytes: int, crc: int) -> bytes: model = manifest["model"] source_revision = bytes.fromhex(manifest["source_revision"]) if len(source_revision) != 20: raise ValueError("source_revision must be a 40-character Git SHA-1") integers = [int(model[name]) for name in MODEL_FIELDS] floats = [ float(model["rms_epsilon"]), float(model["rope_theta"]), float(model["routed_scale"]), float(model["kda_lower_bound"]), ] return HEADER.pack( b"L3RKNN1\0", 1, HEADER.size, 0x01020304, tensor_count, ENTRY.size, int(manifest.get("flags", 0)), HEADER.size, string_offset, string_bytes, payload_offset, payload_bytes, file_bytes, source_revision, crc, *integers, *floats, bytes(24), ) def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("manifest", type=Path) parser.add_argument("output", type=Path) parser.add_argument( "--include-prefix", action="append", default=[], help="pack only tensors whose names start with this prefix (repeatable)", ) args = parser.parse_args() manifest_path = args.manifest.resolve() manifest = json.loads(manifest_path.read_text(encoding="utf-8")) base = manifest_path.parent tensors = manifest.get("tensors", []) if args.include_prefix: tensors = [ item for item in tensors if any(item["name"].startswith(prefix) for prefix in args.include_prefix) ] if not tensors: raise ValueError("--include-prefix did not match any tensors") names = bytearray() prepared = [] for item in tensors: name = item["name"].encode("utf-8") if not name or b"\0" in name: raise ValueError("tensor names must be non-empty UTF-8 without NUL") dims = [int(value) for value in item.get("dims", [])] if len(dims) > 4: raise ValueError(f"tensor {item['name']!r} has more than four dimensions") source = resolve(base, item["path"]) if not source.is_file(): raise FileNotFoundError(source) aux = resolve(base, item["aux_path"]) if item.get("aux_path") else None if aux is not None and not aux.is_file(): raise FileNotFoundError(aux) prepared.append((item, len(names), name, dims, source, aux)) names.extend(name) table_bytes = len(prepared) * ENTRY.size string_offset = HEADER.size + table_bytes payload_offset = align(string_offset + len(names)) cursor = payload_offset layouts = [] for item, name_offset, name, dims, source, aux in prepared: data_offset = align(cursor) data_bytes = source.stat().st_size cursor = data_offset + data_bytes aux_offset = align(cursor) if aux is not None else 0 aux_bytes = aux.stat().st_size if aux is not None else 0 if aux is not None: cursor = aux_offset + aux_bytes layouts.append((item, name_offset, name, dims, source, aux, data_offset, data_bytes, aux_offset, aux_bytes)) file_bytes = align(cursor) payload_bytes = file_bytes - payload_offset args.output.parent.mkdir(parents=True, exist_ok=True) entries = [] with args.output.open("w+b") as output: output.write(bytes(HEADER.size + table_bytes)) output.write(names) pad_to(output, payload_offset) for item, name_offset, name, dims, source, aux, data_offset, data_bytes, aux_offset, aux_bytes in layouts: pad_to(output, data_offset) digest = hashlib.sha256() if copy_file(source, output, digest) != data_bytes: raise IOError(f"short read while packing {source}") if aux is not None: pad_to(output, aux_offset) aux_digest = hashlib.sha256() if copy_file(aux, output, aux_digest) != aux_bytes: raise IOError(f"short read while packing {aux}") padded_dims = dims + [0] * (4 - len(dims)) entries.append(ENTRY.pack( name_offset, len(name), enum_value(DTYPES, item["dtype"], "dtype"), len(dims), enum_value(ROLES, item.get("role", "unknown"), "role"), enum_value(QUANTS, item.get("quant", "none"), "quantization"), enum_value(LAYOUTS, item.get("layout", "row_major"), "layout"), int(item.get("flags", 0)), *padded_dims, data_offset, data_bytes, aux_offset, aux_bytes, float(item.get("scale", 1.0)), float(item.get("zero_point", 0.0)), int(item.get("layer", 0xFFFFFFFF)), int(item.get("op", 0)), int(item.get("expert", 0xFFFFFFFF)), int(item.get("core", 0xFFFFFFFF)), digest.digest(), bytes(24), )) pad_to(output, file_bytes) header = make_header(manifest, len(entries), string_offset, len(names), payload_offset, payload_bytes, file_bytes, 0) crc = zlib.crc32(header) & 0xFFFFFFFF header = make_header(manifest, len(entries), string_offset, len(names), payload_offset, payload_bytes, file_bytes, crc) output.seek(0) output.write(header) for entry in entries: output.write(entry) output.flush() os.fsync(output.fileno()) print(f"packed {len(entries)} tensors, {file_bytes} bytes -> {args.output}") return 0 if __name__ == "__main__": raise SystemExit(main())