std::string result;

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

if (devices[i] == nullptr) {

continue;

}

if (!result.empty()) result += ", ";

result += ggml_backend_dev_name(devices[i]);

}

return result.empty() ? "default" : result;

common_speculative_type type;

common_params_speculative params;

common_speculative_config(common_speculative_type t,

const common_params_speculative & p = common_params_speculative{}) : type(t), params(p) {}

const llama_model * model_tgt,

const llama_model * model_dft) {

const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt);

const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft);

const auto vocab_type_tgt = llama_vocab_type(vocab_tgt);

LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt);

const auto vocab_type_dft = llama_vocab_type(vocab_dft);

LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft);

if (vocab_type_tgt != vocab_type_dft) {

LOG_WRN("%s: draft model vocab type must match target model to use speculation but "

"vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt);

return false;

}

if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||

(llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) {

LOG_WRN("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",

__func__,

llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft),

llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft));

return false;

}

if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||

(llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) {

LOG_WRN("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",

__func__,

llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft),

llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft));

return false;

}

{

const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt);

const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft);

const int vocab_diff = n_vocab_tgt > n_vocab_dft

? n_vocab_tgt - n_vocab_dft

: n_vocab_dft - n_vocab_tgt;

if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {

LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__);

LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",

n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);

return false;

}

for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) {

const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i);

const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);

if (std::strcmp(token_text_tgt, token_text_dft) != 0) {

LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__);

LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i,

common_token_to_piece(vocab_tgt, i).c_str(),

common_token_to_piece(vocab_dft, i).c_str());

return false;

}

}

}

return true;

const common_speculative_type type;

uint32_t n_seq;

size_t n_call_begin = 0; // number of times this implementation was called for refresh.

size_t n_call_draft = 0; // number of times this implementation was called for generation.

size_t n_call_accept = 0; // number of times this implementation was called for accumulation.

size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation.

size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model.

size_t n_gen_tokens = 0; // number of tokens generated by this implementation.

size_t n_acc_tokens = 0; // number of tokens accepted by the target model.

// TODO: track performance of most recent calls

const bool gen_perf = true; // whether to generate performance stats.

int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds.

int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds.

int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds.

common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {}

virtual ~common_speculative_impl() = default;

virtual void begin(llama_seq_id seq_id, const llama_tokens & prompt) = 0;

virtual bool process(const llama_batch & batch) = 0;

virtual void draft(common_speculative_draft_params_vec & dparams) = 0;

virtual void accept(llama_seq_id seq_id, uint16_t n_accepted, bool is_other) = 0;

// true if this implementation requires the target context to extract post-norm embeddings

virtual bool need_embd() const = 0;

// true if this implementation requires the target context to extract pre-norm embeddings

virtual bool need_embd_nextn() const { return false; }

common_params_speculative_draft params;

llama_batch batch;

std::vector<common_sampler_ptr> smpls;

common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq)

: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq)

, params(params.draft)

{

auto * ctx_dft = this->params.ctx_dft;

auto * ctx_tgt = this->params.ctx_tgt;

LOG_INF("%s: adding speculative implementation 'draft-simple'\n", __func__);

LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);

LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__,

this->params.n_gpu_layers,

ggml_type_name(this->params.cache_type_k),

ggml_type_name(this->params.cache_type_v),

ctx_tgt ? "yes" : "no",

ctx_dft ? "yes" : "no",

common_speculative_get_devices_str(this->params.devices).c_str());

batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1);

// TODO: optimize or pass from outside?

// {

// common_params_sampling params;

// params.no_perf = false;

//

// params.top_k = 40;

// params.top_p = 0.9;

//

// params.samplers = {

// COMMON_SAMPLER_TYPE_TOP_K,

// COMMON_SAMPLER_TYPE_TOP_P,

// COMMON_SAMPLER_TYPE_INFILL,

// };

//

// result->smpl = common_sampler_init(llama_get_model(ctx_dft), params);

// }

smpls.resize(n_seq);

for (auto & smpl : smpls) {

common_params_sampling params;

params.no_perf = false;

params.top_k = 10;

params.samplers = {

COMMON_SAMPLER_TYPE_TOP_K,

};

smpl.reset(common_sampler_init(llama_get_model(ctx_dft), params));

}

const bool vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft));

