Skip to content

server : adjust checkpoint logic #11

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
am17an merged 2 commits into from
May 14, 2026
Merged

server : adjust checkpoint logic #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10f5216
server : adjust checkpoint logic
ggerganov May 14, 2026
d7edf57
cont : rm asserts
ggerganov May 14, 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
7 changes: 4 additions & 3 deletions common/arg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string>(value, ',');
params.speculative.types = common_speculative_types_from_names(enabled_types);
const auto types_str = string_split<std::string>(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(
Expand Down Expand Up @@ -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;
Expand Down
25 changes: 25 additions & 0 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<common_adapter_lora_info> & lora) {
std::vector<llama_adapter_lora *> loras;
std::vector<float> scales;
Expand Down Expand Up @@ -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();
}
7 changes: 7 additions & 0 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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();
};
10 changes: 6 additions & 4 deletions src/llama-memory-recurrent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {

@am17an am17an May 14, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added the assert here, it's okay to remove?

@ggerganov ggerganov May 14, 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.

We count on this to return false in the common_context_can_seq_rm() to determine if partial remove is allowed or not. The assert triggers on startup with regular non-MTP recurrent models.

@am17an am17an May 14, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it seems to crash with the assert so I think i'll just merge

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) {
Expand Down
91 changes: 50 additions & 41 deletions tools/server/server-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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));

Expand All @@ -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) {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
}
}

Expand Down Expand Up @@ -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++) {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -3171,22 +3179,23 @@ 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);
slot.spec_i_batch.clear();

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());
Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down
Loading