51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
"""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()
|