LOG_DBG("%s: vocab_cmpt = %d\n", __func__, vocab_cmpt);

if (!vocab_cmpt) {

LOG_ERR("%s: the target and draft vocabs are not compatible\n", __func__);

throw std::runtime_error("draft model vocab type must match target model to use speculation");

}

if (n_seq != llama_n_seq_max(ctx_dft)) {

LOG_ERR("%s: n_seq mismatch: %d != %d\n", __func__, n_seq, llama_n_seq_max(ctx_dft));

throw std::runtime_error("the draft model number of sequences is incompatible with the speculative n_seq");

}

}

~common_speculative_impl_draft_simple() override {

llama_batch_free(batch);

}

void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {

// noop

}

bool process(const llama_batch & batch) override {

auto * ctx_dft = params.ctx_dft;

const int ret = llama_decode(ctx_dft, batch);

if (ret != 0) {

LOG_ERR("%s: failed to decode draft batch, ret = %d\n", __func__, ret);

return false;

}

return true;

}

void draft(common_speculative_draft_params_vec & dparams) override {

auto & ctx_dft = params.ctx_dft;

common_batch_clear(batch);

// keep track of which sequences are still drafting

int n_drafting = 0;

std::vector<bool> drafting(n_seq);

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

auto & dp = dparams[seq_id];

if (!dp.drafting) {

continue;

}

n_drafting++;

drafting[seq_id] = true;

common_sampler_reset(smpls[seq_id].get());

common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true);

}

int ret = llama_decode(ctx_dft, batch);

if (ret != 0) {

LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);

return;

}

int i = 0;

while (n_drafting > 0) {

int i_batch = 0;

common_batch_clear(batch);

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

if (!drafting[seq_id]) {

continue;

}

auto * smpl = smpls[seq_id].get();

common_sampler_sample(smpl, ctx_dft, i_batch, true);

++i_batch;

const auto * cur_p = common_sampler_get_candidates(smpl, true);

for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {

LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",

seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,

common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());

}

// add drafted token for each sequence

const llama_token id = cur_p->data[0].id;

// only collect very high-confidence draft tokens

if (cur_p->data[0].p < params.p_min) {

drafting[seq_id] = false;

n_drafting--;

continue;

}

common_sampler_accept(smpl, id, true);

auto & dp = dparams.at(seq_id);

auto & result = *dp.result;

result.push_back(id);

if ((params.n_max <= (int) result.size()) ||

(dp.n_max > 0 && dp.n_max <= (int) result.size())) {

drafting[seq_id] = false;

n_drafting--;

continue;

}

common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true);

}

if (batch.n_tokens == 0) {

break;

}

// evaluate the drafted tokens on the draft model

ret = llama_decode(ctx_dft, batch);

if (ret != 0) {

LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret);

break;

}

++i;

}

for (auto & dp : dparams) {

if (!dp.drafting) {

continue;

}

if (dp.result->size() < (size_t) params.n_min) {

dp.result->clear();

}

}

}

void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {

// noop

}

bool need_embd() const override {

return false;

}

common_params_speculative_draft params;

llama_batch batch;

std::vector<common_sampler_ptr> smpls;

int32_t n_embd_dec = 0; // draft hidden size

int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size

int32_t n_embd_tgt = 0; // target model hidden size

const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices

uint32_t target_layer_ids_n = 0;

// [per-seq] deferred boundary state

std::vector<std::vector<float>> pending_g_last;

std::vector<llama_pos> pending_pos_last;

// [per-seq] snapshot of the most recent process()'s encoder output

std::vector<std::vector<float>> verify_g; // [n_seq][n_rows * n_embd_dec]

std::vector<llama_pos> verify_pos_first; // [n_seq] — pos of verify_g[seq][0]

std::vector<int32_t> verify_g_rows; // [n_seq] — number of rows

// scratch buffer for concatenated target features [n_tokens, n_embd_enc]

std::vector<float> features_buf;

std::vector<float> g_embd_buf;

common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq)

: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq)

, params(params.draft)

{

LOG_INF("%s: adding speculative implementation 'draft-eagle3'\n", __func__);

LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, params.draft.n_max, params.draft.n_min, params.draft.p_min);

auto * ctx_tgt = this->params.ctx_tgt;

auto * ctx_dft = this->params.ctx_dft;

