Skip to content

speculative: stop plain drafter+target falling back to KV checkpoints - #173

Open
bri-prism wants to merge 3 commits into
prismfrom
fix/spec-partial-seq-removal
Open

bri-prism wants to merge 3 commits into
prismfrom
fix/spec-partial-seq-removal

Conversation

@bri-prism

Copy link
Copy Markdown

Plain drafter+target speculative decoding on hybrid (recurrent + attention) models is classified as unable to do partial sequence removal, and falls back to a whole-KV checkpoint save/restore every round. On a 27B qwen35 that is 74–77% of each round — about 0.5 s per drafted token, ~15x the drafter's own compute — which makes speculation roughly 2x slower than plain decoding no matter how good the drafts are.

Two causes, both required:

1. DRAFT_SIMPLE is missing from need_n_rs_seq() (common/common.h)

The list covers DRAFT_MTP, DRAFT_EAGLE3, DRAFT_DFLASH and DRAFT_DSPARK, so the ordinary drafter+target case gets n_rs_seq = 0. common_context_can_seq_rm() then fails its llama_n_rs_seq() > 0 test, llama_memory_seq_rm() fails on the hybrid memory, and the context is classified COMMON_CONTEXT_SEQ_RM_TYPE_FULL.

LLM_ARCH_QWEN35 is already in llm_arch_supports_rs_rollback() (src/llama-arch.cpp), so the architecture was never the blocker.

2. The draft context explicitly zeroes cparams.n_rs_seq (common/speculative.cpp)

It has to match the target. On partial acceptance both contexts drop the rejected suffix, and a draft context at n_rs_seq = 0 silently fails its seq_rm and retains stale positions, which then trips the M-RoPE monotonic-position check on the next batch (seq_pos_max == batch seq_pos_min). Traced: tgt_ok=1 dft_ok=0 with tgt_max=10 against dft_max=12. The existing comment there already wondered whether this should be set.

No change needed in speculative-simple.cpp

It already removes the rejected draft tokens unconditionally after the use_ckpt_tgt block (seq_rm(seq_id, n_past, -1) on both contexts), and that path is always reached when use_ckpt_tgt is false — the only continue sits inside the checkpoint branch. An earlier version of this fix added a second removal there; ablating it produced identical results (194 tokens, 33.190% acceptance, no abort), so it was dropped.

Measured

Arc B390, 2B-PTQ1_0 drafting 27B-PTQ1_0, greedy, against a 4.43 t/s llama-bench decode-only baseline:

k before after
1 1.137x / 1.356x (two prompts)
2 0.482x 1.085x
4 0.535x 0.727x

Beyond the speedup, the cost model mean_len / (C(k+1) + k·draft_cost) now predicts measured throughput to within 0.96–1.08x across 8 runs (k=1..4 × 2 prompts); before the fix the same expression was off by 3.66–3.98x.

Correctness: generated text is byte-identical to a k=1 reference over the common prefix (the only difference is where the token limit truncates, since blocks emit at different boundaries).

Note on reported acceptance rates

For hybrid models that previously took the checkpoint path, the reported acceptance rate changes meaning. That path continues before n_drafted/n_accept are incremented, so partial-acceptance rounds were never counted and the rate read 100% by construction. With partial rounds counted it is 14–87% depending on k and workload. Any historical 100% figure from a hybrid model is that artifact, not perfect acceptance.

Scope and limitations

Exercised on one architecture (qwen35 hybrid), one backend (Vulkan), one drafter/target pair. Reviewers on other arches or backends should treat the cross-arch behaviour as unverified.

Pure-attention models are unaffected: they already had use_ckpt_tgt == false and always took the pre-existing removal path, so there is no behaviour change for them.

On hybrid (recurrent + attention) models, plain drafter+target speculative
decoding was classified as unable to do partial sequence removal, and fell
back to a whole-KV checkpoint save/restore every round. On a 27B qwen35 that
was 74-77% of each round -- roughly 0.5 s per drafted token, ~15x the
drafter's own compute -- which made speculation about 2x slower than plain
decoding regardless of draft quality.

Two causes, both needed:

