From 64a601b582d9d27fe2128a3c8c90d4fb428fcb64 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:46:12 -0700 Subject: [PATCH] speculative: give the draft context the target's recurrent-state slots Two related defects made every speculative round on a hybrid (recurrent + attention) target fall back to a whole-KV checkpoint save and restore. need_n_rs_seq() tested MTP, EAGLE3, DFLASH and DSPARK but not DRAFT_SIMPLE, so a plain drafter plus target got n_rs_seq = 0 and common_context_can_seq_rm() classified the context as SEQ_RM_TYPE_FULL. The arch was never the blocker: LLM_ARCH_QWEN35 is already in llm_arch_supports_rs_rollback(). The draft context then hardcoded cparams.n_rs_seq = 0 regardless of type, so its llama_memory_seq_rm() silently failed and it retained stale positions. That half affects every drafter type, not only DRAFT_SIMPLE. Nothing is needed in speculative-simple.cpp: the partial-acceptance path there already drops the rejected draft tokens unconditionally once the checkpoint branch is skipped, via llama_memory_seq_rm(ctx, seq_id, n_past, -1) on both contexts. --- common/common.h | 2 +- common/speculative.cpp | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/common/common.h b/common/common.h index 16a94408cc8c..eb50a6cbc520 100644 --- a/common/common.h +++ b/common/common.h @@ -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; }); return needs_rs_seq ? draft.n_max : 0u; diff --git a/common/speculative.cpp b/common/speculative.cpp index 1ff0ddeb7e0c..df5c965406fd 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2416,9 +2416,10 @@ 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; + // the draft context must reserve the same per-sequence recurrent-state slots as the target. + // leaving this at 0 makes common_context_can_seq_rm() classify the context as FULL, so + // llama_memory_seq_rm() silently fails on it and stale positions are retained. + cparams.n_rs_seq = params.speculative.need_n_rs_seq(); cparams.ctx_other = ctx_tgt; std::string model_path;