tools/export_gdn_prefill.py

#!/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()