#include "server-context.h"

#include "server-chat.h"

#include "server-common.h"

#include "server-http.h"

#include "server-task.h"

#include "server-queue.h"

#include "server-schema.h"

#include "build-info.h"

#include "common.h"

#include "fit.h"

#include "llama.h"

#include "log.h"

#include "sampling.h"

#include "speculative.h"

#include "mtmd.h"

#include "mtmd-helper.h"

#include <algorithm>

#include <cstddef>

#include <cinttypes>

#include <exception>

#include <memory>

#include <filesystem>

#include <utility>

#include <fstream>

// fix problem with std::min and std::max

#if defined(_WIN32)

#define WIN32_LEAN_AND_MEAN

#ifndef NOMINMAX

# define NOMINMAX

#endif

#include <windows.h>

#endif

using json = nlohmann::ordered_json;

constexpr int HTTP_POLLING_SECONDS = 1;

static uint32_t server_n_outputs_max(const common_params & params) {

const uint32_t n_batch = params.n_batch;

if (params.embedding ||

(params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {

return n_batch;

}

const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);

const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;

return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));

}

// state diagram: https://github.com/ggml-org/llama.cpp/pull/9283

enum slot_state {

SLOT_STATE_IDLE,

SLOT_STATE_WAIT_OTHER, // after assigning a task, but waiting for parent slot to process prompt

SLOT_STATE_STARTED, // after assigning a task and about to process prompt

SLOT_STATE_PROCESSING_PROMPT,

SLOT_STATE_DONE_PROMPT,

SLOT_STATE_GENERATING,

};

struct server_slot {

int id;

llama_context * ctx_tgt = nullptr;

llama_context * ctx_dft = nullptr;

// multimodal

mtmd_context * mctx = nullptr;

mtmd::batch_ptr mbatch = nullptr;

// speculative decoding

common_speculative * spec;

llama_tokens spec_draft;

llama_tokens spec_prompt;

std::vector<int32_t> spec_i_batch;

common_prompt_checkpoint spec_ckpt;

// TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state

// see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837

std::unique_ptr<const server_task> task;

std::unique_ptr<const server_task> task_prev; // used for debugging

// used to determine the slot that has been used the longest

int64_t t_last_used = -1;

// generation props

int32_t n_ctx = 0; // context size per slot

int32_t n_keep = 0;

int32_t n_decoded = 0;

int32_t n_remaining = -1;

int32_t i_batch = -1;

int32_t n_prompt_tokens_cache = 0;

int32_t n_prompt_tokens_processed = 0;

size_t last_nl_pos = 0;

std::string generated_text;

std::string debug_generated_text;

llama_tokens generated_tokens;

std::vector<completion_token_output> generated_token_probs;

bool has_next_token = true;

bool has_new_line = false;

bool truncated = false;

stop_type stop;

std::string stopping_word;

// state

slot_state state = SLOT_STATE_IDLE;

server_prompt prompt;

bool prompt_save(server_prompt_cache & prompt_cache) const {

if (prompt.tokens.size() == 0) {

return false;

}

GGML_ASSERT(prompt.data.size() == 0);

const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);

const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0;

const size_t cur_size = cur_size_tgt + cur_size_dft;

SRV_WRN(" - saving prompt with length %d, total state size = %.3f MiB (draft: %.3f MiB)\n",

(int) prompt.tokens.size(), cur_size / (1024.0 * 1024.0), cur_size_dft / (1024.0 * 1024.0));

auto * cur = prompt_cache.alloc(prompt, cur_size_tgt, cur_size_dft);

if (cur == nullptr) {

return false;

}

llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);

if (ctx_dft) {

llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);

}

return true;

}

bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) {

bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id);

if (!res) {

SLT_WRN(*this, "%s", "failed to load prompt from cache\n");

}

return res;

}

void prompt_clear(bool allow_processing) {

if (!allow_processing) {

GGML_ASSERT(!is_processing());

}

SLT_INF(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size());

common_context_seq_rm(ctx_tgt, id, -1, -1);

if (ctx_dft) {

common_context_seq_rm(ctx_dft, id, -1, -1);

}

prompt.tokens.clear();

}

