Context-copy for AR-only runtimes (MTPLX_CONTEXT_COPY_AR, opt-in)

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

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.
 

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),
+    )