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