std::vector<common_adapter_lora_info> lora;

int32_t alora_invocation_start = -1;

// sampling

json json_schema;

common_sampler_ptr smpl;

llama_token sampled; // in speculative mode, this is the last accepted token

// stats

size_t n_sent_text = 0; // number of sent text character

int64_t t_start_process_prompt;

int64_t t_start_generation;

int64_t t_print_last = 0;

int32_t n_decoded_last = 0;

double t_prompt_processing = 0.0; // ms

double t_token_generation = 0.0; // ms

std::function<void(int /* id_slot */)> callback_on_release;

// Speculative decoding stats

int32_t n_draft_total = 0; // Total draft tokens generated

int32_t n_draft_accepted = 0; // Draft tokens actually accepted

int32_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model

std::vector<int32_t> n_accepted_per_pos; // Accepted tokens per draft position

void reset() {

SLT_DBG(*this, "%s", "\n");

n_prompt_tokens_cache = 0;

last_nl_pos = 0;

generated_text = "";

has_new_line = false;

truncated = false;

stop = STOP_TYPE_NONE;

stopping_word = "";

n_sent_text = 0;

if (can_speculate()) {

spec_draft.clear();

spec_i_batch.clear();

spec_ckpt.clear();

}

generated_tokens.clear();

generated_token_probs.clear();

json_schema = json();

// clear speculative decoding stats

n_draft_total = 0;

n_draft_accepted = 0;

n_draft_verif_steps = 0;

n_accepted_per_pos.clear();

task_prev = std::move(task);

task.reset();

llama_set_sampler(ctx_tgt, id, nullptr);

// clear alora start

alora_invocation_start = -1;

// clear multimodal state

mbatch.reset();

}

void init_sampler() const {

common_sampler_reset(smpl.get());

if (!task->need_sampling()) {

return;

}

const int64_t t_start = ggml_time_us();

int n_text = 0;

for (int i = 0; i < (int) prompt.tokens.size(); i++) {

const llama_token id = prompt.tokens[i];

if (id != LLAMA_TOKEN_NULL) {

common_sampler_accept(smpl.get(), id, false);

n_text++;

}

}

SLT_TRC(*this, "init sampler, took %0.2f ms, tokens: text = %d, total = %d\n",

(ggml_time_us() - t_start) / 1000.0, n_text, (int) prompt.tokens.size());

}

bool need_embd() const {

GGML_ASSERT(task);

return task->need_embd() || (spec && common_speculative_need_embd(spec));

}

bool need_embd_nextn() const {

GGML_ASSERT(task);

return spec && common_speculative_need_embd_nextn(spec);

}

// if the context does not have a memory module then all embeddings have to be computed within a single ubatch

// also we cannot split if the pooling would require any past tokens

// (MTP supports splitting — uses task->need_embd() not need_embd())

bool can_split() const {

GGML_ASSERT(task);

return

!task->need_embd() ||

(llama_get_memory(ctx_tgt) && llama_pooling_type(ctx_tgt) == LLAMA_POOLING_TYPE_LAST);

}

bool can_batch_with(server_slot & other_slot) const {

GGML_ASSERT(task);

return task->type == other_slot.task->type && are_lora_equal(lora, other_slot.lora);

}

bool has_budget(const common_params & global_params) {

GGML_ASSERT(task);

if (task->params.n_predict == -1 && global_params.n_predict == -1) {

return true; // limitless

}

n_remaining = -1;

if (task->params.n_predict != -1) {

n_remaining = task->params.n_predict - n_decoded;

} else if (global_params.n_predict != -1) {

n_remaining = global_params.n_predict - n_decoded;

}

return n_remaining > 0; // no budget

}

bool is_processing() const {

return state != SLOT_STATE_IDLE;

}

bool can_speculate() const {

return !!spec;

}

void add_token(const completion_token_output & token) {

if (!is_processing()) {

SLT_WRN(*this, "%s", "slot is not processing\n");

return;

}

generated_token_probs.push_back(token);

}

