Save inference parity implementation and evaluation harness

tools/mtplx-cache-reference.py

@@ -0,0 +1,50 @@
+"""Reproduce MTPLX's AR prefix-selection failure without model/GPU work.
+
+Uses the pinned reference's real bank and lookup functions with empty cache
+payloads. This verifies selection policy only, not recurrent-state correctness.
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+from mtplx import generation
+from mtplx.session_bank import SessionBank
+
+
+def main():
+    runtime = SimpleNamespace(model_path=Path("/diagnostic/qwen"), mtp_enabled=True)
+    bank = SessionBank()
+    prompt = list(range(7896))
+    for length, policy in ((7460, "cycle"), (7836, "committed")):
+        bank.put(runtime=runtime, token_ids=prompt[:length], cache=[],
+                 logits=None, hidden=None, mtp_history_policy=policy)
+
+    assert bank.longest_prefix(prompt).prefix_len == 7836
+    assert bank.restore(runtime, prompt, mtp_history_policy="cycle",
+                        cache_factory=list) is None
+    assert bank.last_miss_reason == "policy_mismatch"
+    candidates = bank.near_prefix_candidates(prompt, mtp_history_policy="cycle")
+    assert [(entry.prefix_len, matched) for entry, matched in candidates] == [
+        (7836, 7836), (7460, 7460),
+    ]
+    assert generation._restore_near_prefix_prompt_state(
+        runtime, prompt, base_hidden_variant=None, mtp_hidden_variant=None,
+        mtp_history_policy="cycle", session_bank=bank, template_hash=None,
+        draft_head_identity=None, policy_fingerprint=None, cache_factory=list,
+    ) is None
+
+    # Control: the shorter prefix is usable; merely removing the incompatible
+    # longest candidate makes the unchanged exact-restore function serve it.
+    del bank._entries[tuple(prompt[:7836])]
+    restored = bank.restore(runtime, prompt, mtp_history_policy="cycle",
+                            cache_factory=list)
+    assert restored is not None and restored.entry.prefix_len == 7460
+    print(json.dumps({"ok": True, "reference": "MTPLX SessionBank + generation",
+                      "incompatible_longest_prefix": 7836,
+                      "shadowed_usable_prefix": 7460, "prompt_tokens": len(prompt),
+                      "control_restored_prefix": restored.entry.prefix_len}))
+
+
+if __name__ == "__main__":
+    main()