[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.

Why

Agent clients start every new session with the same head: tool definitions and system text, often several thousand tokens. MTPLX can reuse that head from an earlier session, but on a hybrid model (Qwen3.5, 3.6 and 3.8, which mix attention with GDN layers) only up to a point where it saved a snapshot of the GDN state. A new session borrows those snapshots from an earlier conversation that started with the same head. The session bank keeps at most 8 snapshots per conversation: the oldest one, and the rest close to where that conversation currently ends. Take a 6,000-token head and an earlier conversation that is now 50,000 tokens long: its snapshots sit at its oldest point and around 40,000 to 50,000 tokens, none at token 6,000. The new session can only resume from a snapshot that lies within the shared head, so it falls back to the oldest one and computes the rest of the head again, up to several thousand tokens.

How it works, in plain terms

An attention layer can resume at any token, because its cache simply holds every token. A GDN layer instead carries one running state that summarises everything before it; to resume at a given token you need a snapshot of that state taken at exactly that token. With the switch on, MTPLX finds in the prompt tokens where the head ends (the first turn that is not a system turn), always saves a snapshot there, and never removes it when the bank thins out its snapshots. A new session with the same head then resumes right after the head, also from the SSD cache after a restart, and gives bit for bit the same result as computing the head again. There is no anchor for tokenizers without ChatML turns (such as Gemma 4), for image prompts, or for a head shorter than the restore minimum. With the switch off nothing changes.

What changes

• 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

[youssofal]

Not in 2.12.1, which was a stability release. I'll review this for the next release.

[jvmenen]

A note on how the head anchor reaches the bank. The anchor is recorded in the prompt's checkpoint list, and that list is banked by the store-on-prefill snapshot and the prompt-prefix commit; the generation-final commit inherits it from them. When the prefill admission is near the memory line it skips the prompt's banked copy (prompt_publish_skipped), and then the anchor was dropped with the request: the next session with the same head restored from an older entry's grid checkpoint (1,024 of an 11,843-token head in our logs) or prefilled cold.

#637 fixes that on its own, independent of this PR: the generation-final commit now banks the prompt's checkpoints that lie inside the prefix the final history shares with the prompt. It helps without the anchor as well (the grid checkpoints survive), and with this PR it keeps the anchor through the cold route.

Once #637 is in, I would add the test below here, so the anchor's cold route is covered next to the RAM and SSD cases.

def test_the_anchor_survives_a_cold_turn_without_the_prompt_commit(bank_env):
    """The admission skipped the prompt's banked copy (prompt_publish_skipped):
    the generation-final commit is the turn's only entry, and it keeps the anchor."""
    from mtplx.server.openai import _generation_final_prompt_boundaries

    bank_env.setenv(SWITCH, "1")
    rt = _BankRuntime()
    prompt = _session(first_user_seed=2, length=33_000)
    anchors = generation._session_head_anchors(rt, prompt)
    sink = _anchored_sink(anchors, [], prompt_len=len(prompt))
    _simulated_prefill(sink, 0, len(prompt), anchors)
    final = prompt + _filler(400, seed=5)

    def bank_with(boundaries):
        bank = SessionBank(max_entries=8, max_bytes=1 << 34, per_session_max_bytes=1 << 34)
        _put(bank, final, boundaries)  # no prompt entry: the admission skipped it
        return bank

    kept = _generation_final_prompt_boundaries(
        SimpleNamespace(prompt_gdn_boundaries=tuple(sink)),
        prompt_ids=prompt,
        final_token_ids=final,
    )
    assert HEAD in [record[0] for record in kept]
    assert _new_session_restore_point(bank_with(kept)) == HEAD
    # The old commit passed nothing: no donor, no checkpoint under the head.
    assert _new_session_restore_point(bank_with(None)) == 0

It reuses the bank simulation helpers of tests/test_session_head_anchor.py (_session, _anchored_sink, _simulated_prefill, _put, _new_session_restore_point).