int get_n_draft_max() const {

GGML_ASSERT(task);

if (!can_speculate()) {

return 0;

}

// determine the max draft that fits the current slot state

// note: slot.prompt is not yet expanded with the `id` token sampled above

// also, need to leave space for 1 extra token to allow context shifts

int n_draft_max = n_ctx - prompt.n_tokens() - 2;

if (n_remaining > 0) {

n_draft_max = std::min(n_draft_max, n_remaining - 1);

}

SLT_DBG(*this, "max possible draft: %d\n", n_draft_max);

return n_draft_max;

}

void update_batch(llama_batch & batch) {

if (spec_draft.empty()) {

// no speculative decoding

i_batch = batch.n_tokens;

common_batch_add(batch, sampled, prompt.tokens.pos_next(), { this->id }, true);

SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n",

sampled, n_ctx, prompt.n_tokens(), truncated);

} else {

SLT_DBG(*this, "generate_draft: id=%d, #tokens=%zu, #draft=%zu, pos_next=%d\n",

sampled, prompt.tokens.size(), spec_draft.size(), prompt.tokens.pos_next());

GGML_ASSERT(spec_i_batch.empty());

spec_i_batch.push_back(batch.n_tokens);

for (size_t i = 0; i < spec_draft.size(); i++) {

spec_i_batch.push_back(batch.n_tokens + i + 1);

}

auto pos0 = prompt.tokens.pos_next();

common_batch_add(batch, sampled, pos0++, { this->id }, true);

for (auto token : spec_draft) {

common_batch_add(batch, token, pos0++, { this->id }, true);

}

}

prompt.tokens.push_back(sampled);

prompt.tokens.insert(spec_draft);

}

void release() {

if (is_processing()) {

GGML_ASSERT(task);

SLT_INF(*this, "stop processing: n_tokens = %d, truncated = %d\n", prompt.n_tokens(), truncated);

t_last_used = ggml_time_us();

t_token_generation = (ggml_time_us() - t_start_generation) / 1e3;

state = SLOT_STATE_IDLE;

// do not keep context of the child slots - the parent's context is enough

if (task->is_child()) {

prompt_clear(false);

}

reset();

callback_on_release(id);

}

}

result_timings get_timings() const {

result_timings timings;

timings.cache_n = n_prompt_tokens_cache;

timings.prompt_n = n_prompt_tokens_processed;

timings.prompt_ms = t_prompt_processing;

timings.prompt_per_token_ms = t_prompt_processing / n_prompt_tokens_processed;

timings.prompt_per_second = 1e3 / t_prompt_processing * n_prompt_tokens_processed;

timings.predicted_n = n_decoded;

timings.predicted_ms = t_token_generation;

timings.predicted_per_token_ms = t_token_generation / n_decoded;

timings.predicted_per_second = 1e3 / t_token_generation * n_decoded;

// Add speculative metrics

if (n_draft_total > 0) {

timings.draft_n = n_draft_total;

timings.draft_n_accepted = n_draft_accepted;

}

return timings;

}

size_t find_stopping_strings(const std::string & text, const size_t last_token_size, bool is_full_stop) {

GGML_ASSERT(task);

size_t stop_pos = std::string::npos;

for (const std::string & word : task->params.antiprompt) {

size_t pos;

if (is_full_stop) {

const size_t tmp = word.size() + last_token_size;

const size_t from_pos = text.size() > tmp ? text.size() - tmp : 0;

pos = text.find(word, from_pos);

} else {

// otherwise, partial stop

pos = string_find_partial_stop(text, word);

}

if (pos != std::string::npos && (stop_pos == std::string::npos || pos < stop_pos)) {

if (is_full_stop) {

stop = STOP_TYPE_WORD;

stopping_word = word;

has_next_token = false;

}

stop_pos = pos;

}

}

return stop_pos;

}