GGML_ASSERT(ctx_tgt && ctx_dft && "EAGLE3 requires ctx_tgt and ctx_dft to be set");

const llama_model * model_dft = llama_get_model(ctx_dft);

const llama_model * model_tgt = llama_get_model(ctx_tgt);

target_layer_ids = llama_model_target_layer_ids (model_dft);

target_layer_ids_n = llama_model_target_layer_ids_n(model_dft);

if (target_layer_ids_n != 3) {

throw std::runtime_error("draft model is not eagle3 (expected 3 extract layers, got " +

std::to_string(target_layer_ids_n) + ")");

}

n_embd_tgt = llama_model_n_embd(model_tgt);

n_embd_dec = llama_model_n_embd(model_dft);

n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;

const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);

batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);

// llama_batch_init allocates only one of token/embd; eagle3 decoder needs both.

// TODO: fix, how to call without malloc

batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b);

smpls.resize(n_seq);

for (auto & s : smpls) {

common_params_sampling sparams;

sparams.no_perf = false;

sparams.top_k = 10;

sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };

s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams));

}

// turn on extraction of the target layers' input embeddings

for (uint32_t k = 0; k < target_layer_ids_n; ++k) {

llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);

}

// turn on extraction of the draft model's pre-norm hidden state

// (used both for the encoder output g_embd and the decoder pre-norm output).

llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);

pending_g_last.assign(n_seq, std::vector<float>(n_embd_dec, 0.0f));

pending_pos_last.assign(n_seq, -1);

verify_g.assign(n_seq, std::vector<float>());

verify_pos_first.assign(n_seq, -1);

verify_g_rows.assign(n_seq, 0);

}

~common_speculative_impl_draft_eagle3() override {

if (batch.token != nullptr) {

free(batch.token);

batch.token = nullptr;

}

llama_batch_free(batch);

}

void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {

const int32_t N = (int32_t) prompt.size();

if (N <= 0) {

return;

}

// expected state after prefill: ctx_dft has pos 0..N-2 (last position is deferred to

// draft()'s seed step). Warn only if more than one position is missing.

auto * ctx_dft = this->params.ctx_dft;

const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);

if (pos_max < N - 2) {

LOG_WRN("%s: ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. "

"Drafts may degrade.\n",

__func__, (int) pos_max, N - 2);

}

}

bool process(const llama_batch & batch_in) override {

if (batch_in.n_tokens <= 0) {

return true;

}

if (batch_in.token == nullptr || batch_in.embd != nullptr) {

return true;

}

const int32_t n_tokens = batch_in.n_tokens;

// i_batch_beg[seq] / i_batch_end[seq]: inclusive batch indices of this seq's

// first/last token in batch_in. Assumes per-seq tokens are contiguous within

// the ubatch (server's default ordering).

std::vector<int32_t> i_batch_beg(n_seq, -1);

std::vector<int32_t> i_batch_end(n_seq, -1);

for (int k = 0; k < n_tokens; ++k) {

GGML_ASSERT(batch_in.n_seq_id[k] == 1);

const llama_seq_id seq_id = batch_in.seq_id[k][0];

if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {

continue;

}

i_batch_end[seq_id] = k;

if (i_batch_beg[seq_id] < 0) {

i_batch_beg[seq_id] = k;

}

}

auto * ctx_tgt = this->params.ctx_tgt;

auto * ctx_dft = this->params.ctx_dft;

// Interleave each extract_layer's hidden state into a contiguous buffer of

// shape [n_tokens, target_layer_ids_n * n_embd_tgt]. Then run EAGLE3 encoder

// to get one g_embd row per token.

features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f);

for (uint32_t k = 0; k < target_layer_ids_n; ++k) {

const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);

if (!layer) {

GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]);

}

for (int32_t i = 0; i < n_tokens; ++i) {

float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;

const float * src = layer + (size_t) i * n_embd_tgt;

std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));

}

}

g_embd_buf.resize((size_t) n_tokens * n_embd_dec);

// llama_encode() requires the full encoder batch to fit in n_ubatch.

// Allow batch > ubatch: eagle3's per-token encoder can be chunked safely.

const int32_t n_ubatch_dft = (int32_t) llama_n_ubatch(ctx_dft);

