From f1d7d1e19d6a1306452cc0852e4998faac05361b Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Sun, 4 Oct 2026 00:39:07 +0200 Subject: [PATCH 1/2] feat: context-copy drafting for AR-only runtimes behind MTPLX_CONTEXT_COPY_AR Reuses NgramIndex, block_for_ext and the probation/EMA/suspend policy (now CopyGovernor) of the MTP context-copy lane; verifies a prompt block with one multi-row forward, point-mass acceptance with residual resampling when sampling, KV trim on rejection. Default off. Co-Authored-By: Claude Sonnet 5.5 --- mtplx/context_copy.py | 55 +++++ mtplx/context_copy_ar.py | 181 ++++++++++++++++ mtplx/generation.py | 94 +++++++- tests/test_context_copy_ar.py | 395 ++++++++++++++++++++++++++++++++++ 4 files changed, 715 insertions(+), 10 deletions(-) create mode 100644 mtplx/context_copy_ar.py create mode 100644 tests/test_context_copy_ar.py diff --git a/mtplx/context_copy.py b/mtplx/context_copy.py index 5c722cf7b..d9d94e0c2 100644 --- a/mtplx/context_copy.py +++ b/mtplx/context_copy.py @@ -18,6 +18,10 @@ block can commit far more per verify call (see the benchmarks in the pull request). The two mechanisms compose: copy when a prompt match exists, MTP otherwise. +AR-only runtimes (no MTP head) can opt in to the same lane in generate_ar with +MTPLX_CONTEXT_COPY_AR=1 (see context_copy_ar.py); it reuses NgramIndex, +block_for_ext and CopyGovernor from here. + RAMP (adapting community PR #375 by @johninthewinter): an opt-in fixed/long block-length policy plus a mismatch-tolerant fuzzy re-anchor fallback for when the exact n-gram key misses. OFF BY DEFAULT (MTPLX_RAMP_ENABLED unset or @@ -40,6 +44,20 @@ def context_copy_enabled() -> bool: return (os.environ.get("MTPLX_CONTEXT_COPY") or "").strip() not in {"0", "false", "off"} +def context_copy_ar_enabled() -> bool: + """Context-copy for AR-only runtimes (no MTP head), in ``generate_ar``. + + Opt-in: MTPLX_CONTEXT_COPY_AR=1 (also true/on/yes), and only while the + global MTPLX_CONTEXT_COPY switch is not off. Default off: plain AR decoding + stays byte-identical to before. The copy block is verified with one + multi-row forward, whose batched numerics can differ from row-by-row + decoding in the last bits, so a greedy stream is not guaranteed + bit-identical to the one-token loop. + """ + value = (os.environ.get("MTPLX_CONTEXT_COPY_AR") or "").strip().lower() + return context_copy_enabled() and value in {"1", "true", "on", "yes"} + + def context_copy_batched_enabled() -> bool: """Context-copy block rounds on the BATCHED verify lane (default ON; MTPLX_CONTEXT_COPY_BATCHED=0 disables). @@ -302,6 +320,43 @@ def block_for_ext(ext: int, k_cap: int) -> int: return min(_BLOCK_LADDER[idx], max(4, k_cap)) +class CopyGovernor: + """Acceptance EMA, probation and suspend/backoff of a copy lane. + + The same policy the MTP decode loops keep inline: blocks stay at the + probation size until two sampled rounds hold the 0.5 starting EMA, an EMA + below 0.35 after three rounds suspends copying for ``backoff`` tokens, and + the backoff doubles on each suspension (reset to 64 once a round accepts at + least half its block). + """ + + def __init__(self) -> None: + self.ema = 0.5 + self.seen = 0 + self.suspend_until = 0 + self.backoff = 64 + self.suspensions = 0 + + def suspended(self, n_tokens: int) -> bool: + return n_tokens < self.suspend_until + + def block_cap(self, k: int, probation_k: int) -> int: + """Full block cap once proven, the probation cap before that.""" + return k if (self.seen >= 2 and self.ema >= 0.5) else probation_k + + def record(self, accepted: int, block_len: int, n_tokens: int) -> None: + """Fold in one verified block; ``n_tokens`` counts the stream after it.""" + self.ema = 0.7 * self.ema + 0.3 * (accepted / block_len) + self.seen += 1 + if accepted / block_len >= 0.5: + self.backoff = 64 + if self.seen >= 3 and self.ema < 0.35: + self.suspend_until = n_tokens + self.backoff + self.backoff = min(self.backoff * 2, 4096) + self.ema, self.seen = 0.5, 0 + self.suspensions += 1 + + def ladder_widths(k_caps) -> frozenset[int]: """Verify widths (primary plus block) of every full-length block these caps allow. diff --git a/mtplx/context_copy_ar.py b/mtplx/context_copy_ar.py new file mode 100644 index 000000000..7138e3b07 --- /dev/null +++ b/mtplx/context_copy_ar.py @@ -0,0 +1,181 @@ +"""Context-copy drafting for AR-only runtimes (no MTP head). + +``generate_ar`` normally commits one token per forward pass. With +MTPLX_CONTEXT_COPY_AR=1 it borrows the context-copy lane of the MTP loops: when +the tail of the stream matches an n-gram of the PROMPT, the prompt continuation +is proposed as a block and verified in ONE multi-row target forward pass. + +Everything that decides what to copy is shared with the MTP lanes +(``NgramIndex``, ``block_for_ext``, the probation/EMA/suspend policy in +``CopyGovernor`` and the MTPLX_CONTEXT_COPY_* switches). What is new here is the +verify step for a plain AR model: + +* the target runs on ``[token, *block]`` (``token`` is the sampled-but-not-yet- + forwarded token), so row ``i`` scores ``block[i]`` and row ``len(block)`` + holds the logits after the whole block; +* greedy decoding accepts the longest prefix matching the target argmax; +* sampled decoding accepts each copied token with the target's shaped + probability and draws the first rejection's correction from the residual + (point-mass proposal), so the output law stays the target distribution; +* the KV cache is trimmed back to the accepted prefix, and a stop token inside + the accepted prefix ends the block (it stays out of the cache, as after a + normal AR step). + +The prompt is the only source of blocks (never the generated text), as in the +MTP lanes. Requests with penalties, constraints, guards or hidden-state capture +do not use the lane. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Any + +from .context_copy import ( + CopyGovernor, + JoinedTokens, + NgramIndex, + block_for_ext, + context_copy_ar_enabled, + context_copy_block_k, + context_copy_min_ext, + context_copy_ng_max, + context_copy_ng_min, + context_copy_probation_k, +) + +# Sampled acceptance: (block_logits[len(block), vocab], block) -> (accepted, correction) +AcceptFn = Callable[[Any, Sequence[int]], "tuple[int, int | None]"] + + +@dataclass(frozen=True) +class CopyRoundResult: + tokens: list[int] # accepted block prefix, cut after a stop token + correction: int | None # residual sample at the first rejection (sampled only) + stopped: bool # the last token in ``tokens`` is a stop token + logits: Any # (1, vocab): logits after the committed block + block_len: int + + +def copy_ar_eligible( + *, + mtp_enabled: bool, + penalized: bool, + guarded: bool, + sync_eval: bool, + cache: Any, +) -> bool: + """Whether this ``generate_ar`` request may use the copy lane.""" + if mtp_enabled or penalized or guarded or not sync_eval: + return False + if not context_copy_ar_enabled(): + return False + return bool(cache) and all( + hasattr(entry, "offset") and callable(getattr(entry, "trim", None)) + for entry in cache + ) + + +def accepted_prefix_greedy(block: Sequence[int], target_ids: Sequence[int]) -> int: + """Longest prefix of ``block`` equal to the target argmax rows.""" + accepted = 0 + for proposed, chosen in zip(block, target_ids): + if int(proposed) != int(chosen): + break + accepted += 1 + return accepted + + +class ArContextCopy: + """Prompt n-gram index plus the shared acceptance governor, per request.""" + + def __init__(self, prompt_ids: list[int]) -> None: + self.prompt_ids = prompt_ids + self.k = context_copy_block_k() + self.probation_k = min(self.k, context_copy_probation_k()) + self.min_ext = context_copy_min_ext() + self.index = NgramIndex(context_copy_ng_min(), context_copy_ng_max()) + self.index.sync(prompt_ids) + self.governor = CopyGovernor() + self.probes = self.rounds = self.drafted = self.accepted = 0 + + def propose(self, tokens: list[int], budget: int) -> list[int]: + """Prompt continuation to copy after ``tokens`` (at most ``budget`` long).""" + if budget < 1 or self.governor.suspended(len(tokens)): + return [] + self.probes += 1 + pos, ext = self.index.find( + JoinedTokens(self.prompt_ids, tokens), max_pos=len(self.prompt_ids) + ) + if pos is None or ext < self.min_ext: + return [] + cap = self.governor.block_cap(self.k, self.probation_k) + length = min(block_for_ext(ext, cap), budget) + return [int(t) for t in self.prompt_ids[pos : pos + length]] + + def record(self, block_len: int, accepted: int, n_tokens: int) -> None: + self.rounds += 1 + self.drafted += block_len + self.accepted += accepted + self.governor.record(accepted, block_len, n_tokens) + + def summary(self) -> dict[str, Any]: + return { + "probes": self.probes, + "rounds": self.rounds, + "drafted_tokens": self.drafted, + "accepted_tokens": self.accepted, + "acceptance_rate": self.accepted / self.drafted if self.drafted else 0.0, + "suspensions": self.governor.suspensions, + } + + +def verify_copy_block( + forward: Callable[[Any], Any], + trim_cache: Callable[[int], bool], + cache_offset: int, + token: int, + block: Sequence[int], + stop_token_ids: set[int] | None, + *, + greedy: bool, + accept: AcceptFn | None = None, +) -> CopyRoundResult: + """One target pass over ``[token, *block]``; commit the accepted prefix. + + ``forward(ids)`` runs the target on a ``(1, n)`` int array and returns logits + ``(1, n, vocab)``; ``trim_cache(offset)`` rolls the KV cache back to that + length and reports success. ``accept`` is required when not ``greedy``. + The cache ends holding every committed token except the last, as after a + normal AR step (``token`` and the accepted block, minus a final stop token). + """ + import mlx.core as mx + + ids = mx.array([[int(token), *[int(t) for t in block]]]) + logits = forward(ids) + correction: int | None = None + if greedy: + chosen = mx.argmax(logits[0], axis=-1) + mx.eval(chosen) + accepted = accepted_prefix_greedy(block, chosen.tolist()) + else: + if accept is None: + raise ValueError("sampled copy verification needs an accept function") + accepted, correction = accept(logits[0, : len(block)], block) + stopped = False + stops = stop_token_ids or set() + for i in range(accepted): + if int(block[i]) in stops: + accepted, stopped, correction = i + 1, True, None + break + keep = cache_offset + 1 + accepted - (1 if stopped else 0) + if not trim_cache(keep): + raise RuntimeError("context-copy (AR): KV cache rollback failed") + return CopyRoundResult( + tokens=[int(t) for t in block[:accepted]], + correction=correction, + stopped=stopped, + logits=logits[:, accepted, :], + block_len=len(block), + ) diff --git a/mtplx/generation.py b/mtplx/generation.py index ebbd8eca7..6003c72a7 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -47,6 +47,7 @@ ) from .errors import is_allocation_failure from .progress_heartbeat import tick as _owner_progress_tick +from .context_copy_ar import ArContextCopy, copy_ar_eligible, verify_copy_block from .cache_state import ( CacheSnapshot, detach_array_leaf, @@ -9415,6 +9416,24 @@ def emit_token(token: int) -> None: in ("1", "true", "yes", "on") ) or bool(os.environ.get("MTPLX_EVAL_AUDIT")) + _cc_ar = ( + ArContextCopy(prompt_ids) + if copy_ar_eligible( + mtp_enabled=bool(getattr(rt, "mtp_enabled", False)), + penalized=bool(sampler.presence_penalty or sampler.frequency_penalty), + guarded=( + constraint is not None + or _loop_guard is not None + or _thinking_guard is not None + or ar_return_hidden + ), + sync_eval=_ar_sync_eval, + cache=cache, + ) + else None + ) + _cc_forced: int | None = None # residual correction awaiting its emit + for step in range(max_tokens): if _loop_guard is not None: _guard_transition = _loop_guard.observe(tokens) @@ -9458,15 +9477,18 @@ def emit_token(token: int) -> None: sync_elapsed = time.perf_counter() - sync_started target_eval_time += sync_elapsed target_decode_time += sync_elapsed - token, _ = _sample_from_logits( - logits_row, - sampler, - rng, - token_counts=Counter(tokens) - if (sampler.presence_penalty or sampler.frequency_penalty) - else None, - penalty_overlay=(_ar_steer_overlay(tokens) if _steer_active else None), - ) + if _cc_forced is not None: + token, _cc_forced = _cc_forced, None + else: + token, _ = _sample_from_logits( + logits_row, + sampler, + rng, + token_counts=Counter(tokens) + if (sampler.presence_penalty or sampler.frequency_penalty) + else None, + penalty_overlay=(_ar_steer_overlay(tokens) if _steer_active else None), + ) if first_token_logprobs_top_k is not None and not tokens: first_logprobs = first_token_logprobs( logits[0], token_id=token, top_k=first_token_logprobs_top_k @@ -9497,9 +9519,59 @@ def emit_token(token: int) -> None: } ) break - if step + 1 >= max_tokens or _is_stop(token, stop_token_ids): + if len(tokens) >= max_tokens or _is_stop(token, stop_token_ids): break + if _cc_ar is not None: + _cc_block = _cc_ar.propose(tokens, max_tokens - len(tokens) - 1) + if _cc_block: + _cc_started = time.perf_counter() + _cc_round = verify_copy_block( + lambda ids: rt.forward_ar(ids, cache=cache, return_hidden=False), + lambda keep: _trim_cache_to_offset(cache, keep), + int(cache[0].offset), + token, + _cc_block, + stop_token_ids, + greedy=sampler.temperature <= 0, + accept=lambda rows, block: _point_mass_block_accept( + rows, block, sampler, rng + ), + ) + _cc_elapsed = time.perf_counter() - _cc_started + target_decode_time += _cc_elapsed + target_eval_time += _cc_elapsed + verify_calls += 1 + _cc_ar.record( + _cc_round.block_len, + len(_cc_round.tokens), + len(tokens) + len(_cc_round.tokens), + ) + for _cc_token in _cc_round.tokens: + tokens.append(_cc_token) + emit_token(_cc_token) + events.append({"step": step, "token": _cc_token}) + repetition_result = _trim_repeated_suffix( + tokens, repetition_config, repetition_long_cycle + ) + if repetition_result is not None: + events.append( + { + "step": step, + "repetition_stop": { + "reason": repetition_result.reason, + "block_tokens": repetition_result.block_tokens, + "repeats": repetition_result.repeats, + "trimmed_tokens": repetition_result.repeated_tokens, + }, + } + ) + break + if repetition_result is not None or _cc_round.stopped: + break + logits, _cc_forced = _cc_round.logits, _cc_round.correction + continue + started = time.perf_counter() profiler = _ar_forward_profiler(step) with attention_phase("ar_decode"): @@ -9538,6 +9610,8 @@ def emit_token(token: int) -> None: verify_calls += 1 logits = logits_next[:, -1, :] + if _cc_ar is not None: + events.append({"context_copy_ar": _cc_ar.summary()}) if token_callback is not None and _stream_gate.window > 0: # Armed-stream reconcile (F35): the trim decision is known here — # flush the held tail in full (no trim) or the post-trim remainder. diff --git a/tests/test_context_copy_ar.py b/tests/test_context_copy_ar.py new file mode 100644 index 000000000..5c4fc7e19 --- /dev/null +++ b/tests/test_context_copy_ar.py @@ -0,0 +1,395 @@ +"""Context-copy for AR-only runtimes: drafter, acceptance, KV trim, end to end. + +CPU only: a tiny stub model with a real (list-backed) KV cache drives +``generate_ar``. The stub's next token depends on the position and the last +token, so a wrong cache trim changes the stream. +""" + +from __future__ import annotations + +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest + +from mtplx.context_copy import CopyGovernor, context_copy_ar_enabled +from mtplx.context_copy_ar import ( + ArContextCopy, + accepted_prefix_greedy, + copy_ar_eligible, + verify_copy_block, +) +from mtplx.generation import generate_ar +from mtplx.mtp_patch import MTPContract +from mtplx.runtime import MTPLXRuntime +from mtplx.sampling import SamplerConfig + +VOCAB = 8 +GREEDY = SamplerConfig(temperature=0.0, top_p=1.0, top_k=0) + + +@pytest.fixture(autouse=True) +def _cpu_and_env(monkeypatch): + mx.set_default_device(mx.cpu) + monkeypatch.setenv("MTPLX_CONFIG", "/nonexistent") + for name in ( + "MTPLX_CONTEXT_COPY", + "MTPLX_CONTEXT_COPY_AR", + "MTPLX_CONTEXT_COPY_K", + "MTPLX_CONTEXT_COPY_PROBATION_K", + "MTPLX_CONTEXT_COPY_NGMIN", + "MTPLX_CONTEXT_COPY_NGMAX", + "MTPLX_CONTEXT_COPY_MINEXT", + "MTPLX_ASYNC_AR", + "MTPLX_EVAL_AUDIT", + ): + monkeypatch.delenv(name, raising=False) + yield + mx.set_default_device(mx.gpu) + + +class ListCache: + """Minimal trimmable KV cache: it stores the token ids it has seen.""" + + def __init__(self): + self.tokens: list[int] = [] + + @property + def offset(self) -> int: + return len(self.tokens) + + def trim(self, n: int) -> int: + n = min(int(n), len(self.tokens)) + del self.tokens[len(self.tokens) - n :] + return n + + +class Tokenizer: + def decode(self, tokens, **_kwargs): + return "".join(str(int(t)) for t in tokens) + + +def stub_next(position: int, last: int) -> int: + """Next token given the position of ``last`` and ``last`` itself.""" + return (last + 1 + (1 if position % 5 == 0 else 0)) % VOCAB + + +class DeterministicModel: + """Greedy-peaked logits from stub_next; reads positions from the cache.""" + + def __init__(self): + self.row_calls: list[int] = [] + + def make_cache(self): + return [ListCache()] + + def __call__(self, input_ids, *, cache=None, **_kwargs): + ids = [int(t) for t in np.asarray(input_ids).reshape(-1)] + entry = cache[0] + base = entry.offset + entry.tokens.extend(ids) + self.row_calls.append(len(ids)) + rows = [] + for j, tok in enumerate(ids): + row = [0.0] * VOCAB + row[stub_next(base + j, tok)] = 10.0 + rows.append(row) + return mx.array([rows], dtype=mx.float32) + + +class IidModel(DeterministicModel): + """Context-free next-token distribution, to test sampled exactness.""" + + LOGITS = (1.2, 0.3, 0.0, -0.5) + + def __call__(self, input_ids, *, cache=None, **_kwargs): + ids = np.asarray(input_ids).reshape(-1) + cache[0].tokens.extend(int(t) for t in ids) + return mx.array([[list(self.LOGITS)] * len(ids)], dtype=mx.float32) + + +def runtime(model) -> MTPLXRuntime: + return MTPLXRuntime( + model=model, + tokenizer=Tokenizer(), + model_path=Path("tiny-ar"), + mtp_enabled=False, + contract=MTPContract(), + ) + + +def trajectory(first: int, start_pos: int, n: int) -> list[int]: + out = [first] + for i in range(n - 1): + out.append(stub_next(start_pos + i, out[-1])) + return out + + +def run(prompt, max_tokens, *, stop=None, sampler=GREEDY, seed=0, model=None): + model = model or DeterministicModel() + out = generate_ar( + runtime(model), + list(prompt), + max_tokens=max_tokens, + sampler=sampler, + seed=seed, + stop_token_ids=stop or set(), + ) + return out, model + + +def copy_summary(out): + for event in out.stats.events: + if "context_copy_ar" in event: + return event["context_copy_ar"] + return None + + +def enable(monkeypatch, ngmin="3", ngmax="5"): + monkeypatch.setenv("MTPLX_CONTEXT_COPY_AR", "1") + monkeypatch.setenv("MTPLX_CONTEXT_COPY_NGMIN", ngmin) + monkeypatch.setenv("MTPLX_CONTEXT_COPY_NGMAX", ngmax) + + +# --- switches ------------------------------------------------------------ + + +def test_ar_switch_is_off_by_default_and_respects_global_off(monkeypatch): + assert context_copy_ar_enabled() is False + monkeypatch.setenv("MTPLX_CONTEXT_COPY_AR", "1") + assert context_copy_ar_enabled() is True + monkeypatch.setenv("MTPLX_CONTEXT_COPY", "0") + assert context_copy_ar_enabled() is False + + +def test_eligibility_excludes_mtp_penalties_guards_async_and_untrimmable(monkeypatch): + enable(monkeypatch) + cache = [ListCache()] + base = { + "mtp_enabled": False, + "penalized": False, + "guarded": False, + "sync_eval": True, + "cache": cache, + } + assert copy_ar_eligible(**base) + for key, value in ( + ("mtp_enabled", True), + ("penalized", True), + ("guarded", True), + ("sync_eval", False), + ("cache", []), + ("cache", [object()]), + ): + assert not copy_ar_eligible(**{**base, key: value}) + + +# --- governor (shared policy) --------------------------------------------- + + +def test_governor_probation_then_suspend_with_doubling_backoff(): + gov = CopyGovernor() + assert gov.block_cap(24, 8) == 8 + gov.record(8, 8, 10) + gov.record(8, 8, 20) + assert gov.block_cap(24, 8) == 24 + for n in range(3): + gov.record(0, 8, 100 + n) + assert gov.suspensions == 1 and gov.suspended(150) and not gov.suspended(100 + 2 + 64) + assert gov.backoff == 128 and gov.seen == 0 and gov.block_cap(24, 8) == 8 + + +# --- drafter --------------------------------------------------------------- + + +def test_drafter_copies_prompt_continuation_and_never_generated_text(monkeypatch): + enable(monkeypatch, "3", "4") + prompt = [1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 9, 9] + drafter = ArContextCopy(prompt) + assert drafter.propose([8, 5, 6, 7], 10)[:4] == [0, 1, 2, 3] + assert drafter.propose([8, 5, 6, 7], 2) == [0, 1] # budget cap + assert drafter.propose([8, 5, 6, 7], 0) == [] + assert drafter.propose([4, 5, 7, 7, 7], 10) == [] # no match + # a gram that only recurs inside generated text is not a prompt match + assert drafter.propose([9, 8, 7, 9, 8, 7, 9, 8, 7], 10) == [] + + +def test_drafter_block_length_follows_ladder_and_probation(monkeypatch): + enable(monkeypatch, "2", "2") + prompt = list(range(40)) + drafter = ArContextCopy(prompt) + assert len(drafter.propose([100, 10, 11], 50)) == 8 # ladder rung 0, probation + drafter.governor.seen, drafter.governor.ema = 5, 0.9 + assert len(drafter.propose([100, 10, 11], 50)) == 8 # ext 0 -> still rung 0 + + +# --- greedy acceptance + KV trim ------------------------------------------ + + +def test_accepted_prefix_greedy(): + assert accepted_prefix_greedy([1, 2, 3], [1, 2, 9, 4]) == 2 + assert accepted_prefix_greedy([1, 2, 3], [1, 2, 3, 4]) == 3 + assert accepted_prefix_greedy([1], [2, 5]) == 0 + + +def _verify(block, cache_tokens, *, stop=None): + model = DeterministicModel() + cache = [ListCache()] + cache[0].tokens.extend(cache_tokens) + token = stub_next(len(cache_tokens) - 1, cache_tokens[-1]) + result = verify_copy_block( + lambda ids: model(ids, cache=cache), + lambda keep: cache[0].trim(cache[0].offset - keep) >= 0, + cache[0].offset, + token, + block, + stop, + greedy=True, + ) + return result, cache[0], token + + +def test_verify_trims_cache_to_accepted_prefix_on_rejection(): + base = [3, 4, 5] + token = stub_next(2, 5) + good = trajectory(stub_next(3, token), 4, 2) # two tokens the model agrees with + block = [*good, (good[-1] + 4) % VOCAB, 1] # then a wrong one + result, cache, tok = _verify(block, base) + assert tok == token + assert result.tokens == good and result.correction is None and not result.stopped + assert cache.tokens == [*base, token, *good] + assert cache.offset == len(base) + 1 + len(good) + + +def test_verify_full_accept_keeps_whole_block_and_returns_next_logits(): + base = [3, 4, 5] + token = stub_next(2, 5) + block = trajectory(stub_next(3, token), 4, 4) + result, cache, _ = _verify(block, base) + assert result.tokens == block + assert cache.offset == len(base) + 1 + len(block) + nxt = int(mx.argmax(result.logits[0]).item()) + assert nxt == stub_next(len(base) + len(block), block[-1]) + + +def test_verify_stop_inside_block_cuts_and_keeps_stop_out_of_cache(): + base = [3, 4, 5] + token = stub_next(2, 5) + block = trajectory(stub_next(3, token), 4, 5) + result, cache, _ = _verify(block, base, stop={block[2]}) + assert result.stopped and result.tokens == block[:3] + assert cache.offset == len(base) + 1 + 2 # token + 2 drafts, stop excluded + + +# --- end to end vs plain AR ------------------------------------------------ + + +def _prompt_with_repeat(): + # the model's own trajectory twice, so the tail keeps matching the prompt + first = trajectory(1, 0, 24) + return first + first[:6] + + +def test_greedy_stream_matches_plain_ar_with_fewer_forwards(monkeypatch): + prompt = _prompt_with_repeat() + plain, plain_model = run(prompt, 40) + enable(monkeypatch) + copied, copy_model = run(prompt, 40) + assert copied.tokens == plain.tokens + summary = copy_summary(copied) + assert summary["rounds"] > 0 and summary["accepted_tokens"] > 0 + assert len(copy_model.row_calls) < len(plain_model.row_calls) + assert any(n > 1 for n in copy_model.row_calls[1:]) + + +def test_greedy_with_diverging_prompt_still_matches_plain_ar(monkeypatch): + base = trajectory(1, 0, 30) + prompt = base[:12] + [(t + 3) % VOCAB for t in base[12:18]] + base[18:] + base[:5] + plain, _ = run(prompt, 60) + enable(monkeypatch, "2", "4") + copied, _ = run(prompt, 60) + assert copied.tokens == plain.tokens + assert copy_summary(copied)["rounds"] > 0 + + +@pytest.mark.parametrize("max_tokens", [1, 2, 3, 7, 16]) +def test_max_tokens_is_exact(monkeypatch, max_tokens): + prompt = _prompt_with_repeat() + plain, _ = run(prompt, max_tokens) + enable(monkeypatch) + copied, _ = run(prompt, max_tokens) + assert len(copied.tokens) == max_tokens + assert copied.tokens == plain.tokens + assert copied.finish_reason == plain.finish_reason + + +def test_stop_token_inside_accepted_block_ends_stream_like_plain_ar(monkeypatch): + prompt = _prompt_with_repeat() + plain_full, _ = run(prompt, 40) + stop = {plain_full.tokens[9]} + plain, _ = run(prompt, 40, stop=stop) + enable(monkeypatch) + copied, _ = run(prompt, 40, stop=stop) + assert copied.tokens == plain.tokens + assert copied.tokens[-1] in stop + assert copied.finish_reason == plain.finish_reason + + +def test_disabled_by_default_runs_plain_loop(): + prompt = _prompt_with_repeat() + out, model = run(prompt, 20) + assert copy_summary(out) is None + assert all(n == 1 for n in model.row_calls[1:]) + + +# --- sampled exactness ------------------------------------------------------ + + +def _softmax(logits): + z = np.exp(np.array(logits) - np.max(logits)) + return z / z.sum() + + +def _chi2(counts, probs): + counts = np.asarray(counts, dtype=float) + expected = probs * counts.sum() + return float(((counts - expected) ** 2 / expected).sum()) + + +def test_sampled_copy_blocks_follow_the_target_distribution(monkeypatch): + """Copy rounds fire constantly (prompt over the 4 live tokens, 2-gram keys), + yet the emitted stream must stay i.i.d. from the target distribution.""" + sampler = SamplerConfig(temperature=1.0, top_p=1.0, top_k=0) + probs = _softmax(IidModel.LOGITS) + rng = np.random.default_rng(7) + prompt = [int(t) for t in rng.integers(0, 4, size=64)] + enable(monkeypatch, "2", "3") + singles = np.zeros(4) + pairs = np.zeros((4, 4)) + rounds = accepted = drafted = 0 + for seed in range(400): + out, _ = run(prompt, 24, sampler=sampler, seed=seed, model=IidModel()) + toks = out.tokens + assert len(toks) == 24 + for t in toks: + singles[t] += 1 + for a in range(0, 24, 2): + pairs[toks[a], toks[a + 1]] += 1 + summary = copy_summary(out) + rounds += summary["rounds"] + accepted += summary["accepted_tokens"] + drafted += summary["drafted_tokens"] + assert rounds > 400 and 0 < accepted < drafted # blocks fired, some accepted, some rejected + # 3 dof / 15 dof; thresholds at p = 0.001 (fixed seeds, so deterministic) + assert _chi2(singles, probs) < 16.3 + assert _chi2(pairs.ravel(), np.outer(probs, probs).ravel()) < 37.7 + + +def test_sampled_run_is_reproducible_for_a_seed(monkeypatch): + sampler = SamplerConfig(temperature=1.0, top_p=0.9, top_k=3) + prompt = [int(t) for t in np.random.default_rng(3).integers(0, 4, size=48)] + enable(monkeypatch, "2", "3") + a, _ = run(prompt, 20, sampler=sampler, seed=5, model=IidModel()) + b, _ = run(prompt, 20, sampler=sampler, seed=5, model=IidModel()) + assert a.tokens == b.tokens From ee63be4d461ca37fc5bce3d7fac57372c07f1538 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Sun, 4 Oct 2026 02:45:00 +0200 Subject: [PATCH 2/2] docs: CHANGELOG line for MTPLX_CONTEXT_COPY_AR with measured numbers Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..88e581b94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable user-facing changes to MTPLX. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow [Semantic Versioning](https://semver.org/). +## [Unreleased] + +### Added + +- **Context-copy drafting for runtimes without an MTP head, opt-in.** `MTPLX_CONTEXT_COPY_AR=1` lets `generate_ar` copy a block from the prompt when the stream matches a prompt n-gram, and verify it in one forward pass (default off; same n-gram index, block ladder and probation/suspend policy as the MTP lanes). Measured by @jvmenen on an M5 Pro 64 GB (automatic power mode, thermal pressure 0), Llama 3.1 8B Instruct 4-bit, in process, median of 5 runs, October 4, commit f1d7d1e1: decode tok/s 55.3 to 289.3 (5.2x) returning an edited 2,458-token Python file, 52.4 to 69.2 (1.3x) quoting from a 4,075-token report, 59.5 to 59.3 on free prose (no matches, no copy rounds); greedy streams identical to the one-token loop in 20 of 20 runs; peak memory unchanged. + ## [2.12.2] - 2026-10-03 ### Fixed