#!/usr/bin/env python3 """End-to-end greedy transcript parity for the fixed 512px Bina CoreML runtime.""" from __future__ import annotations import argparse import hashlib import importlib.util import json import time from pathlib import Path import coremltools as ct import numpy as np import torch from PIL import Image from qwen_vl_utils import process_vision_info from transformers import AutoProcessor EOS_TOKEN_ID = 2 PAD_TOKEN_ID = 0 def load_runtime_module(path: Path): spec = importlib.util.spec_from_file_location("bina_coreml_runtime_export", path) if spec is None or spec.loader is None: raise RuntimeError(f"Could not import runtime module from {path}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def letterbox(image: Image.Image, size: int) -> Image.Image: source = image.convert("RGB") source.thumbnail((size, size), Image.Resampling.LANCZOS) canvas = Image.new("RGB", (size, size), "white") canvas.paste(source, ((size - source.width) // 2, (size - source.height) // 2)) return canvas def build_inputs(processor, image: Image.Image, prompt: str): messages = [ { "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": prompt}, ], } ] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) image_inputs, video_inputs = process_vision_info(messages) return processor( text=[text], images=image_inputs, videos=video_inputs, return_tensors="pt", ) def compute_units(name: str): return { "cpu_only": ct.ComputeUnit.CPU_ONLY, "cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU, "all": ct.ComputeUnit.ALL, }[name] def prefix_match(left: list[int], right: list[int]) -> tuple[int, dict | None]: for index, (actual, expected) in enumerate(zip(left, right)): if actual != expected: return index, {"index": index, "coreml": actual, "reference": expected} matched = min(len(left), len(right)) if len(left) != len(right): return matched, { "index": matched, "coreml": left[matched] if matched < len(left) else None, "reference": right[matched] if matched < len(right) else None, } return matched, None def coreml_generate( rt, model, sample, vision_ml, prefill_ml, decode_ml, max_cache_length: int, max_new_tokens: int, ): started = time.perf_counter() pixel_values = sample["pixel_values"].detach().cpu().numpy().astype(np.float32) vision_outputs = vision_ml.predict({"pixel_values": pixel_values}) image_embeds = torch.from_numpy(vision_outputs["image_embeds"]) _, prefill_example = rt.build_prefill_example_from_image_embeds(model, sample, image_embeds) if int(prefill_example[0].shape[1]) >= max_cache_length: raise ValueError( f"prompt length {prefill_example[0].shape[1]} exceeds cache {max_cache_length}" ) prefill_feed = { spec.name: tensor.detach().cpu().to(torch.float32).numpy() for spec, tensor in zip(rt.prefill_input_specs(prefill_example), prefill_example, strict=True) } prefill_outputs = prefill_ml.predict(prefill_feed) current_token = int(prefill_outputs["logits"][:, -1, :].argmax(axis=-1)[0]) full_keys = [ torch.from_numpy(prefill_outputs[f"full_key_{i}"]) for i in range(rt.SuryaCoreMLPrefillFlat.full_layers) ] full_values = [ torch.from_numpy(prefill_outputs[f"full_value_{i}"]) for i in range(rt.SuryaCoreMLPrefillFlat.full_layers) ] conv_states = [ torch.from_numpy(prefill_outputs[f"conv_state_{i}"]) for i in range(rt.SuryaCoreMLPrefillFlat.linear_layers) ] recurrent_states = [ torch.from_numpy(prefill_outputs[f"recurrent_state_{i}"]) for i in range(rt.SuryaCoreMLPrefillFlat.linear_layers) ] generated = [current_token] cache_len = int(sample["input_ids"].shape[1]) rope_delta = rt.sample_rope_delta(model, sample) stop_reason = "eos" if current_token == EOS_TOKEN_ID else "max_new_tokens" for _ in range(max_new_tokens - 1): if current_token == EOS_TOKEN_ID: stop_reason = "eos" break if cache_len >= max_cache_length: stop_reason = "cache_limit" break token = torch.tensor([[current_token]], dtype=torch.long) with torch.no_grad(): inputs_embeds = model.model.language_model.embed_tokens(token) position_ids = torch.full((3, 1, 1), cache_len + rope_delta, dtype=torch.long) cos, sin = model.model.language_model.rotary_emb(inputs_embeds, position_ids) attention_mask = rt.decode_attention_mask(cache_len, max_cache_length) example = (inputs_embeds, cos, sin, attention_mask) + tuple( full_keys + full_values + conv_states + recurrent_states ) feed = { spec.name: tensor.detach().cpu().to(torch.float32).numpy() for spec, tensor in zip(rt.decode_input_specs(example), example, strict=True) } outputs = decode_ml.predict(feed) next_token = int(outputs["logits"][:, -1, :].argmax(axis=-1)[0]) for index in range(rt.SuryaCoreMLDecodeStepFlat.full_layers): full_keys[index][:, :, cache_len : cache_len + 1, :] = torch.from_numpy( outputs[f"new_full_key_{index}"] ) full_values[index][:, :, cache_len : cache_len + 1, :] = torch.from_numpy( outputs[f"new_full_value_{index}"] ) conv_states = [ torch.from_numpy(outputs[f"new_conv_state_{i}"]) for i in range(rt.SuryaCoreMLDecodeStepFlat.linear_layers) ] recurrent_states = [ torch.from_numpy(outputs[f"new_recurrent_state_{i}"]) for i in range(rt.SuryaCoreMLDecodeStepFlat.linear_layers) ] cache_len += 1 current_token = next_token generated.append(next_token) return { "tokens": generated, "stop_reason": stop_reason, "prompt_tokens": int(sample["input_ids"].shape[1]), "mrope_position_delta": rope_delta, "seconds": time.perf_counter() - started, } def reference_generate(model, sample, max_new_tokens: int): started = time.perf_counter() model.model.rope_deltas = None with torch.inference_mode(): output = model.generate( **sample, do_sample=False, max_new_tokens=max_new_tokens, eos_token_id=EOS_TOKEN_ID, pad_token_id=PAD_TOKEN_ID, use_cache=True, ) prompt_length = int(sample["input_ids"].shape[1]) tokens = output[0, prompt_length:].detach().cpu().tolist() return { "tokens": [int(token) for token in tokens], "stop_reason": "eos" if tokens and tokens[-1] == EOS_TOKEN_ID else "max_new_tokens", "seconds": time.perf_counter() - started, } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model-id", required=True) parser.add_argument("--runtime-script", type=Path, required=True) parser.add_argument("--vision-package", type=Path, required=True) parser.add_argument("--prefill-package", type=Path, required=True) parser.add_argument("--decode-package", type=Path, required=True) parser.add_argument("--image", type=Path, action="append", required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--canvas-size", type=int, default=512) parser.add_argument("--prompt-tokens", type=int, default=416) parser.add_argument("--max-cache-length", type=int, default=1024) parser.add_argument("--max-new-tokens", type=int, default=512) parser.add_argument( "--reference-dtype", choices=["bfloat16", "float32"], default="bfloat16", ) parser.add_argument( "--compute-units", choices=["cpu_only", "cpu_and_gpu", "all"], default="cpu_only", ) args = parser.parse_args() rt = load_runtime_module(args.runtime_script.expanduser().resolve()) dtype = torch.bfloat16 if args.reference_dtype == "bfloat16" else torch.float32 processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True) model = rt.load_model(args.model_id, dtype) units = compute_units(args.compute_units) vision_ml = ct.models.MLModel(str(args.vision_package.expanduser().resolve()), compute_units=units) prefill_ml = ct.models.MLModel(str(args.prefill_package.expanduser().resolve()), compute_units=units) decode_ml = ct.models.MLModel(str(args.decode_package.expanduser().resolve()), compute_units=units) rows = [] for image_path in args.image: image_path = image_path.expanduser().resolve() original_bytes = image_path.read_bytes() with Image.open(image_path) as opened: source_size = list(opened.size) fixed = letterbox(opened, args.canvas_size) sample = build_inputs(processor, fixed, rt.PROMPT) if int(sample["input_ids"].shape[1]) != args.prompt_tokens: raise ValueError( f"{image_path.name}: expected fixed prompt length {args.prompt_tokens}, " f"got {sample['input_ids'].shape[1]}" ) reference = reference_generate(model, sample, args.max_new_tokens) coreml = coreml_generate( rt, model, sample, vision_ml, prefill_ml, decode_ml, args.max_cache_length, args.max_new_tokens, ) matched, mismatch = prefix_match(coreml["tokens"], reference["tokens"]) coreml_text = processor.tokenizer.decode( coreml["tokens"], skip_special_tokens=True, clean_up_tokenization_spaces=False ) reference_text = processor.tokenizer.decode( reference["tokens"], skip_special_tokens=True, clean_up_tokenization_spaces=False ) row = { "image": image_path.name, "source_sha256": hashlib.sha256(original_bytes).hexdigest(), "source_size": source_size, "canvas_size": [args.canvas_size, args.canvas_size], "prompt_tokens": coreml["prompt_tokens"], "mrope_position_delta": coreml["mrope_position_delta"], "coreml_tokens": coreml["tokens"], "reference_tokens": reference["tokens"], "coreml_stop_reason": coreml["stop_reason"], "reference_stop_reason": reference["stop_reason"], "matched_prefix_tokens": matched, "token_exact": mismatch is None, "mismatch": mismatch, "coreml_text": coreml_text, "reference_text": reference_text, "transcript_exact": coreml_text == reference_text, "coreml_seconds": coreml["seconds"], "reference_seconds": reference["seconds"], } print(json.dumps(row, ensure_ascii=False), flush=True) rows.append(row) receipt = { "model_id": args.model_id, "source_precision": "bfloat16", "reference_execution_dtype": args.reference_dtype, "coreml_compute_precision": { "vision": "float32", "prefill": "float16", "decode": "float16", }, "compute_units": args.compute_units, "canvas_size": args.canvas_size, "max_cache_length": args.max_cache_length, "eos_token_id": EOS_TOKEN_ID, "tested": len(rows), "token_exact": sum(row["token_exact"] for row in rows), "transcript_exact": sum(row["transcript_exact"] for row in rows), "all_eos": all( row["coreml_stop_reason"] == "eos" and row["reference_stop_reason"] == "eos" for row in rows ), "rows": rows, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print(json.dumps(receipt, ensure_ascii=False, indent=2), flush=True) if ( receipt["token_exact"] != receipt["tested"] or receipt["transcript_exact"] != receipt["tested"] or not receipt["all_eos"] ): raise SystemExit(1) if __name__ == "__main__": main()