export_decoder_qdq_int8.py

from __future__ import annotations

import hashlib
import json
from collections import Counter
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
from onnx import TensorProto, helper, numpy_helper


ROOT = Path(__file__).resolve().parent
OUTPUT_DIR = ROOT / "output"
REPORT_DIR = ROOT / "reports"
GRAPH_NAMES = ("decoder_prefill", "decoder_step")
FORBIDDEN_WEBGPU_OPS = {"DynamicQuantizeLinear", "MatMulInteger"}


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def quantize_columns(values: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Symmetric INT8 weight-only quantization for MatMul's output axis."""
    values = values.astype(np.float32, copy=False)
    scales = np.max(np.abs(values), axis=0) / np.float32(127.0)
    scales = np.where(scales == 0, np.float32(1.0), scales).astype(np.float32)
    quantized = np.clip(np.rint(values / scales), -127, 127).astype(np.int8)
    zero_points = np.zeros(scales.shape, dtype=np.int8)
    return quantized, scales, zero_points


def rewrite_graph(
    source: Path,
    destination: Path,
    *,
    quantize_gather_shapes: frozenset[tuple[int, ...]] = frozenset(),
) -> dict:
    model = onnx.load(source)
    graph = model.graph
    initializers = {value.name: value for value in graph.initializer}
    uses = Counter(name for node in graph.node for name in node.input if name)
    replacements: dict[str, str] = {}
    new_initializers = []
    dq_nodes = []
    quantized_elements = 0
    quantized_matmuls = 0
    quantized_gathers = 0

    for node_index, node in enumerate(graph.node):
        is_matmul = node.op_type == "MatMul" and len(node.input) >= 2
        is_gather = node.op_type == "Gather" and len(node.input) >= 1
        if not is_matmul and not is_gather:
            continue
        weight_input = 1 if is_matmul else 0
        weight_name = node.input[weight_input]
        initializer = initializers.get(weight_name)
        if initializer is None or initializer.data_type != TensorProto.FLOAT:
            continue
        values = numpy_helper.to_array(initializer)
        if values.ndim != 2:
            continue
        if is_gather and tuple(values.shape) not in quantize_gather_shapes:
            continue
        if uses[weight_name] != 1:
            raise RuntimeError(
                f"{source.name}: {weight_name} has {uses[weight_name]} uses; "
                "weight sharing must be handled explicitly"
            )

        quantized, scales, zero_points = quantize_columns(values)
        prefix = f"{weight_name}_qdq_{node_index}"
        quantized_name = f"{prefix}_int8"
        scale_name = f"{prefix}_scale"
        zero_name = f"{prefix}_zero"
        output_name = f"{prefix}_fp32"
        new_initializers.extend(
            (
                numpy_helper.from_array(quantized, quantized_name),
                numpy_helper.from_array(scales, scale_name),
                numpy_helper.from_array(zero_points, zero_name),
            )
        )
        dq_nodes.append(
            helper.make_node(
                "DequantizeLinear",
                [quantized_name, scale_name, zero_name],
                [output_name],
                name=f"{prefix}_dequantize",
                axis=1,
            )
        )
        replacements[weight_name] = output_name
        node.input[weight_input] = output_name
        quantized_elements += values.size
        if is_matmul:
            quantized_matmuls += 1
        else:
            quantized_gathers += 1

    if not replacements:
        raise RuntimeError(f"{source.name}: no constant MatMul weights found")

    retained = [value for value in graph.initializer if value.name not in replacements]
    del graph.initializer[:]
    graph.initializer.extend(retained)
    graph.initializer.extend(new_initializers)
    original_nodes = list(graph.node)
    del graph.node[:]
    graph.node.extend(dq_nodes)
    graph.node.extend(original_nodes)

    model.producer_name = "vibe-manga-baberu-webgpu-qdq"
    model.producer_version = "1"
    onnx.checker.check_model(model)
    onnx.save(model, destination)

    operators = Counter(node.op_type for node in graph.node)
    forbidden = sorted(FORBIDDEN_WEBGPU_OPS.intersection(operators))
    if forbidden:
        raise RuntimeError(f"{destination.name}: forbidden operators: {forbidden}")
    return {
        "source": source.name,
        "bytes": destination.stat().st_size,
        "sha256": sha256(destination),
        "quantized_matmuls": quantized_matmuls,
        "quantized_gathers": quantized_gathers,
        "quantized_elements": quantized_elements,
        "operators": dict(sorted(operators.items())),
        "forbidden_webgpu_operators": forbidden,
    }


def make_inputs(session: ort.InferenceSession) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(20260716)
    feeds: dict[str, np.ndarray] = {}
    for value in session.get_inputs():
        shape = [257 if dimension == "past_len" else dimension for dimension in value.shape]
        if value.name == "vision_embeds":
            feeds[value.name] = rng.normal(0, 0.2, shape).astype(np.float32)
        elif value.name == "token_one_hot":
            token = np.zeros(shape, dtype=np.float32)
            token[0, 0, 4] = 1
            feeds[value.name] = token
        elif value.name == "position_ids":
            feeds[value.name] = np.array([[257]], dtype=np.int32)
        else:
            feeds[value.name] = rng.normal(0, 0.02, shape).astype(np.float32)
    return feeds


def validate(reference: Path, candidate: Path) -> dict:
    reference_session = ort.InferenceSession(
        str(reference), providers=["CPUExecutionProvider"]
    )
    candidate_session = ort.InferenceSession(
        str(candidate), providers=["CPUExecutionProvider"]
    )
    feeds = make_inputs(reference_session)
    expected = reference_session.run(None, feeds)
    actual = candidate_session.run(None, feeds)
    differences = [
        float(np.max(np.abs(expected_value - actual_value)))
        for expected_value, actual_value in zip(expected, actual)
    ]
    expected_token = int(expected[0][0, -1].argmax())
    actual_token = int(actual[0][0, -1].argmax())
    return {
        "max_abs": max(differences),
        "logits_max_abs": differences[0],
        "cache_max_abs": max(differences[1:]),
        "reference_top_token": expected_token,
        "candidate_top_token": actual_token,
        "top_token_matches": expected_token == actual_token,
    }


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    report = {"format": "symmetric per-output-channel INT8 QDQ", "graphs": {}}
    for name in GRAPH_NAMES:
        source = OUTPUT_DIR / f"{name}_fp32.onnx"
        destination = OUTPUT_DIR / f"{name}_qdq_int8.onnx"
        if not source.exists():
            raise SystemExit(f"Missing {source}. Run export_decoder_fp32.py first.")
        print(f"Rewriting {source.name} -> {destination.name}", flush=True)
        metadata = rewrite_graph(source, destination)
        print(f"Validating {destination.name} on ONNX Runtime CPU", flush=True)
        metadata["cpu_parity"] = validate(source, destination)
        report["graphs"][name] = metadata

    report_path = REPORT_DIR / "decoder-qdq-int8-report.json"
    report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {report_path}")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()