optimize_decoder_execution.py

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter, defaultdict
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
DEFAULT_SOURCE = ROOT / ".work/models/shared/decoder_unified_gather_qdq_int8.onnx"
DEFAULT_DESTINATION = ROOT / ".work/models/model-opt/decoder_gather_before_dq_int8.onnx"
DEFAULT_FIXED_KV_DESTINATION = ROOT / ".work/models/model-opt/decoder_gather_dq_fixed_kv_int8.onnx"
DEFAULT_REPORT = ROOT / ".work/reports/model-execution-optimization.json"
NUM_LAYERS = 6
VISION_TOKENS = 256
MAX_NEW_TOKENS = 128
# Prefill produces 257 cache entries. The runtime stops before running a decode
# step for token 128, so the largest present cache produced is length 384.
MAX_CACHE_LENGTH = VISION_TOKENS + MAX_NEW_TOKENS
CACHE_NAMES = [f"{kind}_{axis}{layer}" for kind in ("past", "present") for axis in ("k", "v") for layer in range(NUM_LAYERS)]


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 set_shape(value_info: onnx.ValueInfoProto, shape: list[int | str]) -> None:
    dimensions = value_info.type.tensor_type.shape.dim
    del dimensions[:]
    for value in shape:
        dimension = dimensions.add()
        if isinstance(value, int):
            dimension.dim_value = value
        else:
            dimension.dim_param = value


def rewrite_embedding_gather(model: onnx.ModelProto) -> dict:
    graph = model.graph
    initializers = {value.name: value for value in graph.initializer}
    consumers: dict[str, list[onnx.NodeProto]] = defaultdict(list)
    for node in graph.node:
        for name in node.input:
            consumers[name].append(node)

    embedding_dq = None
    embedding_gather = None
    for node in graph.node:
        if node.op_type != "DequantizeLinear" or len(node.input) < 3:
            continue
        quantized = initializers.get(node.input[0])
        if not quantized or list(quantized.dims) != [14630, 512]:
            continue
        matches = [consumer for consumer in consumers[node.output[0]] if consumer.op_type == "Gather"]
        if len(matches) != 1:
            raise RuntimeError(f"Embedding DQ expected one Gather consumer, found {len(matches)}")
        embedding_dq = node
        embedding_gather = matches[0]
        break
    if embedding_dq is None or embedding_gather is None:
        raise RuntimeError("Could not locate quantized 14630x512 embedding Gather")

    original_output = embedding_gather.output[0]
    quantized_output = f"{original_output}_int8"
    embedding_gather.input[0] = embedding_dq.input[0]
    embedding_gather.output[0] = quantized_output
    gathered_dq = helper.make_node(
        "DequantizeLinear",
        [quantized_output, embedding_dq.input[1], embedding_dq.input[2]],
        [original_output],
        name=f"{embedding_dq.name}_after_gather",
        # Gather axis 0 with token_ids rank 2 moves hidden axis 1 to axis 2.
        axis=2,
    )

    rewritten = []
    for node in graph.node:
        if node is embedding_dq:
            continue
        rewritten.append(node)
        if node is embedding_gather:
            rewritten.append(gathered_dq)
    del graph.node[:]
    graph.node.extend(rewritten)
    return {
        "quantized_elements_selected": 14630 * 512,
        "fp32_bytes_avoided_before_gather": 14630 * 512 * 4,
        "gather_output": original_output,
    }


