Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bri-prism merged 3 commits into from
Sep 16, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b9b74a3
speculative: stop plain drafter+target falling back to KV checkpoints
bri-prism Sep 14, 2026
cc02c36
speculative: keep the checkpoint path when the rollback window does n…
bri-prism Sep 15, 2026
a16ed76
speculative: base the rollback fit on the effective micro-batch
bri-prism Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1308,8 +1308,6 @@ common_init_result::common_init_result(common_params & params, bool model_only)
if (spec_mtp) {
cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
}
cparams_dft.n_rs_seq = 0;

const common_fit_extra_model extra = {
/*.path_model =*/ params_dft.model.path.c_str(),
/*.mparams =*/ &mparams_dft,
Expand Down Expand Up @@ -1721,6 +1719,18 @@ struct llama_context_params common_context_params_to_llama(const common_params &
cparams.n_ctx = params.n_ctx;
cparams.n_seq_max = params.n_parallel;
cparams.n_rs_seq = params.speculative.need_n_rs_seq();
// recurrent/hybrid memory keeps the last n_rs_seq + 1 tokens of a sequence inside one micro-batch, so the window has to fit or we keep the checkpoint path
// llama_context clamps the micro-batch to min(n_batch, n_ubatch), and n_batch to n_ctx, so the check uses the effective size and not the requested one
{
const uint32_t n_batch_eff = (uint32_t) (params.n_ctx > 0 ? std::min(params.n_batch, params.n_ctx) : params.n_batch);
const uint32_t n_ubatch_eff = params.n_ubatch == 0 ? n_batch_eff : std::min(n_batch_eff, (uint32_t) params.n_ubatch);

if (cparams.n_rs_seq > 0 && n_ubatch_eff <= cparams.n_rs_seq + 1) {
COM_WRN("%s: speculative rollback window (%u + 1) does not fit micro-batch size %u, using KV checkpoints instead (raise -ub and -b to at least %u to enable rollback)\n",
__func__, cparams.n_rs_seq, n_ubatch_eff, cparams.n_rs_seq + 2);
cparams.n_rs_seq = 0;
}
}
cparams.n_outputs_max = std::max(params.n_outputs_max, 0);
cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0);
cparams.n_batch = params.n_batch;
Expand Down
2 changes: 1 addition & 1 deletion common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,7 +386,7 @@ struct common_params_speculative {

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;

@bri-prism bri-prism Sep 15, 2026

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.

@bri-prism bri-prism Sep 15, 2026

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.

});

return needs_rs_seq ? draft.n_max : 0u;
Expand Down
4 changes: 1 addition & 3 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2416,9 +2416,7 @@ common_speculative_init_result::common_speculative_init_result(
// the draft context holds as many tokens per sequence as the target context
cparams.n_ctx = llama_n_ctx(ctx_tgt);

// 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

@bri-prism bri-prism Sep 15, 2026

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.

cparams.ctx_other = ctx_tgt;

std::string model_path;
Expand Down
39 changes: 28 additions & 11 deletions examples/speculative-simple/speculative-simple.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,12 @@ int main(int argc, char ** argv) {
llama_context * ctx_dft = params.speculative.draft.ctx_dft;

// check if the context supports partial sequence removal
const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
const common_context_seq_rm_type seq_rm_tgt = common_context_can_seq_rm(ctx_tgt);
const common_context_seq_rm_type seq_rm_dft = ctx_dft ? common_context_can_seq_rm(ctx_dft) : COMMON_CONTEXT_SEQ_RM_TYPE_NO;

// a bounded rollback window (RS) can only undo up to llama_n_rs_seq() tokens, so longer drafts still take a checkpoint (same rule as the server)
bool use_ckpt_tgt = seq_rm_tgt == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
bool use_ckpt_dft = seq_rm_dft == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;

if (use_ckpt_tgt) {
LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n");
Expand Down Expand Up @@ -173,7 +177,7 @@ int main(int argc, char ** argv) {
llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), seq_id),
llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id));

if (use_ckpt_dft) {
if (seq_rm_dft == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) {
ckpt.update_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}

Expand All @@ -195,21 +199,32 @@ int main(int argc, char ** argv) {
};
common_speculative_draft(spec);

// reset the draft context to the checkpoint before verification
if (ctx_dft) {
// only a FULL context was checkpointed above; an RS context is handled after the draft is sized
if (seq_rm_dft == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) {
ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}

llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);
}

// save a checkpoint of the target context before evaluating the draft
// this allows us to restore the state if partial draft acceptance occurs
if (!draft.empty()) {
use_ckpt_tgt = seq_rm_tgt == COMMON_CONTEXT_SEQ_RM_TYPE_FULL ||
(seq_rm_tgt == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt));
const bool ckpt_dft_rs = seq_rm_dft == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft);
use_ckpt_dft = seq_rm_dft == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ckpt_dft_rs;

if (use_ckpt_tgt) {
ckpt.update_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
}

// reset the draft context to the checkpoint before verification
if (ctx_dft) {
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// a FULL draft context was already checkpointed before drafting; only the RS overflow case needs one here
if (ckpt_dft_rs) {
ckpt.update_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}

llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);
}
} else {
// we have a previous (partial) draft to reuse from checkpoint restoration
Expand Down Expand Up @@ -276,7 +291,9 @@ int main(int argc, char ** argv) {
}

if (ctx_dft) {
ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}

llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);
}
Expand Down