Feature request: prompt-lookup (n-gram) drafting as an alternative/hybrid draft source

[ashalliants]

Summary

Feature request: support prompt-lookup / n-gram drafting as an alternative (or supplementary) draft source alongside the DFlash block-diffusion drafter.

Motivation

For edit-heavy / agentic coding workloads (the model re-emits file contents with small changes, multi-turn sessions where output overlaps prior context), a large fraction of generated tokens are verbatim copies of tokens already in context. Prompt-lookup decoding (PLD) drafts those spans for free with simple n-gram matching — no draft model forward passes at all — and typically reaches very high acceptance inside copy regions (vLLM ships this as the [ngram] speculative method; llama.cpp has an equivalent lookup-decoding mode).

Why dflash-mlx is uniquely positioned to do this

I tried to get ngram drafting working on Qwen3.6 via the upstream mlx-lm speculative path (ml-explore/mlx-lm#1297) and it is architecturally blocked there: the generic verifier requires trimmable caches, and the Qwen3.5/3.6 family's gated-delta linear-attention layers use an untrimmable ArraysCache:

ValueError: Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'}).

dflash-mlx already solved the hard part — verification with partial-acceptance rollback over these hybrid-cache models. Swapping (or augmenting) the draft source should be much smaller work than what already exists: an ngram lookup produces a candidate block, and the existing verify path consumes it unchanged.

Evidence there's headroom

Measured on M5 / 32 GB, Qwen3.6-35B-A3B-4bit + z-lab/Qwen3.6-35B-A3B-DFlash, greedy, ~3K-token generations on a code-edit prompt (re-emit a ~100-line file with a small refactor):

• DFlash: 67.0 tok/s @ 74.9% acceptance (code-edit), 64.2 tok/s @ 70.9% (fresh code)
• plain decode: ~54.8 tok/s

74.9% acceptance on a workload where long spans are exact copies suggests the diffusion drafter leaves accuracy on the table precisely where ngram lookup is near-perfect. A hybrid policy could be simple and conservative, e.g.:

1. try n-gram match against context (cheap, exact);
2. if a match of length ≥ k exists, emit the lookup continuation as the draft block (zero draft-model cost);
3. otherwise fall back to the DFlash drafter as today.

That keeps DFlash's behaviour on novel text and upgrades copy regions to near-1.0 acceptance with lower draft cost. It would also benefit models that have no trained DFlash drafter yet, since ngram drafting needs no sidecar model at all.

Possible knobs

• --draft-source {dflash,ngram,hybrid} (default dflash, current behaviour)
• min match length / max draft length for the ngram path (vLLM's prompt_lookup_min experience suggests too-short matches can hurt structured/tool-call output, so a conservative default like ≥ 4–8 matters)

Happy to share benchmark scripts/prompts or run comparisons on Apple Silicon (M5, 32 GB) if that's useful.

[ashalliants]

Correction + measurements after reading the source more carefully: dflash-mlx already ships a prompt-lookup draft source — engine/copyspec.py (CopySpecIndex), tried before the DFlash drafter each cycle with draft_source = "copyspec". I missed it when filing this; apologies for the redundant framing.

So let me refine this request based on what copyspec actually does today, with measurements.

Measured: copyspec barely engages on a copy-heavy workload

Same setup as above (M5 32 GB, Qwen3.6-35B-A3B-4bit, code-edit prompt asking the model to re-emit a ~100-line file with a small refactor), instrumented via SummaryEvent.copyspec_hits / copyspec_tokens:

run	tok/s	acceptance	tokens/cycle	copyspec
first 1024 tokens (mostly reasoning)	102.7	88.4%	8.61	1 hit / 15 tokens (1.5%)
full 3072 tokens (incl. file re-emission)	68.6	74.9%	3.98	1 hit / 15 tokens (0.5%)

The second half of the generation — exactly the region dominated by verbatim copying of the input file — is where throughput drops by a third, tokens/cycle halves, and the adaptive verifier shortened blocks 11 times. Copyspec contributed a single 15-token block across the whole run.

Why it misses (from reading copyspec.py)

1. Full-block-or-nothing: draft_after returns None when token_count < max_tokens — a match must supply the entire block (typically 16 tokens) or it's discarded. Edit workloads are full of 6–15-token copy spans (copy runs break at each changed line), all of which currently get zero benefit.
2. Fixed 6-token exact-match window: COPYSPEC_WINDOW_SIZE = 6 is a hardcoded constant — no way to trade match precision vs. hit rate per workload.
3. No observability: copyspec_hits/copyspec_tokens exist on SummaryEvent but neither the CLI summary line nor the server logs surface them, so users can't tell whether copyspec is helping (I benchmarked this engine for a day without knowing it existed).

Refined request

1. Allow partial-block copy drafts — splice a shorter copyspec match with mask-token tail (or DFlash-drafted tail) instead of discarding it.
2. Expose window size / minimum match length as CLI/server flags.
3. Print copyspec stats in the generate summary line and server logs.

I'm happy to take a swing at a PR for some or all of these (starting with 1 + 3, benchmarked before/after on the workload above) if you'd be open to it — let me know if there's prior context on why partial matches are rejected (e.g. acceptance-quality regressions) before I start.

[ashalliants]

Follow-up with a root-cause: I replayed the captured token stream from the benchmark above through CopySpecIndex offline (exact shipped class, engine-like non-overlapping walk: draft at cursor, commit correct prefix + 1, jump). The full-block rule and window size turn out to be secondary — the one-strike kill switch is what's suppressing copyspec.

What happens on the edit-task run

At generation position 33 of 3072 (early reasoning phase), a coincidental 6-gram match produced a zero-acceptance copy block, and state.copyspec_disabled = True permanently disabled copyspec for the remaining 99% of the generation — during which roughly half of all tokens were verbatim copies sitting in the index.

Replay numbers (same run: Qwen3.6-35B-A3B-4bit, code-edit prompt, 3072 gen tokens)

policy	copy drafts	correctly drafted tokens	% of generation
current (one-strike kill)	3	10	0.3%
no kill switch, w=6	200	1,581	51.5%
no kill switch, w=8	152	1,468	47.8% (zero-acc blocks: 41 → 17)
no kill switch, w=10	125	1,396	45.4% (zero-acc: 12)

Match quality without the kill switch is high: of 1,781 per-position matches, 1,116 were full 15-token correct blocks (median correct prefix = 15). The dud rate at w=6 is 41 zero-acceptance blocks out of 200 drafts — each costs one verify slot but still commits its staged token, so the downside is bounded; permanent disablement to avoid it forfeits ~1,500 free tokens on this workload.

Suggested fix shape

1. Replace the permanent disable with tolerance — e.g. disable only after K consecutive zero-acceptance copy blocks (reset on any accepted copy), or a cooldown that re-enables after N cycles. Either is a few lines in spec_epoch.py.
2. Surface copyspec_hits / copyspec_tokens in the generate summary line and server stats so this is visible to users.
3. (Optional) expose the window size — w=8 cut dud blocks 2.4× for ~7% fewer useful tokens in the sweep above.

I'm putting together a PR for 1 + 2 with before/after benchmarks on the same workload and will link it here.

[ashalliants]

PR up: #36 — implements the cooldown gate (fix 1) + summary-line stats (fix 2) with before/after benchmarks: +16.8% on the edit workload (78.1 → 91.2 tok/s, copyspec 1 → 161 blocks), no regression on fresh code. Window-size knob left as a possible follow-up.

[ashalliants]

Full benchmark harness + captured token stream for the replay numbers above (and for verifying #36 independently): https://gist.github.com/ashalliants/91819d410f6822e406a314740c8b7d0e — copyspec_sim.py + edit_tokens.json reproduce the kill-switch-at-position-33 / 51.5%-copy-draftable analysis in seconds on CPU.

[bstnxbt]

Thanks for the detailed report, replay gist, and PR. The diagnosis is real: on the captured copy-heavy edit task, the existing one-strike copyspec disable suppresses almost all useful prompt-lookup drafting.

I merged #36 briefly, but reverted it from main in f77b1d9 after broader validation showed the cooldown gate is not safe enough yet across supported targets/prompts.

Validation notes from local Apple Silicon runs with explicit 120s cooldown:

• Qwen3.6-27B-4bit + DFlash, edit prompt, 3072 tokens:
  • before perf: replace copyspec one-strike disable with cooldown gate (+17% on edit workloads) #36: 109.2 / 119.1 tok/s, copyspec 0.7%, tokens/cycle 7.03
  • with perf: replace copyspec one-strike disable with cooldown gate (+17% on edit workloads) #36: 61.0 / 109.2 tok/s, copyspec 72.6%, tokens/cycle 6.24
  • result: too much regression risk for main.

Earlier diagnostic runs on Qwen3.6-35B-A3B showed the intended upside on the copy-heavy edit task, but also exposed prompt sensitivity:

• 35B edit 3072: positive, ~224.4 -> 247.4 tok/s avg, copyspec 0.5% -> 49.0%
• 35B fresh 3072: roughly neutral/slightly positive
• 35B README 1024: neutral
• 35B README 3072: negative in that diagnostic run

So the issue should stay open. The next fix needs a safer policy than “cooldown after misses” alone, probably gated by observed copy survival / tokens-per-cycle / model or prompt regime, or a cheaper probe path that does not let copyspec dominate when it lowers throughput.

For now main is back to the prior one-strike behavior to avoid shipping a regression to users while we narrow the policy.

[ashalliants]

Follow-up: why a safe automatic copyspec gate isn't achievable, + an opt-in flag

After the #36 revert I did a much deeper validation pass (M5 32GB, greedy, clean clock — calibrated against your 35B main reference of ~78 tok/s). Summary up front: #36 wasn't a tuning miss — there is no per-cycle gate that's safe across workloads. Copyspec is a workload-specific optimization, so the right thing to ship is an explicit opt-in, which I've done (PR linked below). The conservative one-strike default is justified.

The full picture (force-on ceiling vs one-strike default)

Forcing copyspec fully on vs off, on the two regimes:

target	copy-heavy prompt (edit)	non-copy prompt (fresh)
Qwen3.6-35B-A3B-4bit (MoE, expensive draft)	+~19–22% (77.7→94.7 tok/s via the shipped flag)	−4.5%
Qwen3.6-27B dense (cheap draft)	neutral / slightly negative	−10%

So copyspec helps copy-heavy prompts and hurts non-copy prompts on every model — the magnitude scales with how expensive the draft forward-pass is (big MoE → skipping it via a copied block is a large saving; cheap dense draft → the saving doesn't cover the cost). It is not a clean per-model property: enabling it model-wide would regress that model's own non-copy traffic.

Why no automatic gate works

Two gate designs failed in the same way: a benefit-comparison gate (v1) under-engaged (oscillated on noisy realized-tps and captured only +2.9% of the +18.9%); an aggressive sticky gate (v2) regressed the non-copy control −11%. Root cause, isolated with a fixed-block run:

• With the adaptive verifier turned off (fixed block size), force-on copyspec on the non-copy workload costs −0.2% (zero regression). With it on, −10%.
• So the regression is the adaptive block-size policy reacting to copy blocks: a low-acceptance copy block pushes the verifier toward small blocks, and that decision degrades subsequent dflash cycles. The copy blocks themselves, and the draft cache, are fine.

That's the trap: copyspec's true cost is a cross-cycle externality (its effect on future cycles), which a per-cycle copy-vs-dflash comparison structurally cannot observe. And every static treatment of copy cycles within the adaptive policy sits on a Pareto frontier — there's no safe-everywhere point:

copy cycles → adaptive policy	35B edit	27B fresh
feed normally	+18.9%	−10%
partially excluded (the #36-era isolation)	+13%	−9%
fully skipped	+8.4%	−4% (residual is irreducible cadence disruption — copy blocks commit multiple tokens, shifting where dflash cycles land; can't be undone without un-committing tokens)

Higher 35B engagement wants copy to feed the adaptive policy (high-acceptance copy keeps it in full-block mode); fresh wants copy invisible. Opposite requirements. Checkpoint/restore of adaptive state recovers only the policy half, not the cadence half.

Conclusion + what I'm proposing

Copyspec aggression is worth it only when you know the workload is copy-heavy and the target has an expensive draft pass. That's a deployment-time fact, not a runtime-detectable one. So: keep conservative as the safe default, and add an explicit opt-in.

PR: #41 — adds --copyspec-mode {conservative,aggressive,off} (env DFLASH_COPYSPEC_MODE) on serve/generate/benchmark, default conservative (unchanged one-strike behavior). aggressive keeps copyspec on (the +18.9% path for copy-heavy MoE workloads, documented as potentially regressing non-copy traffic); off disables it. Also restores the | copyspec N blocks / M tokens summary-line stat for visibility.

Full repro harness + captured token streams (updated, on-device — new numbers supersede my earlier offline-replay gist): https://gist.github.com/ashalliants/3b4425dcfed5d34eda375b6f560ed905 (the original offline harness remains at https://gist.github.com/ashalliants/91819d410f6822e406a314740c8b7d0e). Happy to fold any of this into docs if useful.

[ashalliants]

Update: a safe automatic gate is achievable — correcting my earlier conclusion

My previous comment argued no per-cycle gate could be safe because copyspec's cost is a cross-cycle externality. That part holds — but I gave up one experiment too early, and the conclusion was wrong. There is a safe auto-gate, and PR #41 now ships it as --copyspec-mode auto.

What I missed

The regression is entirely the adaptive block-size policy reacting to copy blocks — confirmed with a fixed-block run (adaptive off → copyspec force-on costs −0.2% on the non-copy workload, vs −10% with adaptive on). Crucially, that perturbation lives in the adaptive policy's heuristic state, which is resettable — it has no coherence dependency on committed tokens (unlike the KV/draft caches, which is the case I wrongly conflated it with). Decisive test: run copyspec for a few blocks, then disable it and rebuild the adaptive policy → the non-copy workload recovers to +0.0% (vs −9% if you disable without the reset). The damage isn't irreversible after all.

auto

A windowed A/B latch: periodically measure realized throughput copyspec-off (baseline) vs copyspec-on (probe), latch to whichever is faster, reset the adaptive policy on disengage, and back off re-probing exponentially when copyspec keeps losing. Measured (M5 32GB, greedy, tpc):

target	edit (copy-heavy)	fresh (non-copy)
35B-A3B-4bit (MoE)	+23.5%	−0.2% (neutral)
27B dense	+4.6%	−7.5%*

* one exploratory probe's cost on a short non-copy run on a cheap-draft target; amortizes on longer generations. On the MoE, auto is strictly better than forcing copyspec on (better on both edit and fresh).

So: default stays conservative (unchanged), and auto is safe to enable broadly — it captures the win where it exists and self-disables (cleanly) where it doesn't. PR #41 updated accordingly. Thanks for keeping the issue open.

[ashalliants]

Finalized in #41 — two refinements since the last update

After more validation, the shipped form in #41 is tightened in two ways:

1. Dropped the always-on aggressive mode. The surface is now --copyspec-mode {conservative,off,auto} (default conservative, unchanged). aggressive was auto-dominated on realistic mixed workloads and a footgun when misapplied; auto captures the win without requiring the user to know their workload, so there's no reason to keep the sharp edge.

2. The 27B-dense-fresh number is a short-generation artifact, not a standing regression. On a cheap-draft target copyspec can't win, so auto's single exploratory probe dominates a short run (−7.5% tpc at 1536 tokens). The exponential re-probe backoff amortizes it as the generation grows — at 4096 tokens the same workload is +2.9% (neutral/positive). So on realistic-length generations auto is safe everywhere; I'm not shipping a terminal self-disable precisely because it would reintroduce perf: replace copyspec one-strike disable with cooldown gate (+17% on edit workloads) #36's premature-give-up (quitting during a reasoning preamble before the copy-heavy phase).

Net: auto is the recommended "on" setting — big win on strong-copyspec targets (35B-A3B edit: 77.7 → 93.0 tok/s, +19.7%), neutral elsewhere on real workloads — and conservative remains the safe default. Full numbers + methodology (Apple M5 / 32 GB) and the updated repro gist are in #41.