| """Convert, gate, compare and package the independent V2 KV-cache experiment.""" |
| import argparse |
| import json |
| import shutil |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "src")) |
| import numpy as np |
| import torch |
| from vimeml.deployment import kv_cache as kv, v2 |
| from vimeml.deployment.bundle import fresh_directory, read_json, write_json |
| from vimeml.deployment.validation import numeric_error |
| from vimeml.training.data import file_sha |
| from vimeml.training.evaluate_ime import common_prefix_length |
|
|
|
|
| def align(bundle, directory, reference, output, baseline=None): |
| runtime = kv.Runtime(bundle, directory) |
| old = v2.CoreMLLM(bundle, baseline) if baseline else None |
| errors, invariants, saved = [], [], {} |
| info = read_json(reference / "reference.json") |
| if info.get("format") != "vimeml_v2_mac_preparation_reference_v1" or info["checkpoint_step"] != runtime.metadata["checkpoint_step"]: |
| raise ValueError("Invalid frozen reference identity.") |
| if old and any(old.metadata[key] != runtime.metadata[key] for key in ("checkpoint_sha256", "tokenizer_sha256", "model")): |
| raise ValueError("Stateless baseline model/tokenizer identity differs.") |
| with np.load(reference / "logits.npz", allow_pickle=False) as arrays: |
| for name, case in info["cases"].items(): |
| ids = arrays[name + "_input_ids"][0] |
| expected = old.model(torch.tensor([ids.tolist()])).numpy() if old else arrays[name + "_logits"] |
| patterns = [[len(ids)], [1] * len(ids)] |
| if len(ids) > 3: patterns.append([3, len(ids)-3]) |
| for chunks in patterns: |
| parts, past, state = [], 0, None |
| for size in chunks: |
| logits, state = runtime.step(ids[past:past+size], state, past) |
| parts.append(logits); past += size |
| actual = np.concatenate(parts, axis=1) |
| valid = case["valid_prefix"] |
| error = numeric_error(expected[:, :valid], actual[:, :valid], 3e-4, 3e-4) |
| errors.append({"case": name, "chunks": chunks, **error}) |
| if chunks == [len(ids)]: saved[name] = actual |
| last, _ = runtime.step(ids, last_only=True) |
| errors.append({"case": name, "last_only": True, **numeric_error(expected[:, -1:], last, 3e-4, 3e-4)}) |
| for name in ("right_pad", "changed_future"): |
| invariants.append({"case": name, **numeric_error(saved["short"], saved[name][:, :16], 3e-4, 3e-4)}) |
| |
| |
| prefix = arrays["short_input_ids"][0][:7] |
| _, state = runtime.step(prefix, last_only=True) |
| original = tuple(a.copy() for a in state) |
| for suffix in ([9], [47, 81, 102], [9]): |
| actual, _ = runtime.step(suffix, state, len(prefix)) |
| expected = (old.model(torch.tensor([[*prefix.tolist(), *suffix]])).numpy() if old else |
| v2.BundleLM(bundle).model(torch.tensor([[*prefix.tolist(), *suffix]])).detach().numpy())[:, len(prefix):] |
| errors.append({"branch": suffix, **numeric_error(expected, actual, 3e-4, 3e-4)}) |
| immutable = all(np.array_equal(a, b) for a, b in zip(state, original)) |
| dirty = tuple(a.copy() for a in state) |
| for a in dirty: a[:, :, :, len(prefix):, :] = 999 |
| a, _ = runtime.step([9], state, len(prefix)) |
| b, _ = runtime.step([9], dirty, len(prefix)) |
| invariants.append({"case": "unused_tail", **numeric_error(a, b, 0, 0)}) |
| report = {"format": "vimeml_kv_alignment_v1", "passed": immutable and all(e["outside_tolerance"] == 0 for e in errors+invariants), |
| "coreml_manifest_sha256": runtime.metadata["coreml_manifest_sha256"], "model": runtime.metadata, |
| "comparison": "same-weight stateless Core ML" if old else "frozen Windows FP32", |
| "baseline_manifest_sha256": old.metadata["coreml_manifest_sha256"] if old else None, |
| "tolerances": {"atol": 3e-4, "rtol": 3e-4}, "immutable_branch_inputs": immutable, |
| "errors": errors, "invariants": invariants} |
| fresh_directory(output); write_json(output / "alignment.json", report) |
| return report |
|
|
|
|
| class CachedForward: |
| """Evaluation adapter: share exact batch token prefixes within one call.""" |
| def __init__(self, runtime): |
| self.runtime, self.config = runtime, runtime.config |
|
|
| def __call__(self, inputs): |
| ids = inputs.numpy().astype(np.int32) |
| common = common_prefix_length(ids.tolist()) |
| if common: |
| prefix, state = self.runtime.step(ids[0, :common]) |
| result = [] |
| for row in ids: |
| if common == len(row): logits = prefix |
| elif common: |
| suffix, _ = self.runtime.step(row[common:], state, common) |
| logits = np.concatenate((prefix, suffix), axis=1) |
| else: logits, _ = self.runtime.step(row) |
| result.append(logits) |
| return torch.from_numpy(np.concatenate(result, axis=0)) |
|
|
|
|
| class CachedLM(v2.BundleLM): |
| def __init__(self, bundle, directory): |
| super().__init__(bundle) |
| self.runtime = kv.Runtime(bundle, directory) |
| self.model, self.metadata = CachedForward(self.runtime), self.runtime.metadata |
| self._generating = False |
| self._ids = self._state = None |
|
|
| @torch.inference_mode() |
| def next_logits(self, ids): |
| if self._generating and self._ids and ids[:len(self._ids)] == self._ids and len(ids) > len(self._ids): |
| logits, state = self.runtime.step(ids[len(self._ids):], self._state, len(self._ids), last_only=True) |
| else: |
| logits, state = self.runtime.step(ids, last_only=True) |
| if self._generating: self._ids, self._state = list(ids), state |
| return torch.from_numpy(logits[0,-1]) |
|
|
| def generate(self, *args, **kwargs): |
| self._generating, self._ids, self._state = True, None, None |
| try: return super().generate(*args, **kwargs) |
| finally: self._generating, self._ids, self._state = False, None, None |
|
|
|
|
| def timing(work, repeats=20): |
| for _ in range(3): work() |
| times = [] |
| for _ in range(repeats): |
| start = time.perf_counter(); work(); times.append((time.perf_counter()-start)*1000) |
| return {"n": len(times), "p50_ms": float(np.median(times)), "p95_ms": float(np.percentile(times, 95)), "max_ms": max(times)} |
|
|
|
|
| def benchmark(bundle, directory, baseline, output): |
| runtime, old = kv.Runtime(bundle, directory), v2.CoreMLLM(bundle, baseline) |
| rng = np.random.default_rng(716) |
| rows = [] |
| for length in (8, 32, 64, 120): |
| ids = rng.integers(4, 16384, length).tolist(); ids[0] = 2 |
| _, state = runtime.step(ids[:-1], last_only=True) |
| rows.append({"prefix_tokens": length-1, |
| "stateless": timing(lambda: old.model(torch.tensor([ids]))), |
| "cached_decode": timing(lambda: runtime.step(ids[-1:], state, length-1, last_only=True)), |
| "prefill_last_only": timing(lambda: runtime.step(ids, last_only=True))}) |
| def generated(cached, ids, tokens=8): |
| sequence = list(ids) |
| if cached: logits, state = runtime.step(sequence, last_only=True) |
| else: logits = old.model(torch.tensor([sequence])).numpy()[:, -1:] |
| for i in range(tokens): |
| token = int(logits[0,-1].argmax()); sequence.append(token) |
| if i == tokens-1: break |
| if cached: logits, state = runtime.step([token], state, len(sequence)-1, last_only=True) |
| else: logits = old.model(torch.tensor([sequence])).numpy()[:, -1:] |
| return sequence |
| generation = [] |
| for length in (8, 32, 64): |
| ids = [2, *rng.integers(4, 16384, length-1).tolist()] |
| generation.append({"prefix_tokens": length, "new_tokens": 8, |
| "tokens_equal": generated(True,ids) == generated(False,ids), |
| "cached_total": timing(lambda: generated(True,ids), 10), |
| "stateless_total": timing(lambda: generated(False,ids), 10)}) |
| fresh_directory(output) |
| write_json(output / "benchmark.json", {"scope": "Mac CPU_ONLY Python/Core ML; includes Python and array-validation overhead, not iPhone or keyboard extension", |
| "model": runtime.metadata, "decode": rows, "generation": generation, |
| "cache_bytes_per_snapshot": runtime.manifest["cache_bytes_per_snapshot"]}) |
|
|
|
|
| def package(bundle, source, compiled, comparison, output): |
| m = kv.verify_package(source) |
| gate = read_json(comparison) |
| if m["kind"] != "linear8_fp32_compute" or not (compiled / "coremldata.bin").is_file(): |
| raise ValueError("Expected compiled INT8 KV resources.") |
| if (gate.get("format") != "vimeml_kv_alignment_v1" or not gate.get("passed") |
| or gate.get("comparison") != "same-weight stateless Core ML" |
| or gate["coreml_manifest_sha256"] != file_sha(source / "manifest.json")): |
| raise ValueError("Passed same-weight stateless comparison required before packaging.") |
| fresh_directory(output) |
| resources = output / "Resources"; resources.mkdir() |
| shutil.copytree(compiled, resources / "TinyJapaneseV21KVINT8.mlmodelc") |
| |
| |
| write_json(resources / "VimeLMManifestV21KV.json", {"format": "vime_ios_lm_v2_kv_v1", |
| "architecture": v2.ARCHITECTURE, "model_version": "2.1-extend5-step40000-int8-b32-kv-v1", |
| "minimum_ios": 18, "vocab_size": 16384, "context_length": 128, "compute_units": "CPU_ONLY", |
| "special_ids": {"pad": 0, "unk": 1, "bos": 2, "eos": 3}, |
| "tokenizer_sha256": m["bundle"]["tokenizer_sha256"], "checkpoint_sha256": m["bundle"]["checkpoint_sha256"], |
| "cache_shape": m["cache_shape"], "cache_dtype": "float32", |
| "source_manifest_sha256": file_sha(source / "manifest.json"), |
| "stateless_alignment_sha256": file_sha(comparison), |
| "compiled_files_sha256": {name: item["sha256"] for name,item in v2.inventory(compiled,hashes=True).items()}, |
| "policy": "Experimental opt-in; explicit FP32 request-local immutable caches; baseline remains default."}) |
| write_json(output / "manifest.json", {"format": "vimeml_kv_client_resources_v1", "model": m, |
| "files": v2.inventory(output)}) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--threads", type=int, default=4) |
| stages = p.add_subparsers(dest="stage", required=True) |
| for name in ("convert", "align", "compress", "evaluate", "samples", "benchmark", "package"): |
| s = stages.add_parser(name); s.add_argument("--output", type=Path, required=True) |
| if name != "compress": s.add_argument("--bundle", type=Path, required=True) |
| if name in ("align", "evaluate", "samples", "benchmark"): s.add_argument("--model", type=Path, required=True) |
| if name in ("compress", "package"): s.add_argument("--source", type=Path, required=True) |
| if name == "compress": s.add_argument("--alignment", type=Path, required=True) |
| if name == "align": |
| s.add_argument("--reference", type=Path, required=True); s.add_argument("--baseline", type=Path) |
| if name in ("evaluate", "samples", "benchmark"): s.add_argument("--baseline", type=Path, required=True) |
| if name == "evaluate": s.add_argument("--benchmark", type=Path, required=True) |
| if name == "package": |
| s.add_argument("--compiled", type=Path, required=True) |
| s.add_argument("--comparison", type=Path, required=True) |
| args = p.parse_args() |
| if args.output.exists() or args.threads < 1: p.error("Fresh output and positive thread count required.") |
| torch.set_num_threads(args.threads) |
| if args.stage == "convert": kv.convert(args.bundle,args.output) |
| elif args.stage == "compress": kv.compress(args.source,args.output,args.alignment) |
| elif args.stage == "align": |
| report = align(args.bundle,args.model,args.reference,args.output,args.baseline) |
| print(json.dumps({"passed": report["passed"], "max_abs": max(e["max_abs"] for e in report["errors"])}, indent=2)) |
| if not report["passed"]: raise SystemExit("Alignment failed; evidence preserved.") |
| elif args.stage == "benchmark": benchmark(args.bundle,args.model,args.baseline,args.output) |
| elif args.stage == "package": package(args.bundle,args.source,args.compiled,args.comparison,args.output) |
| else: |
| lm = CachedLM(args.bundle,args.model) |
| if args.stage == "evaluate": v2.evaluate(lm,args.benchmark,args.baseline,args.output) |
| else: v2.samples(lm,args.baseline,args.output) |
| print(f"Output: {args.output.resolve()}") |
|
|
|
|
| if __name__ == "__main__": main() |
|
|
