1"""E0 smoke: real prefill/cache/generation on a small MLX model; identity, prefix-equivalence, quantization, selection."""2import sys, json, time3import numpy as np4import mlx.core as mx5from wp30.backends.mlx_backend import MLXModel, deep_copy_cache, memory_snapshot6from wp30.cache.sparse import SparseKVCache, select_positions, quantize_cache, dense_bytes_formula, cache_tensor_bytes7from wp30.cache.serialize import build_manifest, serialize_cache, deserialize_cache, codec_compress89repo = sys.argv[1] if len(sys.argv) > 1 else "mlx-community/Qwen2.5-1.5B-Instruct-bf16"10m = MLXModel(repo)11print("loaded", repo, "in", round(m.load_ms), "ms; layers", m.info.num_layers, "kv_heads", m.info.num_kv_heads, "head_dim", m.info.head_dim, "dtype", m.info.weight_dtype)1213doc = ("Internal memo. Project Varkolith-7 was proposed by Ms. Teodora Qvistrand on 14 March 2031 with a budget of 4,381,220 CAD. "14 "On 2 April 2031 the budget was revised to 3,907,415 CAD and the project was approved. The site is Lot 83-K in Chibougamau. "15 "The supplier identifier is ZR-77Q-0931. The project was NOT cancelled; the launch was postponed to 19 September 2031. ") * 316question = "What is the approved budget of Project Varkolith-7, in CAD? Answer with the number only."17sys_msg = "You answer questions about the provided document. Be concise."18prefix_msgs = [{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + doc + "\n\n"}]19full_msgs = [{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + doc + "\n\nQuestion: " + question}]20full_ids = m.chat_prompt_tokens(full_msgs)21pre_text = m.tokenizer.apply_chat_template(prefix_msgs, add_generation_prompt=False, tokenize=False)22# strip trailing end-of-turn markers to get the raw prefix text, then longest common token prefix23pre_ids_raw = m.encode(pre_text)24n = 025while n < min(len(pre_ids_raw), len(full_ids)) and pre_ids_raw[n] == full_ids[n]:26 n += 127prefix_ids, rest_ids = full_ids[:n], full_ids[n:]28print("full tokens", len(full_ids), "prefix tokens", len(prefix_ids), "fresh query tokens", len(rest_ids))2930# 1. reference: full fresh prefill + greedy31ref = m.generate_text(full_ids, max_tokens=16)32print("REF answer:", repr(ref["text"]), "prefill_ms", round(ref["prefill_ms"]), "ttft_ms", round(ref["ttft_ms"]))33# reference logits at first generated position34c_full = m.new_cache(); m.prefill(full_ids[:-1], c_full); ref_logits = m.next_logits(c_full, full_ids[-1:])3536# 2. prefix cache + append fresh tokens (multi-token append -> mask path)37c_pre, t_pre = m.prefill(prefix_ids)38print("prefix cache bytes (measured)", cache_tensor_bytes(c_pre), "formula", dense_bytes_formula(m.info.num_layers, len(prefix_ids), m.info.num_kv_heads, m.info.head_dim, 2))39c_b = deep_copy_cache(c_pre)40lg = m.next_logits(c_b, rest_ids)41d = float(mx.abs(lg - ref_logits).max()); print("PREFIX-EQUIVALENCE max|dlogit| =", d, "argmax equal:", int(mx.argmax(lg)) == int(mx.argmax(ref_logits)))42c_b = deep_copy_cache(c_pre); g = m.generate_from_cache(c_b, rest_ids, max_tokens=16); print("prefix-reuse answer:", repr(g["text"]), "ttft_ms", round(g["ttft_ms"]))4344# 3. identity round trip through serialization45man = build_manifest(m.info, prefix_ids, c_pre, "smoke_doc", "B1", {})46payload = serialize_cache(c_pre, man)47gz, tgz = codec_compress(payload, "zstd")48print("serialized bytes", len(payload), "zstd bytes", len(gz), "zstd_ms", round(tgz, 1))49c_l, man2 = deserialize_cache(payload, m.info)50lg2 = m.next_logits(c_l, rest_ids); print("IDENTITY-ROUNDTRIP max|dlogit| =", float(mx.abs(lg2 - lg).max()))5152# 4. 100% retention via SparseKVCache53T = len(prefix_ids)54c_all = select_positions(c_pre, [np.arange(T)] * m.info.num_layers)55lg3 = m.next_logits(c_all, rest_ids); print("100%-RETENTION max|dlogit| =", float(mx.abs(lg3 - lg).max()))56# 4b. sparse mask path vs token-by-token append equivalence57c_s1 = select_positions(c_pre, [np.arange(T)] * m.info.num_layers)58for tok in rest_ids[:-1]:59 m.next_logits(c_s1, [tok])60lg4 = m.next_logits(c_s1, rest_ids[-1:]); print("MASK-PATH vs TOKEN-BY-TOKEN max|dlogit| =", float(mx.abs(lg4 - lg3).max()))6162# 5. quantization 8/4/263for bits in (8, 4, 2):64 cq = quantize_cache(deep_copy_cache(c_pre), bits)65 by = cache_tensor_bytes(cq)66 lgq = m.next_logits(deep_copy_cache(cq), rest_ids)67 g = m.generate_from_cache(cq, rest_ids, max_tokens=16)68 print(f"QUANT {bits}-bit: bytes {by} (ratio {cache_tensor_bytes(c_pre)/by:.2f}) max|dlogit| {float(mx.abs(lgq-lg).max()):.3f} answer {g['text']!r} ttft {g['ttft_ms']:.0f}ms")6970# 6. sink + recent selection at 1/2 and 1/471for frac in (0.5, 0.25):72 k = int(T * frac); sinks = 473 idx = np.concatenate([np.arange(sinks), np.arange(T - (k - sinks), T)])74 cs = select_positions(c_pre, [idx] * m.info.num_layers)75 g = m.generate_from_cache(cs, rest_ids, max_tokens=16)76 print(f"SINK+RECENT keep {k}/{T}: bytes {cache_tensor_bytes(cs)} answer {g['text']!r}")77 rng = np.random.default_rng(0); idx = np.sort(rng.choice(T, k, replace=False))78 cs = select_positions(c_pre, [idx] * m.info.num_layers)79 g = m.generate_from_cache(cs, rest_ids, max_tokens=16)80 print(f"RANDOM keep {k}/{T}: answer {g['text']!r}")81# 7. no-message control82nm_ids = m.chat_prompt_tokens([{"role": "system", "content": sys_msg}, {"role": "user", "content": "Question: " + question}])83g = m.generate_text(nm_ids, max_tokens=16); print("NO-MESSAGE answer:", repr(g["text"]))84print(json.dumps(memory_snapshot()))85