def rewrite_fixed_kv_io(model: onnx.ModelProto) -> dict:
    graph = model.graph
    input_by_name = {value.name: value for value in graph.input}
    output_by_name = {value.name: value for value in graph.output}
    past_names = [f"past_{axis}{layer}" for axis in ("k", "v") for layer in range(NUM_LAYERS)]
    present_names = [f"present_{axis}{layer}" for axis in ("k", "v") for layer in range(NUM_LAYERS)]

    for name in past_names:
        if name not in input_by_name:
            raise RuntimeError(f"Missing cache input {name}")
        set_shape(input_by_name[name], [1, 2, MAX_CACHE_LENGTH, 64])
    for name in present_names:
        if name not in output_by_name:
            raise RuntimeError(f"Missing cache output {name}")
        set_shape(output_by_name[name], [1, 2, MAX_CACHE_LENGTH, 64])

    graph.input.append(helper.make_tensor_value_info("past_length", TensorProto.INT64, [1]))
    graph.initializer.extend(
        [
            numpy_helper.from_array(np.array([0], dtype=np.int64), "fixed_kv_slice_starts"),
            numpy_helper.from_array(np.array([2], dtype=np.int64), "fixed_kv_slice_axes"),
            numpy_helper.from_array(np.array([1], dtype=np.int64), "fixed_kv_slice_steps"),
            numpy_helper.from_array(np.zeros(4, dtype=np.int64), "fixed_kv_pad_begin"),
            numpy_helper.from_array(np.zeros(2, dtype=np.int64), "fixed_kv_pad_end_prefix"),
            numpy_helper.from_array(np.zeros(1, dtype=np.int64), "fixed_kv_pad_end_suffix"),
            numpy_helper.from_array(np.array([MAX_CACHE_LENGTH], dtype=np.int64), "fixed_kv_max_length"),
            numpy_helper.from_array(np.array([2], dtype=np.int64), "fixed_kv_shape_axis"),
        ]
    )

    slice_nodes = []
    sliced_names: dict[str, str] = {}
    for name in past_names:
        sliced_name = f"{name}_valid"
        sliced_names[name] = sliced_name
        slice_nodes.append(
            helper.make_node(
                "Slice",
                [name, "fixed_kv_slice_starts", "past_length", "fixed_kv_slice_axes", "fixed_kv_slice_steps"],
                [sliced_name],
                name=f"fixed_kv_slice_{name}",
            )
        )

    original_nodes = list(graph.node)
    for node in original_nodes:
        for index, name in enumerate(node.input):
            if name in sliced_names:
                node.input[index] = sliced_names[name]

    pad_nodes = []
    for name in present_names:
        producer = next((node for node in original_nodes if name in node.output), None)
        if producer is None:
            raise RuntimeError(f"Missing producer for {name}")
        valid_name = f"{name}_valid"
        producer.output[list(producer.output).index(name)] = valid_name
        for consumer in original_nodes:
            for input_index, input_name in enumerate(consumer.input):
                if input_name == name:
                    consumer.input[input_index] = valid_name
        shape_name = f"{name}_shape"
        length_name = f"{name}_length"
        padding_name = f"{name}_padding"
        pads_name = f"{name}_pads"
        pad_nodes.extend(
            [
                helper.make_node("Shape", [valid_name], [shape_name], name=f"fixed_kv_shape_{name}"),
                helper.make_node(
                    "Gather",
                    [shape_name, "fixed_kv_shape_axis"],
                    [length_name],
                    name=f"fixed_kv_length_{name}",
                    axis=0,
                ),
                helper.make_node(
                    "Sub",
                    ["fixed_kv_max_length", length_name],
                    [padding_name],
                    name=f"fixed_kv_padding_{name}",
                ),
                helper.make_node(
                    "Concat",
                    ["fixed_kv_pad_begin", "fixed_kv_pad_end_prefix", padding_name, "fixed_kv_pad_end_suffix"],
                    [pads_name],
                    name=f"fixed_kv_pads_{name}",
                    axis=0,
                ),
                helper.make_node("Pad", [valid_name, pads_name], [name], name=f"fixed_kv_pad_{name}"),
            ]
        )

    del graph.node[:]
    graph.node.extend(slice_nodes)
    graph.node.extend(original_nodes)
    graph.node.extend(pad_nodes)
    return {
        "max_cache_length": MAX_CACHE_LENGTH,
        "fixed_live_cache_bytes": NUM_LAYERS * 2 * 2 * MAX_CACHE_LENGTH * 64 * 4,
        "dynamic_present_allocation_bytes_across_max_decode": sum(
            NUM_LAYERS * 2 * 2 * length * 64 * 4
            for length in range(VISION_TOKENS + 1, MAX_CACHE_LENGTH + 1)
        ),
    }


def operator_counts(model: onnx.ModelProto) -> dict[str, int]:
    return dict(sorted(Counter(node.op_type for node in model.graph.node).items()))


def source_feeds(session: ort.InferenceSession, past_length: int, *, prefill: bool) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(20260717 + past_length)
    feeds: dict[str, np.ndarray] = {}
    for value in session.get_inputs():
        if value.name == "vision_embeds":
            length = VISION_TOKENS if prefill else 0
            feeds[value.name] = rng.normal(0, 0.2, [1, length, 512]).astype(np.float32)
        elif value.name == "token_ids":
            feeds[value.name] = np.array([[1 if prefill else 4]], dtype=np.int32)
        elif value.name == "position_ids":
            feeds[value.name] = (
                np.arange(VISION_TOKENS + 1, dtype=np.int32)[None, :]
                if prefill
                else np.array([[past_length]], dtype=np.int32)
            )
        else:
            feeds[value.name] = rng.normal(0, 0.02, [1, 2, past_length, 64]).astype(np.float32)
    return feeds


