diff --git a/common/arg.cpp b/common/arg.cpp index b62d0e4b4d25..eac878960132 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3581,8 +3581,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("comma-separated list of types of speculative decoding to use (default: %s)\n", common_speculative_type_name_str(params.speculative.types).c_str()), [](common_params & params, const std::string & value) { - const auto enabled_types = string_split(value, ','); - params.speculative.types = common_speculative_types_from_names(enabled_types); + const auto types_str = string_split(value, ','); + auto types = common_speculative_types_from_names(types_str); + params.speculative.types.insert(params.speculative.types.end(), types.begin(), types.end()); } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_TYPE")); add_opt(common_arg( @@ -4071,7 +4072,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--spec-default"}, string_format("enable default speculative decoding config"), [](common_params & params) { - params.speculative.types = { COMMON_SPECULATIVE_TYPE_NGRAM_MOD }; + params.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_NGRAM_MOD); params.speculative.ngram_mod.n_match = 24; params.speculative.ngram_mod.n_min = 48; params.speculative.ngram_mod.n_max = 64; diff --git a/common/common.cpp b/common/common.cpp index ee94274ca62e..9f597b4feb80 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1455,6 +1455,23 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { return res; } +void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + auto * mem = llama_get_memory(ctx); + if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) { + GGML_ABORT("%s", string_format("failed to remove sequence %d with p0=%d, p1=%d\n", seq_id, p0, p1).c_str()); + } +} + +void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + auto * mem = llama_get_memory(ctx); + llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1); +} + +void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { + auto * mem = llama_get_memory(ctx); + llama_memory_seq_add(mem, seq_id, p0, p1, delta); +} + void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora) { std::vector loras; std::vector scales; @@ -2081,3 +2098,11 @@ void common_prompt_checkpoint::load_dft( GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_dft.size(), n); } } + +void common_prompt_checkpoint::clear_tgt() { + data_tgt.clear(); +} + +void common_prompt_checkpoint::clear_dft() { + data_dft.clear(); +} diff --git a/common/common.h b/common/common.h index 1dd8de472432..222f3177b243 100644 --- a/common/common.h +++ b/common/common.h @@ -896,6 +896,10 @@ enum common_context_seq_rm_type { // note: clears the memory of the context common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx); +// aborts execution on failure +void common_context_seq_rm (llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1); +void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); +void common_context_seq_cp (llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1); // // Batch utils @@ -1077,4 +1081,7 @@ struct common_prompt_checkpoint { llama_context * ctx, llama_seq_id seq_id, llama_state_seq_flags flags) const; + + void clear_tgt(); + void clear_dft(); }; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 109a77be4042..084c5d9ea4f3 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -170,10 +170,12 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // partial rollback via per-token snapshot index (bounded by n_rs_seq) if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) { const llama_pos rollback = cell.pos - (p0 - 1); - GGML_ASSERT(rollback >= 1 && rollback <= (llama_pos) n_rs_seq); - set_rs_idx(seq_id, (uint32_t) rollback); - cell.pos = p0 - 1; - return true; + if (rollback >= 1 && rollback <= (llama_pos) n_rs_seq) { + set_rs_idx(seq_id, (uint32_t) rollback); + cell.pos = p0 - 1; + return true; + } + return false; } // invalidate tails which will be cleared if (p0 <= cell.pos && cell.pos < p1) { diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 530ac383e85b..bc8db4a306ff 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -145,9 +145,9 @@ struct server_slot { SLT_INF(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); - llama_memory_seq_rm(llama_get_memory(ctx_tgt), id, -1, -1); + common_context_seq_rm(ctx_tgt, id, -1, -1); if (ctx_dft) { - llama_memory_seq_rm(llama_get_memory(ctx_dft), id, -1, -1); + common_context_seq_rm(ctx_dft, id, -1, -1); } prompt.tokens.clear(); @@ -517,12 +517,12 @@ struct server_slot { void copy_state_to(server_slot & other) const { GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT); - llama_memory_seq_rm(llama_get_memory(ctx_tgt), other.id, -1, -1); - llama_memory_seq_cp(llama_get_memory(ctx_tgt), id, other.id, -1, -1); + common_context_seq_rm(ctx_tgt, other.id, -1, -1); + common_context_seq_cp(ctx_tgt, id, other.id, -1, -1); if (ctx_dft) { - llama_memory_seq_rm(llama_get_memory(ctx_dft), other.id, -1, -1); - llama_memory_seq_cp(llama_get_memory(ctx_dft), id, other.id, -1, -1); + common_context_seq_rm(ctx_dft, other.id, -1, -1); + common_context_seq_cp(ctx_dft, id, other.id, -1, -1); } other.n_decoded = n_decoded; @@ -788,6 +788,8 @@ struct server_context_impl { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; } + // 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; ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); @@ -802,6 +804,7 @@ struct server_context_impl { auto cparams_mtp = common_context_params_to_llama(params_base); cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_mtp.n_rs_seq = 0; ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp)); if (ctx_dft == nullptr) { @@ -2221,12 +2224,12 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - llama_memory_seq_rm (llama_get_memory(ctx_tgt), slot.id, n_keep , n_keep + n_discard); - llama_memory_seq_add(llama_get_memory(ctx_tgt), slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); + common_context_seq_rm (ctx_tgt, slot.id, n_keep , n_keep + n_discard); + common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); if (ctx_dft) { - llama_memory_seq_rm (llama_get_memory(ctx_dft.get()), slot.id, n_keep , n_keep + n_discard); - llama_memory_seq_add(llama_get_memory(ctx_dft.get()), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard); + common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); } // add generated tokens to cache @@ -2333,14 +2336,23 @@ struct server_context_impl { slot.n_draft_total += draft.size(); // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] + const bool use_ckpt_dft = ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + if (ctx_dft) { - ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + if (use_ckpt_dft) { + ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + } - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), slot.id, ckpt.pos_max + 1, -1); + common_context_seq_rm(ctx_dft.get(), slot.id, ckpt.pos_max + 1, -1); } if (!draft.empty()) { - const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + const bool use_ckpt_tgt = + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt)); + + const bool use_ckpt_dft = + (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get())); if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); @@ -2355,6 +2367,10 @@ struct server_context_impl { (float) ckpt.size() / 1024 / 1024, (float) ckpt.data_dft.size() / 1024 / 1024); } + + if (use_ckpt_dft) { + ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + } } } @@ -2526,12 +2542,12 @@ struct server_context_impl { const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c; - llama_memory_seq_rm (llama_get_memory(ctx_tgt), slot.id, head_p, head_c); - llama_memory_seq_add(llama_get_memory(ctx_tgt), slot.id, head_c, head_c + n_match, kv_shift); + common_context_seq_rm (ctx_tgt, slot.id, head_p, head_c); + common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift); if (ctx_dft) { - llama_memory_seq_rm (llama_get_memory(ctx_dft.get()), slot.id, head_p, head_c); - llama_memory_seq_add(llama_get_memory(ctx_dft.get()), slot.id, head_c, head_c + n_match, kv_shift); + common_context_seq_rm (ctx_dft.get(), slot.id, head_p, head_c); + common_context_seq_add(ctx_dft.get(), slot.id, head_c, head_c + n_match, kv_shift); } for (size_t i = 0; i < n_match; i++) { @@ -2694,18 +2710,10 @@ struct server_context_impl { SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - if (!llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, p0, -1)) { - SLT_WRN(slot, "failed to truncate tokens with position >= %d - clearing the memory\n", p0); - - slot.prompt_clear(true); - - // there is no common part left - slot.n_prompt_tokens_cache = 0; - } else { - if (ctx_dft && !llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), slot.id, p0, -1)) { - GGML_ABORT("failed to truncate draft context\n"); - } - } + common_context_seq_rm(ctx_tgt, slot.id, p0, -1); + if (ctx_dft) { + common_context_seq_rm(ctx_dft.get(), slot.id, p0, -1); + } // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3171,13 +3179,8 @@ struct server_context_impl { // verify and try to accept the draft { - const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; - - // only save the sampler sampler state if we use checkpoints - common_sampler_ptr smpl_save; - if (use_ckpt_tgt) { - smpl_save.reset(common_sampler_clone(slot.smpl.get())); - } + // save the sampler sampler state in case we need to restore it + common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); @@ -3185,8 +3188,14 @@ struct server_context_impl { GGML_ASSERT(accepted.size() >= 1); + const uint32_t n_rollback = slot.spec_draft.size() + 1 - accepted.size(); + + const bool use_ckpt_tgt = + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_rollback > llama_n_rs_seq(ctx_tgt)); + // check for partial draft acceptance - if (accepted.size() < slot.spec_draft.size() + 1) { + if (n_rollback > 0) { if (use_ckpt_tgt) { if (trace > 0) { SLT_INF(slot, "accepted %2zu/%2zu draft tokens (restore checkpoint)\n", accepted.size() - 1, slot.spec_draft.size()); @@ -3202,13 +3211,13 @@ struct server_context_impl { { ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); - llama_memory_seq_rm(llama_get_memory(slot.ctx_tgt), slot.id, ckpt.pos_max + 1, -1); + common_context_seq_rm(slot.ctx_tgt, slot.id, ckpt.pos_max + 1, -1); } if (slot.ctx_dft) { ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); - llama_memory_seq_rm(llama_get_memory(slot.ctx_dft), slot.id, ckpt.pos_max + 1, -1); + common_context_seq_rm(slot.ctx_dft, slot.id, ckpt.pos_max + 1, -1); } slot.prompt.tokens.keep_first(ckpt.n_tokens); @@ -3244,9 +3253,9 @@ struct server_context_impl { slot.sampled = ids.back(); // last accepted token SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft); - llama_memory_seq_rm(llama_get_memory(slot.ctx_tgt), slot.id, slot.prompt.tokens.pos_next(), -1); + common_context_seq_rm(slot.ctx_tgt, slot.id, slot.prompt.tokens.pos_next(), -1); if (slot.ctx_dft) { - llama_memory_seq_rm(llama_get_memory(slot.ctx_dft), slot.id, slot.prompt.tokens.pos_next(), -1); + common_context_seq_rm(slot.ctx_dft, slot.id, slot.prompt.tokens.pos_next(), -1); } for (size_t i = 0; i < ids.size(); ++i) {