1"""Scientific correctness tests for the cache adapters (Section 7.2 of the brief). Runs a real small model."""2import os, json, pytest, numpy as np3import mlx.core as mx4from wp30.backends.mlx_backend import MLXModel, deep_copy_cache5from wp30.cache.sparse import select_positions, quantize_cache, to_sparse_all, cache_tensor_bytes, dense_bytes_formula6from wp30.cache.serialize import build_manifest, serialize_cache, deserialize_cache, IncompatibleCache7from wp30.cache.selection import select_cache, sink_recent_indices, random_indices89REPO = os.environ.get("WP30_TEST_MODEL", "mlx-community/Qwen2.5-1.5B-Instruct-bf16")10DOC = ("Project Varkolith-7 was proposed by Teodora Qvistrand on 14 March 2031 with a budget of 4,381,220 CAD. On 2 April 2031 the budget was "11 "revised to 3,907,415 CAD and the project was approved. The site is Lot 83-K in Chibougamau. The supplier identifier is ZR-77Q-0931. ") * 412Q = "What is the supplier identifier? Answer with the identifier only."1314def kl_div(a, b):15 pa = mx.softmax(a); pb = mx.softmax(b)16 return float((pa * (mx.log(pa + 1e-12) - mx.log(pb + 1e-12))).sum())1718def max_dprob(a, b):19 return float(mx.abs(mx.softmax(a) - mx.softmax(b)).max())2021@pytest.fixture(scope="module")22def m():23 return MLXModel(REPO)2425@pytest.fixture(scope="module")26def m32(m):27 """float32 copy of the model for tolerance-free algorithmic checks."""28 mm = MLXModel(REPO)29 mm.model.set_dtype(mx.float32)30 mx.eval(mm.model.parameters())31 return mm3233def split_ids(model):34 sys_msg = "Answer from the document."35 full = model.chat_prompt_tokens([{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + DOC + "\n\nQuestion: " + Q}])36 pre_txt = model.tokenizer.apply_chat_template([{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + DOC + "\n\n"}], add_generation_prompt=False, tokenize=False)37 pre = model.encode(pre_txt)38 n = 039 while n < min(len(pre), len(full)) and pre[n] == full[n]:40 n += 141 return full, full[:n], full[n:]4243def test_prefix_equivalence_fp32(m32):44 full, pre, rest = split_ids(m32)45 c_full = m32.new_cache(); m32.prefill(full[:-1], c_full); ref = m32.next_logits(c_full, full[-1:])46 c_pre, _ = m32.prefill(pre); lg = m32.next_logits(c_pre, rest)47 # Registered tolerance (DECISIONS D-006): MLX kernels are not bitwise deterministic across prefill chunkings even in48 # fp32 (measured max|dlogit| ~0.1 between one-chunk and token-by-token prefill); the criterion is distributional:49 # KL(ref || cached) < 1e-5, max |dprob| < 1e-4, identical top-1.50 assert kl_div(ref, lg) < 1e-5 and max_dprob(ref, lg) < 1e-451 assert int(mx.argmax(lg)) == int(mx.argmax(ref))5253def test_prefix_equivalence_bf16_top1(m):54 full, pre, rest = split_ids(m)55 c_full = m.new_cache(); m.prefill(full[:-1], c_full); ref = m.next_logits(c_full, full[-1:])56 c_pre, _ = m.prefill(pre); lg = m.next_logits(c_pre, rest)57 assert int(mx.argmax(lg)) == int(mx.argmax(ref))58 assert kl_div(ref, lg) < 1e-3 and max_dprob(ref, lg) < 1e-2 # bf16 registered tolerance; top-1 must agree59 # greedy continuation agrees60 g1 = m.generate_text(full, max_tokens=8)["text"]61 c_pre, _ = m.prefill(pre); g2 = m.generate_from_cache(c_pre, rest, max_tokens=8)["text"]62 assert g1 == g26364def test_no_duplicated_boundary_token(m):65 full, pre, rest = split_ids(m)66 assert pre + rest == full and len(rest) >= 167 bos = m.info.bos_token_id68 if bos is not None:69 assert full.count(bos) <= 17071def test_identity_roundtrip_and_bytes(m):72 full, pre, rest = split_ids(m)73 c_pre, _ = m.prefill(pre)74 lg = m.next_logits(deep_copy_cache(c_pre), rest)75 man = build_manifest(m.info, pre, c_pre, "doc", "B1", {})76 payload = serialize_cache(c_pre, man)77 c2, man2 = deserialize_cache(payload, m.info, man["prefix_token_hash"])78 lg2 = m.next_logits(c2, rest)79 assert float(mx.abs(lg2 - lg).max()) == 0.080 tb = cache_tensor_bytes(c_pre)81 assert tb == dense_bytes_formula(m.info.num_layers, len(pre), m.info.num_kv_heads, m.info.head_dim, 2)82 assert len(payload) >= tb and len(payload) < tb * 1.01 + 655368384def test_loader_rejects_wrong_prefix_and_model(m):85 full, pre, rest = split_ids(m)86 c_pre, _ = m.prefill(pre)87 payload = serialize_cache(c_pre, build_manifest(m.info, pre, c_pre, "doc", "B1", {}))88 with pytest.raises(IncompatibleCache):89 deserialize_cache(payload, m.info, "deadbeef")90 import copy91 other = copy.copy(m.info); other.config_hash = "x"92 with pytest.raises(IncompatibleCache):93 deserialize_cache(payload, other)9495def test_full_retention_control_exact(m):96 full, pre, rest = split_ids(m)97 c_pre, _ = m.prefill(pre); lg = m.next_logits(deep_copy_cache(c_pre), rest)98 T = len(pre)99 c_all = select_positions(c_pre, [np.arange(T)] * m.info.num_layers)100 assert float(mx.abs(m.next_logits(c_all, rest) - lg).max()) == 0.0101 # per-head 2-D indices identical across heads must also be exact102 idx2 = np.tile(np.arange(T), (m.info.num_kv_heads, 1))103 c_all2 = select_positions(c_pre, [idx2] * m.info.num_layers)104 assert float(mx.abs(m.next_logits(c_all2, rest) - lg).max()) == 0.0105106def test_sparse_mask_path_matches_token_by_token_fp32(m32):107 full, pre, rest = split_ids(m32)108 c_pre, _ = m32.prefill(pre)109 T = len(pre)110 idx = sink_recent_indices(T, T // 2)111 a = select_positions(c_pre, [idx] * m32.info.num_layers); lg_a = m32.next_logits(a, rest)112 b = select_positions(c_pre, [idx] * m32.info.num_layers)113 for t in rest[:-1]:114 m32.next_logits(b, [t])115 lg_b = m32.next_logits(b, rest[-1:])116 assert kl_div(lg_a, lg_b) < 1e-5 and int(mx.argmax(lg_a)) == int(mx.argmax(lg_b))117 # logical offset preserved: new tokens positioned at T, T+1, ...118 assert a[0].offset == T + len(rest) and a[0].n_kept == len(idx) + len(rest)119120def test_positions_preserved_vs_relabel(m32):121 """Retaining positions with the original offset must differ from relabelling them as 0..k-1 (which would be StreamingLLM-style)."""122 full, pre, rest = split_ids(m32)123 c_pre, _ = m32.prefill(pre); T = len(pre)124 idx = sink_recent_indices(T, T // 2)125 a = select_positions(c_pre, [idx] * m32.info.num_layers)126 b = select_positions(c_pre, [idx] * m32.info.num_layers, logical_offset=len(idx))127 assert float(mx.abs(m32.next_logits(a, rest) - m32.next_logits(b, rest)).max()) > 0.0128129def test_packed_quantization_bytes_and_execution(m):130 full, pre, rest = split_ids(m)131 c_pre, _ = m.prefill(pre)132 dense = cache_tensor_bytes(c_pre)133 for bits, lo, hi in [(8, 1.6, 2.0), (4, 3.0, 4.0), (2, 5.0, 8.0)]:134 cq = quantize_cache(deep_copy_cache(c_pre), bits)135 b = cache_tensor_bytes(cq)136 assert lo <= dense / b <= hi, (bits, dense / b)137 k0 = cq[0].keys_and_values()[0]138 assert k0[0].dtype == mx.uint32 # packed words, not a wide dtype139 g = m.generate_from_cache(cq, rest, max_tokens=4)140 assert g["output_tokens"] >= 1141142def test_no_shared_mutable_state(m):143 full, pre, rest = split_ids(m)144 c_pre, _ = m.prefill(pre)145 b1 = deep_copy_cache(c_pre); b2 = deep_copy_cache(c_pre)146 m.next_logits(b1, rest)147 assert b2[0].offset == c_pre[0].offset == len(pre)148 assert float(mx.abs(m.next_logits(b2, rest) - m.next_logits(deep_copy_cache(c_pre), rest)).max()) == 0.0149150def test_attention_selection_runs_and_keeps_window(m):151 full, pre, rest = split_ids(m)152 c_pre, _ = m.prefill(pre); T = len(pre)153 out, info = select_cache(m, c_pre, "attention", 0.25, window_token_ids=rest, window_keep=8)154 assert out[0].n_kept == info["keep"] and out[0].offset == T155 rp = np.array(out[0].retained_positions)156 assert rp.ndim == 2 and all((rp[h] >= T - 8).sum() == 8 for h in range(rp.shape[0]))157 g = m.generate_from_cache(out, rest, max_tokens=6)158 assert g["output_tokens"] >= 1159160def test_no_message_control_does_not_know_answer(m):161 ids = m.chat_prompt_tokens([{"role": "system", "content": "Answer."}, {"role": "user", "content": Q}])162 g = m.generate_text(ids, max_tokens=12)163 assert "ZR-77Q-0931" not in g["text"]164