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()