void print_timings_tg() {

if (n_decoded < 100) {

return;

}

const int64_t t_now = ggml_time_us();

if (t_now - t_print_last < 3*1000*1000) {

return;

}

const double n_gen_second = 1e3 / (t_token_generation) * (n_decoded);

const double n_gen_second_win = 1e6 / (t_now - t_print_last) * (n_decoded - n_decoded_last);

t_print_last = t_now;

n_decoded_last = n_decoded;

SLT_INF(*this, "n_decoded = %6d, tg = %6.2f t/s, tg_3s = %6.2f t/s\n", n_decoded, n_gen_second, n_gen_second_win);

}

void print_timings_pp() const {

const double n_prompt_second = 1e3 / t_prompt_processing * n_prompt_tokens_processed;

const double f_progress = (float) prompt.n_tokens() / task->n_tokens();

if (t_prompt_processing < 3000.0) {

return;

}

SLT_INF(*this, "prompt processing, n_tokens = %6d, progress = %.2f, t = %6.2f s / %.2f tokens per second\n",

n_prompt_tokens_processed, f_progress, t_prompt_processing / 1e3, n_prompt_second);

}

void print_timings() const {

const double t_prompt = t_prompt_processing / n_prompt_tokens_processed;

const double n_prompt_second = 1e3 / t_prompt_processing * n_prompt_tokens_processed;

const double t_gen = t_token_generation / n_decoded;

const double n_gen_second = 1e3 / t_token_generation * n_decoded;

SLT_INF(*this,

"prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n",

t_prompt_processing, n_prompt_tokens_processed, t_prompt, n_prompt_second);

SLT_INF(*this,

" eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n",

t_token_generation, n_decoded, t_gen, n_gen_second);

SLT_INF(*this,

" total time = %10.2f ms / %5d tokens\n",

t_prompt_processing + t_token_generation, n_prompt_tokens_processed + n_decoded);

SLT_INF(*this,

" graphs reused = %10d\n",

llama_perf_context(ctx_tgt).n_reused);

if (n_draft_total > 0) {

const float draft_ratio = (float) n_draft_accepted / n_draft_total;

const double mean_acc_len = n_draft_verif_steps > 0 ? 1.0 + (double) n_draft_accepted / (double) n_draft_verif_steps : 1.0;

std::string acceptance_rates_per_pos;

if (n_draft_verif_steps > 0) {

for (size_t i = 0; i < n_accepted_per_pos.size(); ++i) {

if (i > 0) {

acceptance_rates_per_pos += ", ";

}

acceptance_rates_per_pos += string_format("%.3f", (double) n_accepted_per_pos[i] / (double) n_draft_verif_steps);

}

}

SLT_INF(*this,

"draft acceptance = %0.5f (%5d accepted / %5d generated), mean acceptance length = %5.2f, acceptance rate per position = (%s)\n",

draft_ratio, n_draft_accepted, n_draft_total, mean_acc_len, acceptance_rates_per_pos.c_str());

}

common_speculative_print_stats(spec);

}

json to_json(bool only_metrics = false) const {

json res;

res = {

{"id", id},

{"n_ctx", n_ctx},

{"speculative", can_speculate()},

{"is_processing", is_processing()},

};

const auto & ptask = task ? task : task_prev;

if (ptask) {

res["id_task"] = ptask->id;

res["n_prompt_tokens"] = (int32_t) prompt.tokens.size();

res["n_prompt_tokens_processed"] = n_prompt_tokens_processed;

res["n_prompt_tokens_cache"] = n_prompt_tokens_cache;

res["params"] = ptask->params.to_json(only_metrics);

res["next_token"] = {

{

{"has_next_token", has_next_token},

{"has_new_line", has_new_line},

{"n_remain", n_remaining},

{"n_decoded", n_decoded},

}

};

if (!only_metrics) {

res["prompt"] = ptask->tokens.detokenize(ctx_tgt, true);

res["generated"] = generated_text.empty() ? debug_generated_text : generated_text;

}

}

return res;

}

void copy_state_to(server_slot & other) const {

GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT);

common_context_seq_rm(ctx_tgt, other.id, -1, -1);

common_context_seq_cp(ctx_tgt, id, other.id, -1, -1);

