export_decoder_unified.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
import torch

from export_decoder_fp32 import (
    HEAD_DIM,
    HIDDEN_SIZE,
    NUM_KV_HEADS,
    NUM_LAYERS,
    OPSET,
    OUTPUT_DIR,
    REPORT_DIR,
    VISION_TOKENS,
    VOCAB_SIZE,
    DecoderBase,
    compare_outputs,
    flatten_cache,
    inspect_graph,
    load_model,
    make_one_hot,
    make_past,
    output_names,
    past_names,
)
from export_decoder_qdq_int8 import rewrite_graph


class DecoderUnified(DecoderBase):
    def forward(
        self,
        vision_embeds: torch.Tensor,
        token_one_hot: torch.Tensor,
        position_ids: torch.Tensor,
        past_k0: torch.Tensor,
        past_k1: torch.Tensor,
        past_k2: torch.Tensor,
        past_k3: torch.Tensor,
        past_k4: torch.Tensor,
        past_k5: torch.Tensor,
        past_v0: torch.Tensor,
        past_v1: torch.Tensor,
        past_v2: torch.Tensor,
        past_v3: torch.Tensor,
        past_v4: torch.Tensor,
        past_v5: torch.Tensor,
    ):
        keys = (past_k0, past_k1, past_k2, past_k3, past_k4, past_k5)
        values = (past_v0, past_v1, past_v2, past_v3, past_v4, past_v5)
        token_embed = torch.matmul(token_one_hot, self.lm_head.weight)
        inputs_embeds = torch.cat((vision_embeds, token_embed), dim=1)
        outputs = self.decoder(
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            past_key_values=tuple(zip(keys, values)),
            use_cache=True,
            return_dict=True,
        )
        final_hidden_state = outputs.last_hidden_state[:, -1:, :]
        return (self.project_logits(final_hidden_state),) + flatten_cache(
            outputs.past_key_values
        )


def empty_vision() -> torch.Tensor:
    return torch.empty(1, 0, HIDDEN_SIZE, dtype=torch.float32)


def empty_past() -> tuple[torch.Tensor, ...]:
    return tuple(
        torch.empty(1, NUM_KV_HEADS, 0, HEAD_DIM, dtype=torch.float32)
        for _ in range(NUM_LAYERS * 2)
    )


def export_graph(wrapper: DecoderUnified, destination: Path) -> None:
    torch.manual_seed(31)
    vision = torch.randn(1, VISION_TOKENS, HIDDEN_SIZE, dtype=torch.float32)
    token = make_one_hot(1)
    past = make_past(1)
    positions = torch.arange(1, VISION_TOKENS + 2, dtype=torch.int32).unsqueeze(0)
    dynamic_axes = {
        "vision_embeds": {1: "vision_len"},
        "position_ids": {1: "input_len"},
        **{name: {2: "past_len"} for name in past_names()},
        **{name: {2: "total_len"} for name in output_names()[1:]},
    }
    torch.onnx.export(
        wrapper,
        (vision, token, positions, *past),
        destination,
        input_names=["vision_embeds", "token_one_hot", "position_ids", *past_names()],
        output_names=output_names(),
        dynamic_axes=dynamic_axes,
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
    )


def feeds(values: tuple[torch.Tensor, ...]) -> dict[str, np.ndarray]:
    return {name: value.numpy() for name, value in zip(
        ["vision_embeds", "token_one_hot", "position_ids", *past_names()], values
    )}


def validate_mode(
    wrapper: DecoderUnified,
    session: ort.InferenceSession,
    label: str,
    values: tuple[torch.Tensor, ...],
) -> dict:
    with torch.inference_mode():
        expected = wrapper(*values)
    actual = session.run(None, feeds(values))
    return compare_outputs(label, expected, actual)


def validate_graph(wrapper: DecoderUnified, path: Path) -> dict:
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    torch.manual_seed(37)
    prefill = (
        torch.randn(1, VISION_TOKENS, HIDDEN_SIZE, dtype=torch.float32),
        make_one_hot(1),
        torch.arange(VISION_TOKENS + 1, dtype=torch.int32).unsqueeze(0),
        *empty_past(),
    )
    step = (
        empty_vision(),
        make_one_hot(4),
        torch.tensor([[VISION_TOKENS + 1]], dtype=torch.int32),
        *make_past(VISION_TOKENS + 1),
    )
    return {
        "prefill_zero_past": validate_mode(wrapper, session, "unified-prefill", prefill),
        "step_empty_vision": validate_mode(wrapper, session, "unified-step", step),
    }


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(0)
    torch.set_grad_enabled(False)
    wrapper = DecoderUnified(load_model()).eval()
    fp32_path = OUTPUT_DIR / "decoder_unified_fp32.onnx"
    qdq_path = OUTPUT_DIR / "decoder_unified_qdq_int8.onnx"
    print(f"Exporting {fp32_path}", flush=True)
    export_graph(wrapper, fp32_path)
    fp32_parity = validate_graph(wrapper, fp32_path)
    print(f"Quantizing {qdq_path}", flush=True)
    qdq_metadata = rewrite_graph(fp32_path, qdq_path)
    qdq_parity = validate_graph(wrapper, qdq_path)
    report = {
        "description": (
            "One complete six-layer Baberu decoder graph shared by prefill and token "
            "steps; weights, dimensions, vocabulary, and generation logic are unchanged"
        ),
        "fp32": {"graph": inspect_graph(fp32_path), "parity": fp32_parity},
        "qdq_int8": {
            "graph": inspect_graph(qdq_path),
            "quantization": qdq_metadata,
            "parity": qdq_parity,
        },
    }
    destination = REPORT_DIR / "decoder-unified-report.json"
    destination.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()