Session bank: anchor a GDN boundary at the end of the prompt head (opt-in, depends on #550)

[jvmenen]

Depends on #550 (and through it on #549). This branch is #550 plus one commit; only the last commit (perf(session-bank): anchor a GDN boundary at the end of the prompt head (opt-in)) is new here. The feature itself does not need the batch-invariant lane or the in-forward boundaries: without them the anchor is an ordinary mandatory span end of the tail ladder. The dependency is in the tests: the bit-for-bit restore tests run a tiny Qwen3.5-MoE on the lane with in-forward boundaries and reuse the fixtures of tests/test_gdn_inforward_boundaries.py from #550. Once #550 is merged, the diff reduces to that commit.

Summary

Opt-in MTPLX_SESSION_HEAD_ANCHOR=1: on hybrid (GDN) models, every prefill that goes to the session bank records a GDN restore boundary exactly where the prompt's fixed head (tool schemas and system turn) ends, and boundary thinning never drops it. A new session with the same head then restores the whole head from any banked entry of an earlier session, instead of only up to the oldest kept boundary. Default off; off keeps planning and retention unchanged.

Motivation

Agent harnesses start every session with the same head: tool schemas and the system text, often several thousand tokens. Sessions diverge at the first user turn. On a hybrid model a restore needs a recurrent (GDN) state at or below the shared prefix, and the bank only holds states at recorded boundaries. Thinning keeps at most 8 boundaries per entry: the oldest record and a geometric ladder near the tail. The donor for a new session is usually a long conversation, whose kept boundaries sit near its tail, far past the head, so a new session restores only from the oldest boundary and prefills the rest of the head again.

A boundary exactly at the end of the head, kept through thinning and inheritance, makes the whole head reusable.

Change

• New mtplx/session_head_anchor.py: turn_markers(tokenizer) finds the ChatML turn-open token (atomic <|im_start|>) and the system role token; session_head_length(prompt_ids, markers) returns the number of tokens before the first turn that is not a system turn. Read from the prompt tokens, so it is token-exact for every ChatML template, harness and endpoint (chat completions and /v1/messages). No anchor for tokenizers without ChatML turns (for example Gemma 4), for image prompts, or for a head shorter than MTPLX_SESSION_BLOCK_PREFIX_MIN_MATCH_TOKENS (no restore could use it).
• mtplx/generation.py:
  • GdnBoundarySink: the boundary list a prefill appends to, plus its anchors. A plain list stays a sink without anchors.
  • _mandatory_prefill_edges merges the existing stable-prefix edge with the anchors; the cold loops, the warm suffix loop and _predicted_first_prefill_span plan the anchor as a span end (inside the forward where the model publishes in-forward boundary hooks).
  • _thin_gdn_boundary_records(..., keep=...), the sink's thinning and _inherited_gdn_boundaries never drop a record at an anchor.
• mtplx/cache_bank/cold_tier.py: with the switch on, the SSD block-prefix lane restores to the exact match instead of the block-aligned one (the RAM lane already does), so an unaligned anchor is reachable after a restart. The prefix decode takes any length, and a boundary restore picks the newest persisted boundary at or below the match.
• mtplx/runtime_options.py: session_head_anchor_enabled() (MTPLX_SESSION_HEAD_ANCHOR, default off).

How to enable: MTPLX_SESSION_HEAD_ANCHOR=1, with a session bank on a hybrid model.

Tests

New tests/test_session_head_anchor.py (43 tests):

• Head detection: first non-system turn, consecutive system turns, no head without a system turn or a later turn, atomic turn-open only; the real chat templates of locally cached Qwen packs with and without thinking, system and tools (skipped when the tokenizer is not in ~/.mtplx/models), /v1/messages, Gemma without ChatML turns.
• Planning and retention: stable prefix and anchors merge into mandatory edges; the anchor is recorded inside the last wide forward; thinning, the sink and inheritance keep it within the cap.
• A bank simulation of a long chat with the real planner, retention and SessionBank: a new session restores at the anchor with the switch on, at the oldest boundary with it off; the anchor survives the SSD tier.
• On a tiny Qwen3.5-MoE with the invariant lane (Metal): the anchor leaves the prefill bit-identical (logits, hidden, cache); a new session restored at the anchor from RAM or from the SSD gives the cold prefill's logits and hidden bit for bit; an anchor inside a warm suffix does not change the result or the forward count.
• tests/test_qwen4_ple_first_gather_early.py (updated for the planner helper), tests/test_gdn_inforward_boundaries.py, tests/test_batch_invariant_prefill.py and tests/test_session_bank.py pass.
• Full suite with MTPLX_CONFIG=/nonexistent: 9,535 passed, 67 skipped, 1 failed. The failure is test_laguna_model.py::test_laguna_s_2_1_ar_route_skips_qwen_performance_hooks, which fails the same way on main on a Mac with less than 85.3 GiB (addressed in Tests no longer read the user's ~/.mtplx/config.toml or depend on the Mac's memory size #535). Ruff: new files clean, no new findings in the changed files.

🤖 Generated with Claude Code