if (ctx_dft) {

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;

other.n_remaining = n_remaining;

other.i_batch = i_batch;

other.t_start_process_prompt = t_start_process_prompt;

other.t_prompt_processing = t_prompt_processing;

other.n_prompt_tokens_cache = n_prompt_tokens_cache;

other.n_prompt_tokens_processed = n_prompt_tokens_processed;

other.prompt = prompt.clone();

other.init_sampler();

}

// returns 0 on success

// caller need to update prompt.tokens after a successful call to keep track of the processing progress

int process_mtmd_chunk(size_t idx, size_t & n_tokens_out) {

GGML_ASSERT(mctx);

const auto & input_tokens = task->tokens;

const auto & chunk = input_tokens.find_chunk(idx);

int32_t res = 0;

auto try_decode = [&]() -> int32_t {

if (mbatch) {

float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get());

if (embd) {

void * cb_data = spec;

static auto cb = [](llama_batch batch, void * user_data) {

common_speculative * spec = static_cast<common_speculative *>(user_data);

if (!common_speculative_process(spec, batch)) {

return 1;

}

return 0;

};

llama_pos new_n_past; // unused for now

res = mtmd_helper_decode_image_chunk(

mctx,

ctx_tgt,

chunk.get(),

embd,

prompt.tokens.pos_next(),

id,

llama_n_batch(ctx_tgt),

&new_n_past,

cb,

cb_data

);

if (res != 0) {

SLT_ERR(*this, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res);

return -1;

}

n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get());

return 0; // success

}

}

return 1; // (non-error) need to create & encode batch

};

// if the batch is already exist, try searching & encode

res = try_decode();

if (res == 0) {

return 0;

}

if (res < 0) {

// fatal error

return res;

}

// otherwise, the batch is either uninitialized or is used up

// we need to create & encode a new batch

mbatch.reset(mtmd_batch_init(mctx));

res = mtmd_batch_add_chunk(mbatch.get(), chunk.get());

GGML_ASSERT(res == 0); // we should never have an empty batch

// try batching as much as possible

int n_added = 1;

size_t idx_cur = idx;

while (res == 0) {

auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur);

if (next_chunk == nullptr) {

break;

}

res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get());

n_added += (res == 0 ? 1 : 0);

idx_cur = next_idx;

SLT_DBG(*this, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res);

// if res != 0, batch is full or chunk is not compatible -> this loop breaks

}

// TODO @ngxson : move this log line to debug when it become more stable

SLT_INF(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);

res = mtmd_batch_encode(mbatch.get());

if (res != 0) {

SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);

return -1;

}

return try_decode();

}

};

//

// server_metrics

//

struct server_metrics {

int64_t t_start = 0;

uint64_t n_prompt_tokens_processed_total = 0;

uint64_t t_prompt_processing_total = 0;

uint64_t n_tokens_predicted_total = 0;

uint64_t t_tokens_generation_total = 0;

uint64_t n_tokens_max = 0;

uint64_t n_prompt_tokens_processed = 0;

uint64_t t_prompt_processing = 0;

uint64_t n_tokens_predicted = 0;

uint64_t t_tokens_generation = 0;

uint64_t n_decode_total = 0;

uint64_t n_busy_slots_total = 0;

void init() {

t_start = ggml_time_us();

}

void on_prompt_eval(const server_slot & slot) {

n_prompt_tokens_processed_total += slot.n_prompt_tokens_processed;

n_prompt_tokens_processed += slot.n_prompt_tokens_processed;

t_prompt_processing += slot.t_prompt_processing;

t_prompt_processing_total += slot.t_prompt_processing;

n_tokens_max = std::max(n_tokens_max, (uint64_t) slot.prompt.n_tokens());

}

void on_prediction(const server_slot & slot) {

n_tokens_predicted_total += slot.n_decoded;

n_tokens_predicted += slot.n_decoded;

t_tokens_generation += slot.t_token_generation;

t_tokens_generation_total += slot.t_token_generation;

}

void on_decoded(const std::vector<server_slot> & slots) {

n_decode_total++;

for (const auto & slot : slots) {

if (slot.is_processing()) {

n_busy_slots_total++;

}

n_tokens_max = std::max(n_tokens_max, (uint64_t) slot.prompt.n_tokens());

}

}

void reset_bucket() {

n_prompt_tokens_processed = 0;

t_prompt_processing = 0;

n_tokens_predicted = 0;

t_tokens_generation = 0;

}

};