1. need_n_rs_seq() listed DRAFT_MTP, DRAFT_EAGLE3, DRAFT_DFLASH and
   DRAFT_DSPARK but not DRAFT_SIMPLE, so the ordinary drafter+target case got
   n_rs_seq = 0. common_context_can_seq_rm() then failed its
   llama_n_rs_seq() > 0 test, llama_memory_seq_rm() failed on the hybrid
   memory, and the context was classified COMMON_CONTEXT_SEQ_RM_TYPE_FULL.
   LLM_ARCH_QWEN35 is already in llm_arch_supports_rs_rollback(), so the
   architecture was never the blocker.

2. The draft context explicitly zeroed cparams.n_rs_seq. It has to match the
   target: on partial acceptance both contexts drop the rejected suffix, and a
   draft context at n_rs_seq = 0 silently fails its seq_rm and keeps stale
   positions, which then trips the M-RoPE monotonic-position check on the next
   batch (seq_pos_max == batch seq_pos_min). Traced: tgt_ok=1 dft_ok=0 with
   tgt_max=10 against dft_max=12.

No change is needed in speculative-simple.cpp -- it already removes the
rejected draft tokens unconditionally after the use_ckpt_tgt block, and that
path is always reached when use_ckpt_tgt is false. An earlier version of this
fix added a second removal there; ablating it gave identical results, so it
was dropped.

Measured on an Arc B390, 2B-PTQ1_0 drafting 27B-PTQ1_0, greedy, against a
4.43 t/s llama-bench decode-only baseline:

  k=1   1.137x / 1.356x   (two prompts)
  k=2   0.482x -> 1.085x
  k=4   0.535x -> 0.727x

Generated text is byte-identical to a k=1 reference over the common prefix.

Also note the reported acceptance rate changes meaning for hybrid models that
previously took the checkpoint path: that path continues before n_drafted and
n_accept are incremented, so partial-acceptance rounds were never counted and
the rate read 100%. With partial rounds counted it is 14-87% depending on k
and workload.
@bri-prism

Copy link
Copy Markdown
Author

Second-backend evidence for the mechanism, from Metal. The deterministic half reproduces cleanly; the throughput half of my run is not usable and I am not quoting it.

Setup. M5 Pro, Metal. Target preview-27B-PQ2_0 (qwen35), drafter latest-4B-PQ2_0, --spec-type draft-simple, --spec-draft-n-max 2, greedy, -s 1234, same prompt, -n 128. Arm A is 6ac5eb087 unpatched, already built. Arm B is this PR's head b9b74a327, built with the same flags read out of arm A's CMakeCache.txt (Release, GGML_METAL=ON, GGML_BLAS=ON, LLAMA_CURL=OFF). Three interleaved reps per arm, A,B,A,B,A,B.

The fallback is eliminated on Metal. Arm A prints speculative decoding will use checkpoints (context does not support partial sequence removal) at startup; arm B does not. 3/3 reps each way. That is a different backend, a different target instance and a different drafter pair from the B390 work, so the mechanism is not Vulkan-specific.

The 100% acceptance figure is an artifact, confirmed independently. Identical in all three reps per arm:

arm A (unpatched) arm B (this PR)
n_drafted 82 141
n_accept 82 58
accept 100.000% 41.135%

Same prompt, same seed. The 100% is the checkpoint path miscounting, not acceptance. This matters beyond this PR: any acceptance figure taken from a hybrid target on a pre-fix build is an artifact, and any C(n) or tau derived from one inherits the error.

What I am not claiming. My throughput numbers are thermally contaminated and say nothing either way:

        rep1     rep2     rep3
arm A  18.280   18.226   10.831   t/s
arm B  17.809   12.096   12.936

Arm A spreads 40.7% across three reps of a byte-identical run, so the between-arm difference is far inside the noise. A usable throughput comparison needs cooldown between legs and fresh-leg comparison, which I have not done. The deterministic counters above were stable across every rep, which is why I am reporting those and not these.

This run also cannot answer whether speculation pays on Metal. It is patched-vs-unpatched, not spec-vs-autoregressive; there is no matched AR baseline here.