for (int32_t i = 0; i < n_tokens; i += n_ubatch_dft) {

const int32_t n_chunk = std::min(n_ubatch_dft, n_tokens - i);

llama_batch enc_batch = {

/*.n_tokens =*/ n_chunk,

/*.token =*/ nullptr,

/*.embd =*/ features_buf.data() + (size_t) i * n_embd_enc,

/*.pos =*/ nullptr,

/*.n_seq_id =*/ nullptr,

/*.seq_id =*/ nullptr,

/*.logits =*/ nullptr,

};

const int32_t rc = llama_encode(ctx_dft, enc_batch);

if (rc != 0) {

LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",

__func__, rc, (int) n_chunk, (int) i);

return false;

}

// g_embd has shape [n_chunk, n_embd_dec] in ctx_dft's pre-norm embeddings buffer.

const float * g_embd_chunk = llama_get_embeddings_nextn(ctx_dft);

GGML_ASSERT(g_embd_chunk && "EAGLE3 encoder produced no output.");

std::memcpy(g_embd_buf.data() + (size_t) i * n_embd_dec,

g_embd_chunk,

(size_t) n_chunk * n_embd_dec * sizeof(float));

}

const float * g_embd = g_embd_buf.data();

const size_t row_bytes = (size_t) n_embd_dec * sizeof(float);

// EAGLE3 decoder input convention: at memory pos P the input pair is

// (token[P+1], g_embd[P]). This shifts the token index "left by one" relative to g_embd.

//

// Per seq, in order:

// (a) cross-ubatch bridge — when applicable, write the previously-deferred

// pos using this ubatch's first token + pending_g_last.

// (b) main write loop — for k in [beg, end-1], write (token[k+1], g_embd[k])

// at pos[k]. The last training pos (k=end) is left unwritten = new

// deferred boundary, completed by the next process() or draft() call.

// (c) refresh deferred state — stash this ubatch's full g_embd into verify_g,

// update pending_g_last / pending_pos_last to the last row.

common_batch_clear(batch);

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

const int32_t beg = i_batch_beg[seq_id];

const int32_t end = i_batch_end[seq_id];

if (beg < 0 || end < 0) {

continue;

}

// cross-ubatch bridge — complete the prior ubatch's deferred boundary.

// Fires iff all three preconditions hold:

// 1) pending_pos_last >= 0

// 2) pending_pos_last + 1 == pos[beg]

// 3) pending_pos_last > dft_pos_max // TODO: is this check needed?

const llama_pos pending_pos = pending_pos_last[seq_id];

if (pending_pos >= 0 && pending_pos + 1 == batch_in.pos[beg]) {

const llama_pos dft_pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);

if (pending_pos > dft_pos_max) {

common_batch_add(batch, batch_in.token[beg], pending_pos, { seq_id }, /*logits=*/ false);

std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,

pending_g_last[seq_id].data(), row_bytes);

}

}

for (int32_t k = beg; k < end; ++k) {

common_batch_add(batch, batch_in.token[k + 1], batch_in.pos[k], { seq_id }, /*logits=*/ false);

std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,

g_embd + (size_t) k * n_embd_dec, row_bytes);

}

// refresh deferred state

const int32_t n_rows = end - beg + 1;

verify_pos_first[seq_id] = batch_in.pos[beg];

pending_pos_last[seq_id] = batch_in.pos[end];

verify_g_rows[seq_id] = n_rows;

verify_g[seq_id].resize((size_t) n_rows * n_embd_dec, 0.0f);

std::memcpy(verify_g[seq_id].data(), g_embd + (size_t) beg * n_embd_dec, row_bytes * n_rows);

std::memcpy(pending_g_last[seq_id].data(), g_embd + (size_t) end * n_embd_dec, row_bytes);

}

if (batch.n_tokens > 0) {

const int32_t rc = llama_decode(ctx_dft, batch);

if (rc != 0) {

LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",

__func__, rc, (int) batch.n_tokens, (int) batch_in.pos[0]);

return false;

}

}

return true;

}