//

// server_context_impl (private implementation)

//

struct server_context_impl {

friend struct server_context;

public:

// only use these pointers outside of this class:

// - when not in sleeping state

// - and, with thread-safe APIs (e.g., tokenizer calls)

llama_model * model_tgt = nullptr;

mtmd_context * mctx = nullptr;

const llama_vocab * vocab = nullptr;

server_queue queue_tasks;

server_response queue_results;

// note: chat_params must not be refreshed upon existing sleeping state

server_chat_params chat_params;

server_state_callback_t callback_state = [](server_state, json) -> void {};

server_context_impl() {

mtmd_helper_log_set(common_log_default_callback, nullptr);

}

~server_context_impl() {

if (!sleeping) {

// destroy() is already called when entering sleeping state

// we don't call it again here to avoid double free

destroy();

}

}

private:

// note: accessing these fields outside of this class is not thread-safe

// use server_context methods instead

common_params params_base;

// note: keep these alive - they determine the lifetime of the model, context, etc.

common_init_result_ptr llama_init;

llama_context * ctx_tgt = nullptr;

llama_batch batch {};

llama_model_ptr model_dft;

llama_context_ptr ctx_dft;

common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;

common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;

common_speculative_ptr spec;

bool add_bos_token = true;

int32_t n_ctx; // total context for all clients / slots

// set to llama_model_n_swa(model)

// if swa_full is enabled, this is set to 0 to simulate a non-SWA model

int32_t n_swa;

// slots / clients

std::vector<server_slot> slots;

int trace = 0;

int slots_debug = 0;

int n_empty_consecutive = 0;

std::unique_ptr<server_prompt_cache> prompt_cache;

server_metrics metrics;

json json_ui_settings = json::object();

// Necessary similarity of prompt for slot selection

float slot_prompt_similarity = 0.0f;

std::string model_name; // name of the loaded model, to be used by API

std::set<std::string> model_aliases; // additional names for the model

std::set<std::string> model_tags; // informational tags

bool sleeping = false;

void destroy() {

spec.reset();

ctx_dft.reset();

model_dft.reset();

llama_init.reset();

ctx_tgt = nullptr;

model_tgt = nullptr;

mtmd_free(mctx);

mctx = nullptr;

llama_batch_free(batch);

}

void handle_sleeping_state(bool new_state) {

GGML_ASSERT(sleeping != new_state);

if (new_state) {

SRV_INF("%s", "server is entering sleeping state\n");

destroy();

} else {

SRV_INF("%s", "server is exiting sleeping state\n");

if (!load_model(params_base)) {

GGML_ABORT("failed to reload model after sleeping");

}

}

sleeping = new_state;

}

// load the model and initialize llama_context

// this may also be called to resume from sleeping state

bool load_model(common_params & params) {

bool is_resume = sleeping;

SRV_INF("loading model '%s'\n", params.model.path.c_str());

params_base = params;

params_base.n_outputs_max = server_n_outputs_max(params_base);

std::string & mmproj_path = params_base.mmproj.path;

bool has_mmproj = !mmproj_path.empty();

mtmd_context_params mparams = mtmd_context_params_default();

if (has_mmproj) {

mparams.use_gpu = params_base.mmproj_use_gpu;

mparams.print_timings = false;

mparams.n_threads = params_base.cpuparams.n_threads;

mparams.flash_attn_type = params_base.flash_attn_type;

mparams.warmup = params_base.warmup;

mparams.image_min_tokens = params_base.image_min_tokens;

mparams.image_max_tokens = params_base.image_max_tokens;

mparams.batch_max_tokens = params_base.mtmd_batch_max_tokens;

mparams.media_marker = get_media_marker();

}

// optionally get the memory usage of mmproj

if (has_mmproj && params_base.fit_params) {

auto mmproj_mem = mtmd_get_memory_usage(mmproj_path.c_str(), mparams);

if (!mmproj_mem.empty()) {

size_t total = 0;

for (auto & [dev, size] : mmproj_mem) {

total += size;

}

SRV_INF("[mtmd] estimated worst-case memory usage of mmproj is %.2f MiB\n", total / (1024.0 * 1024.0));

GGML_ASSERT(!params_base.fit_params_target.empty());

for (auto & [dev, size] : mmproj_mem) {

for (size_t i = 0; i < ggml_backend_dev_count(); i++) {

if (ggml_backend_dev_get(i) == dev) {

if (i < params_base.fit_params_target.size()) {

SRV_DBG("[mtmd] adding %.2f MiB to fit_params_target for device %s\n", size / (1024.0 * 1024.0), ggml_backend_dev_name(dev));

params_base.fit_params_target[i] += size;

}

break;

}

}

}

} else {

SRV_ERR("%s", "[mtmd] failed to get memory usage of mmproj\n");

}

}

// optionally reserve VRAM for the draft / MTP context before fitting the target model

if (params_base.fit_params) {

const bool spec_mtp = std::find(params_base.speculative.types.begin(),

params_base.speculative.types.end(),

COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end();

const bool has_draft = params_base.speculative.has_dft();

if (has_draft || spec_mtp) {

common_params params_dft = params_base;

bool measure_model_bytes = true;

if (has_draft) {

const auto & params_spec = params_base.speculative.draft;

params_dft.devices = params_spec.devices;

params_dft.model = params_spec.mparams;

params_dft.n_gpu_layers = params_spec.n_gpu_layers;

params_dft.cache_type_k = params_spec.cache_type_k;

params_dft.cache_type_v = params_spec.cache_type_v;

params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides;

} else {

// MTP draft context lives on the target model, only context+compute are new

measure_model_bytes = false;

}

params_dft.n_outputs_max = params_base.n_parallel;

auto mparams_dft = common_model_params_to_llama(params_dft);

auto cparams_dft = common_context_params_to_llama(params_dft);

if (spec_mtp) {

cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;

cparams_dft.type_k = params_base.speculative.draft.cache_type_k;

cparams_dft.type_v = params_base.speculative.draft.cache_type_v;

}

cparams_dft.n_rs_seq = 0;

std::vector<ggml_backend_dev_t> devs;

uint32_t hp_ngl = 0;

uint32_t hp_nct = 0;

uint32_t hp_nex = 0;

try {

auto dmd = common_get_device_memory_data(

params_dft.model.path.c_str(), &mparams_dft, &cparams_dft,

devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);

GGML_ASSERT(!params_base.fit_params_target.empty());

size_t total = 0;

std::vector<ggml_backend_dev_t> tgt_devices = params.devices;

if (tgt_devices.empty()) {

for(size_t i = 0; i < ggml_backend_dev_count(); ++i) {

tgt_devices.push_back(ggml_backend_dev_get(i));

}

}

for (size_t j = 0; j < devs.size(); ++j) {

const size_t bytes = (measure_model_bytes ? dmd[j].model : 0) + dmd[j].context + dmd[j].compute;

total += bytes;

for (size_t i = 0; i < tgt_devices.size(); i++) {

if (tgt_devices[i] == devs[j]) {

SRV_DBG("[spec] adding %.2f MiB to fit_params_target for device %s\n",

bytes / (1024.0 * 1024.0), ggml_backend_dev_name(devs[j]));

params_base.fit_params_target[i] += bytes;

break;

}

}

}

SRV_INF("[spec] estimated memory usage of %s is %.2f MiB\n",

has_draft ? "draft model" : "MTP context",

total / (1024.0 * 1024.0));

} catch (const std::exception & e) {

SRV_WRN("[spec] failed to measure %s memory: %s\n",

has_draft ? "draft model" : "MTP context", e.what());

}

}

}

llama_init = common_init_from_params(params_base);

model_tgt = llama_init->model();

ctx_tgt = llama_init->context();

if (model_tgt == nullptr) {

SRV_ERR("failed to load model, '%s'\n", params_base.model.path.c_str());