def candidate_feeds(source: dict[str, np.ndarray], past_length: int) -> dict[str, np.ndarray]:
    feeds: dict[str, np.ndarray] = {}
    for name, value in source.items():
        if name.startswith("past_"):
            fixed = np.zeros([1, 2, MAX_CACHE_LENGTH, 64], dtype=np.float32)
            fixed[:, :, :past_length, :] = value
            feeds[name] = fixed
        else:
            feeds[name] = value
    feeds["past_length"] = np.array([past_length], dtype=np.int64)
    return feeds


def validate_cpu(source: Path, candidate: Path, *, fixed_kv: bool) -> dict:
    source_session = ort.InferenceSession(str(source), providers=["CPUExecutionProvider"])
    candidate_session = ort.InferenceSession(str(candidate), providers=["CPUExecutionProvider"])
    checks = {}
    for label, past_length, prefill in (("prefill", 0, True), ("step_257", 257, False), ("step_383", 383, False)):
        source_input = source_feeds(source_session, past_length, prefill=prefill)
        expected = source_session.run(None, source_input)
        actual = candidate_session.run(
            None,
            candidate_feeds(source_input, past_length) if fixed_kv else source_input,
        )
        active_length = (VISION_TOKENS + 1) if prefill else past_length + 1
        differences = [float(np.max(np.abs(expected[0] - actual[0])))]
        padding_nonzero = 0
        for expected_cache, actual_cache in zip(expected[1:], actual[1:]):
            actual_valid = actual_cache[:, :, :active_length, :] if fixed_kv else actual_cache
            differences.append(float(np.max(np.abs(expected_cache - actual_valid))))
            if fixed_kv:
                padding_nonzero += int(np.count_nonzero(actual_cache[:, :, active_length:, :]))
        checks[label] = {
            "max_abs": max(differences),
            "logits_max_abs": differences[0],
            "top_token_source": int(expected[0][0, -1].argmax()),
            "top_token_candidate": int(actual[0][0, -1].argmax()),
            "padding_nonzero": padding_nonzero,
        }
        if checks[label]["max_abs"] != 0 or padding_nonzero != 0:
            raise RuntimeError(f"CPU parity failed for {label}: {checks[label]}")
    return checks


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
    parser.add_argument("--destination", type=Path, default=DEFAULT_DESTINATION)
    parser.add_argument("--fixed-kv-destination", type=Path, default=DEFAULT_FIXED_KV_DESTINATION)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    arguments = parser.parse_args()
    arguments.destination.parent.mkdir(parents=True, exist_ok=True)
    arguments.fixed_kv_destination.parent.mkdir(parents=True, exist_ok=True)
    arguments.report.parent.mkdir(parents=True, exist_ok=True)

    model = onnx.load(arguments.source)
    before = operator_counts(model)
    embedding = rewrite_embedding_gather(model)
    model.producer_name = "vibe-manga-baberu-webgpu-model-execution-opt"
    model.producer_version = "1"
    onnx.checker.check_model(model)
    onnx.save(model, arguments.destination)
    gather_parity = validate_cpu(arguments.source, arguments.destination, fixed_kv=False)

    fixed_kv_model = onnx.load(arguments.destination)
    fixed_kv = rewrite_fixed_kv_io(fixed_kv_model)
    fixed_kv_model.producer_version = "1-fixed-kv-experiment"
    onnx.checker.check_model(fixed_kv_model)
    onnx.save(fixed_kv_model, arguments.fixed_kv_destination)
    fixed_kv_parity = validate_cpu(
        arguments.source,
        arguments.fixed_kv_destination,
        fixed_kv=True,
    )
    report = {
        "source": {
            "path": str(arguments.source.relative_to(ROOT)),
            "bytes": arguments.source.stat().st_size,
            "sha256": sha256(arguments.source),
            "operators": before,
        },
        "gather_optimized": {
            "path": str(arguments.destination.relative_to(ROOT)),
            "bytes": arguments.destination.stat().st_size,
            "sha256": sha256(arguments.destination),
            "operators": operator_counts(model),
        },
        "fixed_kv_experiment": {
            "path": str(arguments.fixed_kv_destination.relative_to(ROOT)),
            "bytes": arguments.fixed_kv_destination.stat().st_size,
            "sha256": sha256(arguments.fixed_kv_destination),
            "operators": operator_counts(fixed_kv_model),
        },
        "capability": {
            "layers": 6,
            "hidden_size": 512,
            "kv_heads": 2,
            "vocabulary": 14630,
            "max_new_tokens": MAX_NEW_TOKENS,
            "weights_requantized": False,
        },
        "embedding_gather_before_dequantize": embedding,
        "fixed_kv_io": fixed_kv,
        "cpu_parity": {
            "gather_optimized": gather_parity,
            "fixed_kv_experiment": fixed_kv_parity,
        },
    }
    arguments.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()