Engram-protocol/engram
0
1"""2ENGRAM Protocol — Demo Agent Session3 4 5End-to-end demonstration:6 1. Load model via llama-cpp-python (D1)7 2. Generate with a prompt → measure cold TTFT8 3. Extract KV cache → compress → serialize to .eng9 4. Index in EGR manifold index10 5. Reset model → restore from .eng → measure cached TTFT11 6. Print speedup ratio12 13D6: Target >10x TTFT reduction at 16K context on Llama 3.1 8B.14 Cold baseline: ~1,500-5,000ms. Cached target: <500ms.15 Anything below 4x at 16K is a failure.16"""17 18from __future__ import annotations19 20import argparse21import sys22import time23from pathlib import Path24 25 26def _run_dry_run(args: argparse.Namespace) -> int:27 """Run full pipeline with synthetic tensors — no model file needed."""28 import os29 import tempfile30 31 import torch32 33 from kvcos.core.cache_spec import LLAMA_3_1_8B34 from kvcos.core.serializer import EngramSerializer35 from kvcos.core.types import CompressionMethod, StateExtractionMode36 from kvcos.core.manifold_index import IndexEntry, ManifoldIndex37 from kvcos.core.state_extractor import MARStateExtractor38 from kvcos.storage.local import LocalStorageBackend39 40 spec = LLAMA_3_1_8B41 ctx_len = args.context42 model_name = spec["model_id"]43 44 # ── Synthetic KV tensors ──────────────────────────────────45 torch.manual_seed(42)46 shape = (spec["n_layers"], spec["n_kv_heads"], ctx_len, spec["head_dim"])47 keys = torch.randn(shape, dtype=torch.float16)48 values = torch.randn(shape, dtype=torch.float16)49 50 tensor_mb = keys.numel() * keys.element_size() / 1024 / 102451 52 with tempfile.TemporaryDirectory() as tmp:53 tmp_dir = Path(tmp)54 55 # ── Serialize to .eng ────────────────────────────────56 serializer = EngramSerializer()57 eng_path = tmp_dir / "dry_run.eng"58 59 t0 = time.perf_counter()60 result = serializer.serialize(61 keys=keys, values=values,62 agent_id="dry-run-agent",63 task_description="dry run benchmark",64 model_id=model_name,65 output_path=eng_path,66 compression=CompressionMethod.Q8_0,67 )68 serialize_ms = (time.perf_counter() - t0) * 100069 70 # ── Load back ────────────────────────────────────────71 t0 = time.perf_counter()72 k_out, v_out, meta = serializer.deserialize(eng_path)73 deserialize_ms = (time.perf_counter() - t0) * 100074 75 assert k_out.shape == keys.shape, f"Shape mismatch: {k_out.shape} vs {keys.shape}"76 77 # ── EGR granular timing ──────────────────────────────78 extractor = MARStateExtractor(79 mode=StateExtractionMode.SVD_PROJECT,80 rank=min(160, spec["head_dim"]),81 )82 dim = extractor.output_dim(spec)83 index = ManifoldIndex(dim=dim)84 storage = LocalStorageBackend(data_dir=tmp_dir)85 86 # Index: extract + serialize + store + add87 t0 = time.perf_counter()88 extraction = extractor.extract(keys, spec)89 t_extract = time.perf_counter()90 91 eng2 = tmp_dir / "indexed.eng"92 serializer.serialize(93 keys=keys, values=values,94 agent_id="dry-run-agent",95 task_description="dry run benchmark",96 model_id=model_name,97 output_path=eng2,98 compression=CompressionMethod.Q8_0,99 cache_id="dry-run-001",100 )101 t_serialize = time.perf_counter()102 103 idx_meta = serializer.read_metadata_only(eng2)104 storage.store_file("dry-run-001", eng2, idx_meta)105 t_store = time.perf_counter()106 107 from datetime import datetime, timezone108 entry = IndexEntry(109 cache_id="dry-run-001",110 task_description="dry run benchmark",111 model_id=model_name,112 created_at=datetime.now(timezone.utc).isoformat(),113 context_len=ctx_len,114 l2_norm=extraction.l2_norm,115 )116 index.add(extraction.state_vec, entry)117 t_add = time.perf_counter()118 119 extract_ms = (t_extract - t0) * 1000120 ser_ms = (t_serialize - t_extract) * 1000121 store_ms = (t_store - t_serialize) * 1000122 add_ms = (t_add - t_store) * 1000123 index_ms = (t_add - t0) * 1000124 125 # Retrieve: extract query + search + load126 torch.manual_seed(99)127 query_keys = torch.randn(shape, dtype=torch.float16)128 129 t0 = time.perf_counter()130 q_ext = extractor.extract(query_keys, spec)131 t_qext = time.perf_counter()132 133 results = index.search(q_ext.state_vec, top_k=1)134 t_search = time.perf_counter()135 136 # Load matched engram137 stored_path = storage.get_path("dry-run-001")138 k_loaded, v_loaded, _ = serializer.deserialize(stored_path)139 t_load = time.perf_counter()140 141 q_extract_ms = (t_qext - t0) * 1000142 search_ms = (t_search - t_qext) * 1000143 load_ms = (t_load - t_search) * 1000144 retrieve_ms = (t_load - t0) * 1000145 146 # ── Simulate TTFT estimates ──────────────────────────147 cold_ms = ctx_len * 0.1 # simulated148 cached_ms = deserialize_ms149 egr_overhead = extract_ms + search_ms # overhead added to warm path150 speedup = cold_ms / cached_ms if cached_ms > 0 else float("inf")151 eng_size_mb = os.path.getsize(eng_path) / 1024 / 1024152 153 # ── Output ───────────────────────────────────────────154 sep = "=" * 35155 print(sep)156 print("ENGRAM Protocol \u2014 EGR Demo")157 print(f"Model: {model_name}")158 print(f"Context: {ctx_len} tokens")159 print(sep)160 print(f"Cold TTFT: {cold_ms:.1f}ms (simulated)")161 print(f"Cached TTFT: {cached_ms:.1f}ms (deserialize)")162 print(f"Speedup: {speedup:.1f}x")163 print(f"D6 target: >10x at 16K tokens")164 status = "PASS" if speedup > 10 else "FAIL"165 print(f"Status: {status}")166 print(f"EGR overhead: {egr_overhead:.1f}ms (extract+search)")167 print(f".eng file: {eng_path.name} ({eng_size_mb:.1f}MB)")168 print(f"Tensor shape: {list(shape)} ({tensor_mb:.0f}MB per K/V)")169 print(sep)170 print()171 print("Index breakdown:")172 print(f" SVD extract: {extract_ms:8.1f}ms")173 print(f" Serialize .eng: {ser_ms:8.1f}ms")174 print(f" Store backend: {store_ms:8.1f}ms")175 print(f" FAISS add(): {add_ms:8.1f}ms")176 print(f" TOTAL: {index_ms:8.1f}ms")177 print()178 print("Retrieve breakdown:")179 print(f" SVD extract: {q_extract_ms:8.1f}ms")180 print(f" FAISS search(): {search_ms:8.1f}ms")181 print(f" Load+deser: {load_ms:8.1f}ms")182 print(f" TOTAL: {retrieve_ms:8.1f}ms")183 print()184 print("Verification:")185 print(f" Round-trip shape: {'OK' if k_out.shape == keys.shape else 'FAIL'}")186 print(f" Retrieval result: {'OK' if len(results) >= 1 else 'FAIL'}")187 print(f" .eng valid: {'OK' if eng_path.exists() else 'FAIL'}")188 189 return 0 if speedup > 10 else 1190 191 192def main():193 parser = argparse.ArgumentParser(194 description="ENGRAM Protocol — Demo Agent Session",195 epilog="D6: >10x TTFT reduction at 16K context on Llama 3.1 8B",196 )197 parser.add_argument(198 "--model", "-m", default=None,199 help="Path to GGUF model file (required unless --dry-run)",200 )201 parser.add_argument(202 "--context", "-c", type=int, default=4096,203 help="Context length to fill (tokens). Default: 4096",204 )205 parser.add_argument(206 "--n-ctx", type=int, default=16384,207 help="Max context window for model. Default: 16384",208 )209 parser.add_argument(210 "--data-dir", type=str, default=None,211 help="ENGRAM data directory. Default: ~/.engram/data",212 )213 parser.add_argument(214 "--dry-run", action="store_true",215 help="Run full pipeline with synthetic tensors (no model needed)",216 )217 parser.add_argument(218 "--verbose", "-v", action="store_true",219 help="Enable verbose output",220 )221 args = parser.parse_args()222 223 if args.dry_run:224 return _run_dry_run(args)225 226 if not args.model:227 parser.error("--model is required unless --dry-run is specified")228 229 print("=" * 70)230 print("ENGRAM Protocol — Demo Agent Session")231 print("KV cache fingerprinting for persistent semantic retrieval")232 print("=" * 70)233 print()234 235 # ── Setup ─────────────────────────────────────────────────236 from kvcos.core.config import get_config237 from kvcos.core.serializer import EngramSerializer238 from kvcos.core.types import CompressionMethod, StateExtractionMode239 from kvcos.core.manifold_index import ManifoldIndex240 from kvcos.core.retriever import EGRRetriever241 from kvcos.core.state_extractor import MARStateExtractor242 from kvcos.storage.local import LocalStorageBackend243 from integrations.llama_cpp_bridge import LlamaCppBridge244 245 config = get_config()246 data_dir = Path(args.data_dir) if args.data_dir else config.data_dir247 248 # ── Step 1: Load Model ────────────────────────────────────249 print(f"[1/6] Loading model: {args.model}")250 bridge = LlamaCppBridge(251 model_path=args.model,252 n_ctx=args.n_ctx,253 n_gpu_layers=0, # D1254 verbose=args.verbose,255 )256 spec = bridge.load_model()257 print(f" Model: {spec['model_id']}")258 print(f" Architecture: {spec['n_layers']}L / {spec['n_heads']}H / {spec['n_kv_heads']}KV / {spec['head_dim']}D")259 print(f" Context window: {args.n_ctx}")260 print()261 262 # ── Step 2: Generate + Cold TTFT ──────────────────────────263 filler = "The quick brown fox jumps over the lazy dog. " * 100264 target_tokens = args.context265 prompt = filler[:target_tokens * 4]266 267 print(f"[2/6] Cold prefill ({target_tokens} target tokens)...")268 t0 = time.perf_counter()269 cold = bridge.measure_cold_ttft(prompt)270 print(f" Cold TTFT: {cold.ttft_ms:.1f}ms ({cold.context_len} tokens)")271 print()272 273 # ── Step 3: Extract + Serialize ───────────────────────────274 print("[3/6] Extracting KV cache...")275 try:276 parsed = bridge.extract_kv_cache()277 print(f" Keys shape: {list(parsed.keys.shape)}")278 print(f" Values shape: {list(parsed.values.shape)}")279 print(f" Cells: {parsed.n_cells}")280 except Exception as e:281 print(f" KV extraction failed: {e}")282 print(" This is expected if the blob format doesn't match.")283 print(" Falling back to save_state/load_state raw blob path.")284 parsed = None285 print()286 287 print("[3b/6] Saving raw state blob...")288 raw_state = bridge.llm.save_state()289 raw_blob = bytes(raw_state.llama_state)290 print(f" Raw state size: {len(raw_blob) / 1024 / 1024:.1f} MB")291 292 if parsed is not None:293 print("[3c/6] Serializing to .eng format...")294 serializer = EngramSerializer()295 eng_path = data_dir / "demo" / "session_001.eng"296 result = serializer.serialize(297 keys=parsed.keys,298 values=parsed.values,299 agent_id="demo-agent",300 task_description="demo session - cold prefill benchmark",301 model_id=spec["model_id"],302 output_path=eng_path,303 compression=CompressionMethod.Q8_0,304 )305 print(f" .eng file: {result['path']}")306 print(f" Size: {result['size_bytes'] / 1024 / 1024:.1f} MB")307 print(f" Compression ratio: {result['compression_ratio']:.2f}x")308 print()309 310 # ── Step 4: Index in EGR ──────────────────────────────────311 if parsed is not None:312 print("[4/6] Indexing in EGR manifold index...")313 storage = LocalStorageBackend(data_dir=data_dir)314 extractor = MARStateExtractor(315 mode=StateExtractionMode.SVD_PROJECT,316 rank=min(160, spec["head_dim"]),317 )318 dim = extractor.output_dim(spec)319 index = ManifoldIndex(dim=dim)320 retriever = EGRRetriever(extractor, index, storage)321 322 cache_id = retriever.index_engram(323 keys=parsed.keys,324 values=parsed.values,325 spec=spec,326 agent_id="demo-agent",327 task_description="demo session - cold prefill benchmark",328 model_id=spec["model_id"],329 )330 print(f" Indexed: {cache_id}")331 print(f" State vector dim: {dim}")332 print(f" Index entries: {index.n_entries}")333 else:334 print("[4/6] Skipped (KV extraction failed)")335 print()336 337 # ── Step 5: Restore + Cached TTFT ─────────────────────────338 print("[5/6] Restoring from cached state...")339 t0 = time.perf_counter()340 cached = bridge.measure_cached_ttft(raw_blob)341 print(f" Cached TTFT: {cached.ttft_ms:.1f}ms")342 print()343 344 # ── Step 6: Results ───────────────────────────────────────345 cold_ms = cold.ttft_ms346 cached_ms = cached.ttft_ms347 speedup = cold_ms / cached_ms if cached_ms > 0 else float("inf")348 349 eng_path_str = result["path"] if parsed else "N/A"350 eng_size_kb = result["size_bytes"] / 1024 if parsed else 0351 352 sep = "=" * 35353 print(sep)354 print("ENGRAM Protocol — EGR Demo")355 print(f"Model: {spec['model_id']}")356 print(f"Context: {cold.context_len} tokens")357 print(sep)358 print(f"Cold TTFT: {cold_ms:.1f}ms")359 print(f"Cached TTFT: {cached_ms:.1f}ms")360 print(f"Speedup: {speedup:.1f}x")361 print(f"D6 target: >10x at 16K tokens")362 status = "PASS" if speedup > 10 else "FAIL"363 print(f"Status: {status}")364 print(f".eng file: {eng_path_str} ({eng_size_kb:.1f}KB)")365 print(sep)366 367 return 0 if speedup >= 4 else 1368 369 370if __name__ == "__main__":371 sys.exit(main())372