Consistency check against an independent Metal measurement. A separate measurement on the same target gives C(n) = 1.635 / 2.066 / 2.518 / 3.089 for 2/3/4/5 verify rows. k=2 is 3 verify rows, so the bar is C(3) = 2.066. Arm B drafts 141 and accepts 58 over roughly 70 rounds, so mean emitted is about 1.82 per round, below the bar, predicting a net-negative configuration. Consistent with seeing no gain here, and an argument that this particular pair and depth was never going to show one.

Caveat carried from that work and applicable here: the 4B was not trained as a drafter for this 27B, so 41.135% is not this family's acceptance.

Verification hygiene, since a stale binary produced a false confirmation elsewhere this week: arm B's binary mtime was 10:52:05 against a 10:52:14 wall clock at first use, with zero build errors. An earlier attempt at this A/B was invalid and discarded: I omitted --spec-type, leaving params.speculative.types empty, and need_n_rs_seq() is std::any_of over that list, so it returns false regardless of the fix and both arms printed the checkpoint line. Worth knowing for anyone reproducing this.

@bri-prism

Copy link
Copy Markdown
Author

Correction to my earlier comment, plus the separate "does speculation pay" measurement it was missing. Posting as a new comment rather than editing, so what I claimed and what I withdrew both stay visible.

Withdrawn: the C(3) cross-check

My previous comment used an independently measured C(3) = 2.066 as the bar for this configuration. That was wrong. C(n) is not drafter-invariant — the drafter's forward pass happens inside the round, so it is inside C. That 2.066 was measured with a 700 MB drafter; this run uses latest-4B-PQ2_0. A four-point fit of the same series decomposes as roughly 0.481 AR-steps per verify row plus a 0.642 AR-step fixed term, and the fixed term is almost entirely the drafter. With a 4B in that slot the intercept grows while the slope stays put, so the real bar here is meaningfully above 2.066 and the predicted ratio is worse than the 0.88x I quoted.

The direction of the conclusion is unchanged and in fact strengthened. The number should not be quoted.

Narrowed: the 100% acceptance artifact

I wrote that any acceptance figure from a hybrid target on a pre-fix build is an artifact. Too broad. It applies to the speculative types that were missing from need_n_rs_seq(), which is what this PR adds DRAFT_SIMPLE to. A draft-dspark run on the same target and the same unpatched build reports ordinary acceptance (42.9 / 44.4 / 25.8 / 21.5% across four prompt classes), because DRAFT_DSPARK was already in that predicate.

So the corrected claim: on a build predating this PR, draft-simple on a hybrid target reports a fabricated 100% acceptance. The other types are unaffected. The artifact is still worth knowing about, since it silently invalidates any tau or C(n) derived from such a run, but it is narrower than I stated.

Unchanged: the mechanism result

Nothing above touches the first comment's primary finding, which was deterministic and reproduced 3/3 per arm: arm A prints speculative decoding will use checkpoints (context does not support partial sequence removal), arm B does not, and the counters go from 82/82/100.000% to 141/58/41.135% on the same prompt and seed. That is this PR working on a second backend.

New, and a separate claim: speculation does not pay in this configuration

My earlier run could not answer this, since patched-vs-unpatched is not spec-vs-autoregressive. Measured now, bracketed, on the same box:

AR_pre   26.50 ± 0.31 t/s     (llama-bench tg128, -p 0 -n 128 -ngl 99 -r 3)
SPEC     17.734 t/s           (k=2, draft-simple, checkpoint line absent)
AR_post  26.71 ± 0.10 t/s
drift    +0.8% between brackets
ratio    0.667x

The spec leg sits between two AR measurements that agree within 0.8%, so it is not thermally suppressed relative to its own baseline. Both legs are decode-only, so the comparison is like-for-like.

0.667x. Speculation loses to plain decode by a third here, with the fix applied. That is not a defect in this PR: the drafter is a 4B against a 27B target, roughly 15% of the target's parameters, evaluated twice per round. The fixed per-round cost dominates before acceptance gets a chance to matter.