void draft(common_speculative_draft_params_vec & dparams) override {

auto & ctx_dft = params.ctx_dft;

common_batch_clear(batch);

// keep track of which sequences are still drafting

int n_drafting = 0;

std::vector<bool> drafting(n_seq);

const size_t row_bytes = (size_t) n_embd_dec * sizeof(float);

// Complete the deferred boundary pair (dp.id_last, pending_g_last) at memory

// pos pending_pos_last. dp.id_last is target's freshest sample (= corrected

// token after verify, or first generated token after prefill), matching the

// EAGLE3 input convention (token[P+1], g_embd[P]) at pos P.

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

auto & dp = dparams[seq_id];

if (!dp.drafting) {

continue;

}

if (pending_pos_last[seq_id] < 0) {

continue;

}

n_drafting++;

drafting[seq_id] = true;

common_sampler_reset(smpls[seq_id].get());

llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, pending_pos_last[seq_id], -1);

common_batch_add(batch, dp.id_last, pending_pos_last[seq_id], { seq_id }, true);

std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,

pending_g_last[seq_id].data(),

row_bytes);

}

if (batch.n_tokens == 0) {

return;

}

int ret = llama_decode(ctx_dft, batch);

if (ret != 0) {

LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);

return;

}

int i = 0;

while (n_drafting > 0) {

int i_batch = 0;

common_batch_clear(batch);

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

if (!drafting[seq_id]) {

continue;

}

auto * smpl = smpls[seq_id].get();

common_sampler_sample(smpl, ctx_dft, i_batch, true);

// pre-norm hidden state of this position becomes g_embd for the next step

const float * prenorm = llama_get_embeddings_nextn_ith(ctx_dft, i_batch);

++i_batch;

const auto * cur_p = common_sampler_get_candidates(smpl, true);

for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {

LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",

seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,

common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());

}

const llama_token id = cur_p->data[0].id;

// only collect very high-confidence draft tokens

// (configurable via --spec-draft-p-min, set to 0.0 to disable early-stop)

if (cur_p->data[0].p < params.p_min) {

drafting[seq_id] = false;

n_drafting--;

continue;

}

common_sampler_accept(smpl, id, true);

auto & dp = dparams.at(seq_id);

auto & result = *dp.result;

result.push_back(id);

if (params.n_max <= (int) result.size()) {

drafting[seq_id] = false;

n_drafting--;

continue;

}

common_batch_add(batch, id, pending_pos_last[seq_id] + (i + 1), { seq_id }, true);

std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec, prenorm, row_bytes);

}

if (batch.n_tokens == 0) {

break;

}

ret = llama_decode(ctx_dft, batch);

if (ret != 0) {

LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret);

break;

}

++i;

}

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

auto & dp = dparams[seq_id];

if (!dp.drafting) {

continue;

}

if (dp.result->size() < (size_t) params.n_min) {

dp.result->clear();

}

}

}

void accept(llama_seq_id seq_id, uint16_t n_accepted, bool /*is_other*/) override {

if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {

return;

}

const int32_t n_rows = verify_g_rows[seq_id];

if (n_rows <= 0) {

return;

}

const int32_t i_g = std::min<int32_t>(n_accepted, n_rows - 1);

pending_pos_last[seq_id] = verify_pos_first[seq_id] + i_g;

std::memcpy(pending_g_last[seq_id].data(),

verify_g[seq_id].data() + (size_t) i_g * n_embd_dec,

(size_t) n_embd_dec * sizeof(float));

}

bool need_embd() const override {

return false;

}

common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft)

llama_batch batch;

std::vector<common_sampler_ptr> smpls;

// backend sampler chain per seq, attached to ctx_dft

std::vector<llama_sampler *> backend_chains;

int32_t n_embd = 0;

bool is_mem_shared = false;

// Per-sequence cross-batch carryover: pair (h_p, x_{p+1}) at MTP pos p+1.

// The last h-row of one process() call needs the first token of the NEXT

// call to pair with, so it's stashed here until that next call fires.

std::vector<std::vector<float>> pending_h; // [n_seq][n_embd]

std::vector<int32_t> i_batch_beg;

std::vector<int32_t> i_batch_end;

// Hidden rows from the most recent target verification batch, grouped by seq.

// Row 0 corresponds to the sampled token, row N to the Nth accepted draft token.

std::vector<std::vector<float>> verify_h;

std::vector<int32_t> verify_h_rows;

// Per-seq draft length from the last draft() call, used in accept() to

// roll back ctx_dft's recurrent state past the AR draft's redundant

// pre-advancement before process() mirrored the verify batch.

std::vector<uint16_t> last_n_drafted;

common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq)

: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq)

