diff --git a/common/common.cpp b/common/common.cpp index 03b4ecb3a46f..4045903ae1a3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -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, @@ -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; 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..05d9ff8a8037 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -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 cparams.ctx_other = ctx_tgt; std::string model_path; diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 487ae03abfa7..d492c1773d24 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -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"); @@ -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); } @@ -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 @@ -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); }