Caveats, all of which cut against reading this as a general result:

  • The 4B was never trained as a drafter for this 27B, so 41.135% is not this family's acceptance ceiling.
  • One prompt. Acceptance is known to swing about 1.5x across prompt classes on this target while C(n) stays within 1.8%, so a single prompt characterises C but not tau.
  • Absolute throughput was measured at load 3.6-4.4, not on a quiet machine, so both numbers are depressed by roughly 5% against a quiet-box figure. The ratio is the durable part, not the absolutes.

For contrast, a properly matched drafter on this same target and hardware reaches 1.109 / 1.164 / 1.050 / 0.879 at k = 1/2/3/4 on a code prompt. The lever is drafter size relative to target, not this PR.

@bri-prism

Copy link
Copy Markdown
Author

Addendum: a second, independent cause of the same symptom, which this PR does not fix.

A colleague isolated it more cleanly than my A/B did. Same binary, same target, only the effective spec-type changed — no rebuild, so the binary is held fixed, which my patched-vs-unpatched comparison cannot claim:

llama_model_load: error loading model: unknown model architecture: 'dspark'
speculative decoding will use checkpoints (context does not support partial sequence removal)

The drafter failed to load, and the run did not stop. It silently downgraded into the checkpoint path and went on to emit counters that look like a normal measurement.

So the fabricated-100%-acceptance symptom has at least two distinct causes:

  1. DRAFT_SIMPLE missing from need_n_rs_seq() — what this PR fixes.
  2. A drafter that silently failed to load — not fixed by this PR, and not fixable there. The only tell is the unknown model architecture line sitting immediately above the checkpoint line, which is easy to miss because the interesting output is below it.

Worth distinguishing from the failure mode I hit while setting up my own runs, which is the benign one: omitting --spec-type entirely fails loudly with no implementations specified for speculative decoding and exit 1. That one cannot silently contaminate anything.

Practical guidance for anyone debugging a suspicious acceptance number on a hybrid target: check for a model-load error above the checkpoint line before concluding anything about which defect you are looking at.

@bri-prism

Copy link
Copy Markdown
Author

Additional failure mode: a failed drafter load silently downgrades into the checkpoint path

Reported by a session running this on Metal, and worth recording here because it makes the defect reachable without anyone choosing draft-simple.

When the draft model fails to load, common_speculative_init_from_params() logs and returns at common/speculative.cpp:2435 without establishing a type. Auto-detection (common_speculative_types_from_gguf) had nothing to read, so the type falls back to the default, need_n_rs_seq() returns 0, and the target context is built with n_rs_seq = 0 — which is exactly the misclassification this PR fixes. Observed, consecutive lines:

llama_model_load: error loading model: unknown model architecture: 'dspark'
speculative decoding will use checkpoints (context does not support partial sequence removal)

Same binary, same target, and the checkpoint line appears the moment the type stops being a block-drafter type. That isolates spec-type as the sole variable with the binary held fixed — a cleaner demonstration of the scoping than the measurements in the PR description.

Two consequences for reviewers:

  1. The scoping is by spec-type, not by binary. Pre-fix, --spec-type draft-simple on a rollback-capable arch takes the checkpoint path everywhere — llama-cli, the server, and examples/speculative-simple. The four block-drafter types (DRAFT_MTP, DRAFT_EAGLE3, DRAFT_DFLASH, DRAFT_DSPARK) were already in need_n_rs_seq() and never took it. "Tool X is safe" is not a correct summary; "type Y is safe" is.

  2. A drafter that fails to load produces contaminated numbers plus a wrong spec-type, and the only tell is a log line above the one being read. Anyone benchmarking with a drafter their build cannot load will silently measure the checkpoint path and attribute the result to speculation.

This PR fixes the misclassification, so the downgrade path is no longer expensive. It does not make a failed drafter load noisy — arguably it should error rather than continue, but that is a separate change and not proposed here.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Small micro-batches can now abort, and mixed speculative strategies can exceed the configured rollback bound.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Enables bounded recurrent-state rollback for plain drafter/target speculative decoding on hybrid models.

Changes:

  • Enables rollback snapshots for DRAFT_SIMPLE.
  • Configures the draft context with matching rollback capacity.
File summaries
File Description
common/common.h Enables rollback sizing for simple drafting.
common/speculative.cpp Applies rollback sizing to the draft context.
Review details
  • Files reviewed: 2/2 changed files
  • Comments generated: 3
  • Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.


💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread common/common.h
uint32_t need_n_rs_seq() const {
bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
return t == COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE || t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This holds, thanks. The ubatch assertion was already reachable for MTP, EAGLE3 and DFlash since the target context took n_rs_seq from need_n_rs_seq(); this PR made it reachable from the plain drafter path too. cc02c36 checks the fit once in common_context_params_to_llama and keeps the checkpoint path with a warning when n_ubatch <= n_rs_seq + 1, so a small -ub configuration keeps working the way it did before. The draft context inherits that value instead of overriding it.

Comment thread common/common.h
uint32_t need_n_rs_seq() const {
bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
return t == COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE || t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. speculative-simple now keeps the seq_rm type and decides per round: a checkpoint is taken when the context is FULL, or when it is RS and the draft is longer than llama_n_rs_seq(), the same rule the server applies. The draft context only gets the extra checkpoint in the RS overflow case since a FULL draft context is already checkpointed before drafting, and the restore path guards load_dft the same way.

Comment thread common/speculative.cpp Outdated
…ot fit the ubatch

Recurrent and hybrid memory keep the last n_rs_seq + 1 tokens of a sequence inside one ubatch and assert on it. Decide the fit once in common_context_params_to_llama, warn, and fall back to n_rs_seq = 0 so small -ub configurations keep working. The draft context inherits that value instead of overriding it.

speculative-simple now checkpoints per round when a draft is longer than llama_n_rs_seq(), as the server does, so an n-gram speculator ahead of draft-simple no longer relies on a seq_rm that cannot succeed.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Stale checkpoint reuse, incomplete effective micro-batch validation, and inaccurate memory fitting can cause incorrect state or runtime failures.

Get a fresh assessment by requesting another Copilot review.

Review details
  • Files reviewed: 4/4 changed files
  • Comments generated: 3
  • Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

Comment thread common/common.cpp Outdated
Comment on lines +1725 to +1728
if (cparams.n_rs_seq > 0 && (uint32_t) params.n_ubatch <= cparams.n_rs_seq + 1) {
COM_WRN("%s: speculative rollback window (%u + 1) does not fit ubatch size %d, using KV checkpoints instead (set -ub to at least %u to enable rollback)\n",
__func__, cparams.n_rs_seq, params.n_ubatch, cparams.n_rs_seq + 2);
cparams.n_rs_seq = 0;

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, the guard read the requested size. a16ed76 computes the effective micro-batch the way llama_context does, min(n_batch, n_ubatch) with n_batch clamped to n_ctx when it is set, and checks that. --batch-size 8 --ubatch-size 512 now falls back to checkpoints instead of reaching the assertion.

Comment thread common/speculative.cpp
// note: for small models maybe we can set this to the maximum possible draft from all speculative types
// the extra memory for small models is likely negligible?
cparams.n_rs_seq = 0;
// n_rs_seq stays as common_context_params_to_llama set it: the draft context needs the same rollback window as the target, with n_rs_seq == 0 its seq_rm fails silently on partial acceptance and keeps stale positions

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, the fitter and the real draft context had drifted apart. Removed the n_rs_seq override in the fit path so both contexts are measured with the same recurrent-state shape.

Comment on lines +204 to +206
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, that load was reading a flag from the previous round. The pre-draft reset now loads only for a FULL context, which is the one checkpointed a few lines above; the RS overflow case is still handled after the draft is sized, and the partial-acceptance restore keeps using the current round's value.

llama_context clamps the micro-batch to min(n_batch, n_ubatch) and n_batch to n_ctx, so checking the requested -ub let --batch-size 8 --ubatch-size 512 past the guard and into the same assertion. Check the effective size instead.

The memory fitter no longer forces the draft context's n_rs_seq to zero: recurrent state allocates mem_size * (1 + n_rs_seq) rows, so a fit that measured a different value underestimated the draft context.

In speculative-simple the pre-draft reset now loads a checkpoint only for a FULL context. use_ckpt_dft carries the previous round's RS-overflow decision, and no checkpoint is taken for an RS context at that point, so the load could roll the draft context back to a stale position.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants