| """V2 incremental attention with explicit, immutable FP32 cache snapshots. |
| |
| Caches belong to one request. Reuse only exact token prefixes, with absolute |
| learned positions; never slide, persist across requests, or mutate a branch. |
| """ |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from vimeml.deployment import v2 |
| from vimeml.deployment.bundle import fresh_directory, read_json, write_json |
| from vimeml.deployment.coreml import coremltools, inspect_spec |
| from vimeml.training.data import file_sha |
|
|
| FORMAT = "vimeml_coreml_v2_kv_v1" |
|
|
|
|
| class CacheGraph(nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.model = model.eval() |
| c = model.config |
| self.width, self.heads, self.head_dim = c.d_model, c.n_heads, c.d_model // c.n_heads |
| mask = torch.full((c.context_length, c.context_length), -float("inf")) |
| self.register_buffer("causal_mask", torch.triu(mask, diagonal=1)) |
|
|
| def forward(self, input_ids, cache_length, key_cache, value_cache): |
| length = input_ids.shape[1] |
| past = cache_length[0] |
| positions = torch.arange(length, device=input_ids.device, dtype=torch.int32) + past |
| hidden = self.model.token_embedding(input_ids) + self.model.position_embedding(positions) |
| |
| mask = self.causal_mask[positions.long()] |
| keys, values = [], [] |
| for i, block in enumerate(self.model.blocks): |
| normalized = block.attention_norm(hidden) |
| qkv = block.attention.qkv(normalized).reshape(1, length, 3, self.heads, self.head_dim) |
| query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) |
| indices = (positions.long(),) |
| key = key_cache[i].permute(2, 0, 1, 3).index_put(indices, key.permute(2, 0, 1, 3)).permute(1, 2, 0, 3) |
| value = value_cache[i].permute(2, 0, 1, 3).index_put(indices, value.permute(2, 0, 1, 3)).permute(1, 2, 0, 3) |
| attention = torch.softmax(torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim) + mask, dim=-1) |
| attended = torch.matmul(attention, value).transpose(1, 2).reshape(1, length, self.width) |
| hidden = hidden + block.attention.projection(attended) |
| hidden = hidden + block.ffn(block.ffn_norm(hidden)) |
| keys.append(key) |
| values.append(value) |
| |
| |
| logits = self.model.lm_head(self.model.final_norm(hidden)) |
| return logits, torch.stack(keys), torch.stack(values) |
|
|
|
|
| def cache_shape(config): |
| return (config.n_layers, 1, config.n_heads, config.context_length, config.d_model // config.n_heads) |
|
|
|
|
| def empty_cache(config): |
| return tuple(np.zeros(cache_shape(config), dtype=np.float32) for _ in range(2)) |
|
|
|
|
| @torch.inference_mode() |
| def trace(model): |
| graph = CacheGraph(model).eval() |
| keys, values = (torch.from_numpy(a) for a in empty_cache(model.config)) |
| traced = torch.jit.trace(graph, (torch.tensor([[2, 4, 7]], dtype=torch.int32), |
| torch.tensor([0], dtype=torch.int32), keys, values)) |
| |
| |
| generator = torch.Generator().manual_seed(716) |
| ids = torch.randint(model.config.vocab_size, (1, model.config.context_length), generator=generator, dtype=torch.int32) |
| expected = model(ids) |
| for fn in (graph, traced): |
| state = (keys, values) |
| start = 0 |
| for size in (1, 3, model.config.context_length - 4): |
| actual, k, v = fn(ids[:, start:start+size], torch.tensor([start], dtype=torch.int32), |
| *state) |
| torch.testing.assert_close(actual, expected[:, start:start+size], atol=3e-5, rtol=3e-4) |
| state = k, v |
| start += size |
| actual, _, _ = fn(ids, torch.tensor([0], dtype=torch.int32), keys, values) |
| torch.testing.assert_close(actual, expected, atol=3e-5, rtol=3e-4) |
| return traced |
|
|
|
|
| def finish(output, model, metadata): |
| fresh_directory(output) |
| model.save(str(Path(output) / "model.mlpackage")) |
| files = v2.inventory(Path(output) / "model.mlpackage") |
| write_json(Path(output) / "manifest.json", {"format": FORMAT, "status": "complete", |
| "architecture": v2.ARCHITECTURE, **metadata, "environment": v2.environment(), |
| "spec": inspect_spec(model), "package_files": files, |
| "package_bytes": sum(item["bytes"] for item in files.values())}) |
|
|
|
|
| def verify_package(directory): |
| directory = Path(directory) |
| m = read_json(directory / "manifest.json") |
| if (m.get("format") != FORMAT or m.get("status") != "complete" |
| or m.get("architecture") != v2.ARCHITECTURE or v2.inventory(directory / "model.mlpackage") != m["package_files"]): |
| raise ValueError("Invalid KV package or changed file inventory.") |
| return m |
|
|
|
|
| def convert(bundle, output): |
| if Path(output).exists(): |
| raise ValueError("Output exists; choose a new experiment.") |
| lm, ct = v2.BundleLM(bundle), coremltools() |
| c = lm.model.config |
| result = ct.convert(trace(lm.model), source="pytorch", convert_to="mlprogram", |
| minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT32, |
| inputs=[ct.TensorType(name="input_ids", shape=(1, ct.RangeDim(1, c.context_length, default=1)), dtype=np.int32), |
| ct.TensorType(name="cache_length", shape=(1,), dtype=np.int32), |
| ct.TensorType(name="key_cache", shape=cache_shape(c), dtype=np.float32), |
| ct.TensorType(name="value_cache", shape=cache_shape(c), dtype=np.float32)], |
| outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="new_key_cache", dtype=np.float32), |
| ct.TensorType(name="new_value_cache", dtype=np.float32)], skip_model_load=True) |
| finish(output, result, {"kind": "fp32", "minimum_ios": 18, "bundle": lm.metadata, |
| "cache_shape": cache_shape(c), "cache_dtype": "float32", "compute_units": "CPU_ONLY", |
| "cache_bytes_per_snapshot": 2 * int(np.prod(cache_shape(c))) * 4, |
| "interface": "batch1; absolute learned positions; 0<=past, 1<=Q, past+Q<=128; logits [1,Q,V] for new tokens only. Explicit immutable K/V snapshots.", |
| "conversion_source_sha256": file_sha(Path(__file__)), |
| "policy": "No training/weight changes. Request-local exact token prefix reuse; no sliding or cross-request retention."}) |
|
|
|
|
| def compress(source, output, alignment): |
| if Path(output).exists(): |
| raise ValueError("Output exists; choose a new experiment.") |
| m, gate = verify_package(source), read_json(alignment) |
| if (m["kind"] != "fp32" or not gate.get("passed") or gate.get("format") != "vimeml_kv_alignment_v1" |
| or gate.get("coreml_manifest_sha256") != file_sha(Path(source) / "manifest.json")): |
| raise ValueError("A passed alignment report for this exact FP32 KV package is required.") |
| ct = coremltools() |
| from coremltools.optimize import coreml as opt |
| model = ct.models.MLModel(str(Path(source) / "model.mlpackage"), skip_model_load=True) |
| config = opt.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_block", block_size=32, weight_threshold=2048) |
| optimization, selection = v2.parameter_selection(opt, model, config) |
| result = opt.linear_quantize_weights(model, config=optimization) |
| finish(output, result, {key: m[key] for key in ("minimum_ios", "bundle", "cache_shape", "cache_dtype", "compute_units", "cache_bytes_per_snapshot", "interface", "policy")} | { |
| "kind": "linear8_fp32_compute", "source_manifest_sha256": file_sha(Path(source) / "manifest.json"), |
| "fp32_alignment_sha256": file_sha(alignment), "compression": {"method": "linear_symmetric", "bits": 8, |
| "block_size": 32, "selection": selection}, "conversion_source_sha256": file_sha(Path(__file__))}) |
|
|
|
|
| class Runtime: |
| def __init__(self, bundle, directory): |
| lm = v2.BundleLM(bundle) |
| self.config, self.metadata = lm.model.config, dict(lm.metadata) |
| self.manifest = verify_package(directory) |
| if self.manifest["bundle"]["bundle_manifest_sha256"] != self.metadata["bundle_manifest_sha256"]: |
| raise ValueError("KV model belongs to another bundle.") |
| if (self.manifest["cache_shape"] != list(cache_shape(self.config)) |
| or self.manifest["cache_dtype"] != "float32" or self.manifest["compute_units"] != "CPU_ONLY"): |
| raise ValueError("KV package cache/compute contract differs.") |
| ct = coremltools() |
| self.model = ct.models.MLModel(str(Path(directory) / "model.mlpackage"), compute_units=ct.ComputeUnit.CPU_ONLY) |
| self.metadata.update(precision=self.manifest["kind"], compute_units="CPU_ONLY", |
| coreml_manifest_sha256=file_sha(Path(directory) / "manifest.json")) |
| self.calls = self.processed_tokens = 0 |
|
|
| def step(self, ids, state=None, past=0, last_only=False): |
| raw = np.asarray(ids) |
| if (not raw.size or not np.issubdtype(raw.dtype, np.integer) |
| or raw.ndim not in (1, 2) or (raw.ndim == 2 and raw.shape[0] != 1) |
| or (raw < 0).any() or (raw >= self.config.vocab_size).any()): |
| raise ValueError("Invalid incremental token tensor.") |
| ids = raw.astype(np.int32).reshape(1, -1) |
| if (not isinstance(past, int) or past < 0 or past + ids.shape[1] > self.config.context_length |
| or (state is None and past != 0)): |
| raise ValueError("Invalid incremental input/cache position.") |
| state = empty_cache(self.config) if state is None else state |
| if len(state) != 2 or any(not isinstance(a, np.ndarray) or a.shape != cache_shape(self.config) |
| or a.dtype != np.float32 or not np.isfinite(a).all() for a in state): |
| raise ValueError("Invalid cache tensor.") |
| result = self.model.predict({"input_ids": ids, "cache_length": np.array([past], np.int32), |
| "key_cache": state[0], "value_cache": state[1]}) |
| self.calls += 1 |
| self.processed_tokens += ids.size |
| |
| |
| logits = np.array(result["logits"], dtype=np.float32, copy=True) |
| caches = tuple(np.array(result[name], dtype=np.float32, copy=True) |
| for name in ("new_key_cache", "new_value_cache")) |
| if (logits.shape != (1, ids.size, self.config.vocab_size) |
| or not np.isfinite(logits).all() |
| or any(a.shape != cache_shape(self.config) or not np.isfinite(a).all() for a in caches)): |
| raise ValueError("Invalid Core ML cache/logits output.") |
| return (logits[:, -1:].copy() if last_only else logits), caches |
|
|