, params(params.draft)

{

auto * ctx_tgt = this->params.ctx_tgt;

auto * ctx_dft = this->params.ctx_dft;

GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set");

n_embd = llama_model_n_embd_out(llama_get_model(ctx_dft));

GGML_ASSERT(n_embd == llama_model_n_embd(llama_get_model(ctx_tgt)) &&

"MTP input row width must match the target h_nextn width");

LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__);

LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling);

LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__,

this->params.n_gpu_layers,

ggml_type_name(this->params.cache_type_k),

ggml_type_name(this->params.cache_type_v),

ctx_tgt ? "yes" : "no",

ctx_dft ? "yes" : "no",

common_speculative_get_devices_str(this->params.devices).c_str());

const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);

batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd, /*n_seq_max=*/ 1);

// llama_batch_init allocates only one of token/embd; MTP needs both.

// TODO: fix, how to call without malloc

batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b);

smpls.resize(n_seq);

for (auto & s : smpls) {

common_params_sampling sparams;

sparams.no_perf = false;

sparams.top_k = 10;

sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };

s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams));

}

// offload draft sampling to the backend

backend_chains.assign(n_seq, nullptr);

if (this->params.backend_sampling) {

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());

llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));

if (!llama_set_sampler(ctx_dft, seq_id, chain)) {

LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id);

llama_sampler_free(chain);

chain = nullptr;

}

backend_chains[seq_id] = chain;

}

}

llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);

llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);

is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;

pending_h.assign(n_seq, std::vector<float>(n_embd, 0.0f));

i_batch_beg.assign(n_seq, -1);

i_batch_end.assign(n_seq, -1);

verify_h.assign(n_seq, {});

verify_h_rows.assign(n_seq, 0);

last_n_drafted.assign(n_seq, 0);

}

~common_speculative_impl_draft_mtp() override {

auto * ctx_dft = this->params.ctx_dft;

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {

if (backend_chains[seq_id] == nullptr) {

continue;

}

if (ctx_dft) {

llama_set_sampler(ctx_dft, seq_id, nullptr);

}

llama_sampler_free(backend_chains[seq_id]);

}

backend_chains.clear();

if (batch.token != nullptr) {

free(batch.token);

batch.token = nullptr;

}

llama_batch_free(batch);

}

void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {

const int32_t N = (int32_t) prompt.size();

if (N <= 0) {

return;

}

auto * ctx_dft = this->params.ctx_dft;

const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);

if (pos_max < N - 1 && !is_mem_shared) {

LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - "

"process() hook may not have run on every prefill ubatch "

"(need_embd / logits=1 on every prompt position?). "

"Drafts may degrade.\n",

__func__, (int) pos_max, N - 1);

}

}

bool process(const llama_batch & batch_in) override {

if (batch_in.n_tokens <= 0) {

return true;

}

// TODO: how to make it work with vision tokens?

if (batch_in.token == nullptr || batch_in.embd != nullptr) {

return true;

}

const int32_t n_tokens = batch_in.n_tokens;

// remember the frist and last batch index for each sequence

std::fill(i_batch_beg.begin(), i_batch_beg.end(), -1);

std::fill(i_batch_end.begin(), i_batch_end.end(), -1);

for (int k = 0; k < n_tokens; ++k) {

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {

GGML_ASSERT(batch_in.n_seq_id[k] == 1);

if (batch_in.seq_id[k][0] == seq_id) {

i_batch_end[seq_id] = k;

if (i_batch_beg[seq_id] < 0) {

i_batch_beg[seq_id] = k;

}

}

}

}

auto * ctx_tgt = this->params.ctx_tgt;

auto * ctx_dft = this->params.ctx_dft;

const size_t row_bytes = (size_t) n_embd * sizeof(float);

// if kv is shared with target (e.g Gemma4), then we can skip this catch-up decode

if (!is_mem_shared) {

common_batch_clear(batch);

for (int k = 0; k < n_tokens; ++k) {

common_batch_add(batch, batch_in.token[k], batch_in.pos[k], { batch_in.seq_id[k][0] }, 0);

}

// shift the tgt embeddings to the right by one position

// assumes that the tokens in the batch are sequential for each sequence

// i.e. we cannot have seq_id like this: [0, 0, 0, 1, 1, 0, 1, 1]

// ^--- this is a problem