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;