[jvmenen]

Why

Models without an MTP head, such as Llama 3.1 8B or Qwen3-8B, generate one token per forward pass, also when the answer is mostly text that is already in the prompt: rewriting a file with a small change, extracting passages from a document, quoting. Models with an MTP head already skip through such copied text with context-copy; models without one do not get that speedup at all. On Llama 3.1 8B a full-file code edit runs at 56 tokens per second, even though most of the file is copied unchanged.

How it works, in plain terms

Context-copy looks up where the last few generated tokens also occur in the prompt, and proposes the text that follows there as a block of guesses. The model then checks the whole block in one forward pass instead of one pass per token, and keeps the part that matches what it would have produced. With greedy decoding that is exactly the same text as without context-copy. With sampling, a correction at the first rejected token keeps the output distribution the same. The cache entries of rejected tokens are removed again, so the next step starts from the accepted text. This PR adds that loop to the decode path for models without an MTP head, behind MTPLX_CONTEXT_COPY_AR=1 (off by default). It does not run with penalties, constrained decoding or guards, and hybrid models whose recurrent state cannot be rolled back keep plain decoding. On the full-file edit above, decoding goes from 56 to 275 tokens per second (4.9x); extraction from a report gains 1.3x, and prose without source text is unchanged.

What changes

The same lane in generate_ar, behind MTPLX_CONTEXT_COPY_AR=1 (off by default, and still subject to MTPLX_CONTEXT_COPY).

• New mtplx/context_copy_ar.py. One forward over [token, *block], then acceptance and cleanup:
  • Greedy: the exact-match prefix is accepted.
  • Sampled: point-mass acceptance with a residual correction at the first rejection, so the output law is the target distribution.
  • The KV cache is trimmed to the accepted prefix.
  • A stop token inside the block ends the stream.
  • The block is capped so max_tokens is exact.
• Reused from mtplx/context_copy.py. NgramIndex, the block drafter and the MTPLX_CONTEXT_COPY_* knobs. The probation and suspend policy that the MTP loops keep inline is extracted as CopyGovernor; the MTP lanes themselves are untouched.
• generation.py. About 60 lines in generate_ar.
• Not used with: MTP runtimes, presence/frequency penalties, constrained decoding, loop or thinking guards, hidden-state capture, async AR, or caches without trim. Hybrid recurrent models fall back to plain AR.
• CHANGELOG line under [Unreleased].

Results

M5 Pro 64 GB, mlx-community/Llama-3.1-8B-Instruct-4bit loaded with mtp=False, in process. Off and on interleaved, 5 runs per cell; decode tok/s, median:

Workload	Temp	off	on	Ratio	Acceptance
code edit, full file (1,301 tokens, ends at stop)	0	56.0	275.2	4.9×	0.86
code edit, full file	1.0	55.5	269.1	4.9×	0.87
extraction from a 4,075-token report	0	52.4	69.2	1.3×	0.56
prose without source text	0	59.5	59.3	1.0×	no copy rounds

• All 20 greedy streams are token-identical to the plain loop.
• Peak memory is identical off and on.
• On prose the lane probes at every token but never finds a block, at no measurable cost.

Tests

tests/test_context_copy_ar.py (CPU):

• switches and eligibility
• governor and drafter
• KV trim on rejection, on full accept and on a stop inside the block
• stream equality against plain AR with a position-dependent stub, including max_tokens 1 to 16 and stop tokens
• fewer forwards than plain AR
• a chi-square check on a toy model that the sampled output follows the target distribution

The existing context-copy tests are unchanged and pass.

Notes

• Off by default.
• Open question: whether to move the MTP loops' inline copies of the governor logic onto CopyGovernor. That is left out here to keep the MTP lanes free of risk.

🤖 Generated with Claude Code