#!/usr/bin/env python3 """Export a fixed-token Ling GDN recurrence graph for RK3588 prefill.""" from __future__ import annotations import argparse import hashlib import json from pathlib import Path import numpy as np import torch from torch import nn from torch.nn import functional as F class GdnPrefill(nn.Module): def __init__(self, heads: int, tokens: int, head_dim: int = 128): super().__init__() self.heads = heads self.tokens = tokens self.head_dim = head_dim self.query_scale = head_dim**-0.5 def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, decay: torch.Tensor, beta: torch.Tensor, state: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: shape = (1, self.heads, self.head_dim) recurrent = state.reshape(1, self.heads, self.head_dim, self.head_dim) recurrent = recurrent.transpose(-2, -1) outputs = [] for token in range(self.tokens): q = query[token].reshape(shape) k = key[token].reshape(shape) v = value[token].reshape(shape) d = decay[token].reshape(shape) b = beta[token].reshape(1, self.heads) factor = F.elu(d, alpha=1.0) + 1.0 decayed = recurrent * factor.unsqueeze(-1) predicted = torch.matmul(k.unsqueeze(-2), decayed).squeeze(-2) delta = b.unsqueeze(-1) * (v - predicted) recurrent = decayed + k.unsqueeze(-1) * delta.unsqueeze(-2) output = torch.matmul( (q * self.query_scale).unsqueeze(-2), recurrent ).squeeze(-2) outputs.append(output.reshape(1, self.heads * self.head_dim)) packed_output = torch.cat(outputs, dim=0) final_state = recurrent.transpose(-2, -1).contiguous() final_state = final_state.reshape(1, self.heads * self.head_dim * self.head_dim) return packed_output, final_state def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as source: for chunk in iter(lambda: source.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def inputs_for(heads: int, tokens: int, head_dim: int) -> dict[str, torch.Tensor]: generator = torch.Generator().manual_seed(20260909 + heads + tokens) width = heads * head_dim query = torch.randn((tokens, width), generator=generator) key = torch.randn((tokens, width), generator=generator) query = query.reshape(tokens, heads, head_dim) key = key.reshape(tokens, heads, head_dim) query = query * torch.rsqrt((query * query).sum(-1, keepdim=True) + 1e-6) key = key * torch.rsqrt((key * key).sum(-1, keepdim=True) + 1e-6) return { "query": query.reshape(tokens, width), "key": key.reshape(tokens, width), "value": torch.randn((tokens, width), generator=generator) * 0.2, "decay": -5.0 * torch.sigmoid( torch.randn((tokens, width), generator=generator) * 0.2 ), "beta": torch.sigmoid( torch.randn((tokens, heads), generator=generator) * 0.2 ), "state": torch.randn( (1, heads * head_dim * head_dim), generator=generator ) * 0.01, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--heads", type=int, choices=(5, 6), required=True) parser.add_argument("--tokens", type=int, default=32) parser.add_argument("--head-dim", type=int, default=128) parser.add_argument("--skip-build", action="store_true") parser.add_argument("--quantize", action="store_true") parser.add_argument("--calibration-samples", type=int, default=32) parser.add_argument("--verbose", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() if args.tokens < 1: raise ValueError("tokens must be positive") if args.calibration_samples < 1: raise ValueError("calibration samples must be positive") output_dir = args.output_dir.resolve() data_dir = output_dir / "data" data_dir.mkdir(parents=True, exist_ok=True) model = GdnPrefill(args.heads, args.tokens, args.head_dim).eval() inputs = inputs_for(args.heads, args.tokens, args.head_dim) input_names = list(inputs) output_names = ["output", "new_state"] with torch.no_grad(): reference = model(**inputs) for name, tensor in inputs.items(): tensor.detach().cpu().numpy().astype(np.float32).tofile(data_dir / f"{name}.f32") for name, tensor in zip(output_names, reference, strict=True): tensor.detach().cpu().numpy().astype(np.float32).tofile(data_dir / f"{name}.f32") stem = f"ling3_gdn_prefill_t{args.tokens}_h{args.heads}" onnx_path = output_dir / f"{stem}.onnx" torch.onnx.export( model, tuple(inputs.values()), onnx_path, input_names=input_names, output_names=output_names, opset_version=17, do_constant_folding=True, ) dataset_path = output_dir / "dataset.txt" calibration_dir = output_dir / "calibration" if args.quantize: calibration_dir.mkdir(parents=True, exist_ok=True) dataset_lines: list[str] = [] for sample in range(args.calibration_samples): sample_inputs = inputs_for(args.heads, args.tokens, args.head_dim) paths: list[str] = [] for input_index, (name, tensor) in enumerate(sample_inputs.items()): generator = np.random.default_rng(20260909 + sample * 17 + input_index) array = tensor.detach().cpu().numpy().astype(np.float32) if sample > 0: array *= generator.uniform(0.5, 2.0) path = calibration_dir / f"sample_{sample:03d}_{name}.npy" np.save(path, array) paths.append(str(path)) dataset_lines.append(" ".join(paths)) dataset_path.write_text("\n".join(dataset_lines) + "\n", encoding="utf-8") rknn_path = output_dir / f"{stem}_{'w8a8' if args.quantize else 'fp16'}_rk3588.rknn" build_results: dict[str, int] = {} if not args.skip_build: from rknn.api import RKNN converter = RKNN(verbose=args.verbose) try: config = {"target_platform": "rk3588", "optimization_level": 3} if args.quantize: config.update( quantized_dtype="w8a8", quantized_algorithm="normal", quantized_method="channel", ) build_results["config"] = converter.config(**config) build_results["load_onnx"] = converter.load_onnx(model=str(onnx_path)) if any(build_results.values()): raise RuntimeError(f"RKNN setup failed: {build_results}") build_results["build"] = converter.build( do_quantization=args.quantize, dataset=str(dataset_path) if args.quantize else None, ) if build_results["build"] != 0: raise RuntimeError(f"RKNN build failed: {build_results}") build_results["export"] = converter.export_rknn(str(rknn_path)) if build_results["export"] != 0: raise RuntimeError(f"RKNN export failed: {build_results}") finally: converter.release() manifest = { "architecture": "Ling GDN fixed-token prefill recurrence", "heads": args.heads, "head_dim": args.head_dim, "tokens": args.tokens, "quantization": "w8a8" if args.quantize else "fp16", "calibration_samples": args.calibration_samples if args.quantize else 0, "inputs": {name: list(value.shape) for name, value in inputs.items()}, "outputs": { name: list(value.shape) for name, value in zip(output_names, reference, strict=True) }, "onnx": {"path": str(onnx_path), "sha256": sha256(onnx_path)}, "rknn": ( {"path": str(rknn_path), "sha256": sha256(rknn_path)} if rknn_path.exists() else None ), "build_results": build_results, } (output_dir / "manifest.json").write_text( json.dumps(manifest, indent=2) + "\n", encoding="utf-8" ) print(json.dumps(manifest, indent=2)) if __name__ == "__main__": main()