Repository: ggerganov/llama.cpp
File: tools/server/server.cpp (C++)

Code before:
                if (slot.state == SLOT_STATE_DONE_PROMPT) {
                    if (slot.task_type == SERVER_TASK_TYPE_EMBEDDING) {
                        // prompt evaluated for embedding
                        send_embedding(slot, batch_view);
                        slot.release();
                        slot.i_batch = -1;
                        continue; // continue loop of slots
                    }

                    if (slot.task_type == SERVER_TASK_TYPE_RERANK) {
                        send_rerank(slot, batch_view);
                        slot.release();
                        slot.i_batch = -1;
                        continue; // continue loop of slots
                    }

                    // prompt evaluated for next-token prediction
                    slot.state = SLOT_STATE_GENERATING;

                    // make a checkpoint of the parts of memory that cannot be rolled back.
                    // checkpoints are needed only if:
                    // - the model uses SWA and we are not using `swa_full`
                    // - the model architecture is marked as recurrent or hybrid
                    bool do_checkpoint = (llama_model_is_recurrent(model) || llama_model_is_hybrid(model)) ||
                         (llama_model_n_swa(model) > 0 && !params_base.swa_full);

                    if (do_checkpoint && params_base.n_ctx_checkpoints > 0) {
                        if (slot.ctx_checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
                            // make room for the new checkpoint, if needed
                            const auto & cur = slot.ctx_checkpoints.back();
                            SLT_WRN(slot, "erasing old context checkpoint (pos_min = %d, pos_max = %d, size = %.3f MiB)\n",
                                    cur.pos_min, cur.pos_max, (float) cur.data.size() / 1024 / 1024);

                            slot.ctx_checkpoints.erase(slot.ctx_checkpoints.begin());
                        }

                        const size_t checkpoint_size = llama_state_seq_get_size_ext(ctx, slot.id, LLAMA_STATE_SEQ_FLAGS_CHECKPOINT_ONLY);

                        auto & cur = slot.ctx_checkpoints.emplace_back(ctx_checkpoint{
                            /*.pos_min = */ llama_memory_seq_pos_min(llama_get_memory(ctx), slot.id),
                            /*.pos_max = */ llama_memory_seq_pos_max(llama_get_memory(ctx), slot.id),
                            /*.data    = */ std::vector<uint8_t>(checkpoint_size),
                        });

                        llama_state_seq_get_data_ext(ctx, cur.data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_CHECKPOINT_ONLY);

                        SLT_WRN(slot, "saved context checkpoint %d of %d (pos_min = %d, pos_max = %d, size = %.3f MiB)\n",
                                (int) slot.ctx_checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, cur.pos_max, (float) cur.data.size() / 1024 / 1024);
                    }
                } else if (slot.state != SLOT_STATE_GENERATING) {
                    continue; // continue loop of slots

Review comment (suggestion):
I'm a bit torn on this logic for determining when to do checkpoints. It should be centred around the memory module or the context, rather than the model.

Just making a note for the future - no need to change anything in this PR.

Code after:

                if (slot.state == SLOT_STATE_DONE_PROMPT) {
                    if (slot.task_type == SERVER_TASK_TYPE_EMBEDDING) {
                        // prompt evaluated for embedding
                        send_embedding(slot, batch_view);
                        slot.release();
                        slot.i_batch = -1;
                        continue; // continue loop of slots
                    }

                    if (slot.task_type == SERVER_TASK_TYPE_RERANK) {
                        send_rerank(slot, batch_view);
                        slot.release();
                        slot.i_batch = -1;
                        continue; // continue loop of slots
                    }

                    // prompt evaluated for next-token prediction
                    slot.state = SLOT_STATE_GENERATING;

                    // make a checkpoint of the parts of the memory that cannot be rolled back.
                    // checkpoints are created only if:
                    // - the model uses SWA and we are not using `swa_full`
                    // - the model architecture is marked as recurrent or hybrid
                    //
                    // TODO: try to make this conditional on the context or the memory module, instead of the model type
                    const bool do_checkpoint =
                        (llama_model_is_recurrent(model) || llama_model_is_hybrid(model)) ||
                        (llama_model_n_swa(model) > 0 && !params_base.swa_full);

                    if (do_checkpoint && params_base.n_ctx_checkpoints > 0) {
                        while (slot.ctx_checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
                            // make room for the new checkpoint, if needed
                            const auto & cur = slot.ctx_checkpoints.front();
                            SLT_WRN(slot, "erasing old context checkpoint (pos_min = %d, pos_max = %d, size = %.3f MiB)\n",
                                    cur.pos_min, cur.pos_max, (float) cur.data.size() / 1024 / 1024);

                            slot.ctx_checkpoints.erase(slot.ctx_checkpoints.begin());
                        }

                        const size_t checkpoint_size = llama_state_seq_get_size_ext(ctx, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);

                        auto & cur = slot.ctx_checkpoints.emplace_back(ctx_checkpoint{
                            /*.pos_min = */ llama_memory_seq_pos_min(llama_get_memory(ctx), slot.id),
                            /*.pos_max = */ llama_memory_seq_pos_max(llama_get_memory(ctx), slot.id),
                            /*.data    = */ std::vector<uint8_t>(checkpoint_size),
                        });

                        llama_state_seq_get_data_ext(ctx, cur.data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);

                        SLT_WRN(slot, "saved context checkpoint %d of %d (pos_min = %d, pos_max = %d, size = %.3f MiB)\n",