Commit the final Qwen MTP primary before saving state

tools/issue106-reference.py

@@ -0,0 +1,38 @@
+"""Run the pinned MTPLX final-state exit matrix without loading model weights."""
+import importlib.util
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+root = Path(sys.argv[1]).resolve()
+assert subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
+assert os.environ["MTPLX_CONTEXT_COPY"] == "0"
+sys.path.insert(0, str(root))
+spec = importlib.util.spec_from_file_location("final_state", root / "tests/test_mtp_final_state_covers_committed_tokens.py")
+m = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(m)
+for label, stop_index, depth, omniscient, temperature, maximum in [
+    ("bonus primary", 8, 2, False, 0.0, 16),
+    ("deferred greedy correction", 8, 3, False, 0.0, 16),
+    ("accepted draft stop", 8, 3, True, 0.0, 16),
+    ("sampled correction", 8, 3, False, 0.7, 16),
+    ("first primary", 6, 3, False, 0.0, 16),
+    ("fresh primary depth 1", 7, 1, False, 0.0, 16),
+    ("max tokens depth 3", 10000, 3, False, 0.0, 7),
+    ("max tokens depth 2", 10000, 2, False, 0.0, 5),
+]:
+    sampler = m.SamplerConfig(temperature=temperature, top_k=1 if temperature else 0)
+    out = m.generate_mtpk(m._runtime(m._ScriptedModel(stop_index, omniscient)), m.PROMPT,
+                         speculative_depth=depth, max_tokens=maximum, sampler=sampler, **m.COMMON)
+    gap, origin, safe = m._gap(m.PROMPT, out)
+    assert gap == 0 and safe
+    assert out.final_state.final_committed_mtp_cache[0].offset == len(m.PROMPT) + len(out.tokens) - 1
+    if stop_index < 1000:
+        assert out.tokens[-1] == m.STOP
+    print(json.dumps(dict(case="mtp_exit", label=label, stop_index=stop_index, depth=depth,
+                          omniscient=omniscient, temperature=temperature, max_tokens=maximum,
+                          tokens=out.tokens, gap=gap, origin=origin, safe=safe)), flush=True)
+m.test_ar_final_state_control()
+print(json.dumps(dict(case="ar_control", passed=True)), flush=True)