Felipe97/llama-cpp-compiled
01.1k
1#include "speculative.h"2 3#include "common.h"4#include "ggml.h"5#include "ggml-cpp.h"6#include "llama.h"7#include "log.h"8#include "ngram-cache.h"9#include "ngram-map.h"10#include "ngram-mod.h"11#include "sampling.h"12 13#include "../src/llama-ext.h" // staging API: llama_set_embeddings_nextn / llama_get_embeddings_nextn_ith (used by MTP)14 15#include <algorithm>16#include <cassert>17#include <cmath>18#include <cstring>19#include <iomanip>20#include <map>21#include <cinttypes>22 23#define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)24#define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)25#define SPC_INF(fmt, ...) LOG_INF("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)26#define SPC_WRN(fmt, ...) LOG_WRN("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)27#define SPC_ERR(fmt, ...) LOG_ERR("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)28#define SPC_CNT(fmt, ...) LOG_CNT("" fmt, __VA_ARGS__)29 30#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 12831#define SPEC_VOCAB_CHECK_START_TOKEN_ID 532 33const std::map<std::string, common_speculative_type> common_speculative_type_from_name_map = {34 {"none", COMMON_SPECULATIVE_TYPE_NONE},35 {"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE},36 {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},37 {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP},38 {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},39 {"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK},40 {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},41 {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},42 {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},43 {"ngram-mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD},44 {"ngram-cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE}45};46 47static std::string common_speculative_get_devices_str(const std::vector<ggml_backend_dev_t> & devices) {48 std::string result;49 for (size_t i = 0; i < devices.size(); i++) {50 if (devices[i] == nullptr) {51 continue;52 }53 if (!result.empty()) result += ", ";54 result += ggml_backend_dev_name(devices[i]);55 }56 return result.empty() ? "default" : result;57}58 59struct common_speculative_config {60 common_speculative_type type;61 common_params_speculative params;62 63 common_speculative_config(common_speculative_type t,64 const common_params_speculative & p = common_params_speculative{}) : type(t), params(p) {}65};66 67static bool common_speculative_are_compatible(68 const llama_model * model_tgt,69 const llama_model * model_dft) {70 const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt);71 const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft);72 73 const auto vocab_type_tgt = llama_vocab_type(vocab_tgt);74 SPC_DBG("vocab_type tgt: %d\n", vocab_type_tgt);75 76 const auto vocab_type_dft = llama_vocab_type(vocab_dft);77 SPC_DBG("vocab_type dft: %d\n", vocab_type_dft);78 79 if (vocab_type_tgt != vocab_type_dft) {80 SPC_WRN("draft model vocab type must match target model to use speculation but "81 "vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt);82 return false;83 }84 85 if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||86 (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) {87 SPC_WRN("draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",88 llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft),89 llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft));90 return false;91 }92 93 if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||94 (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) {95 SPC_WRN("draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",96 llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft),97 llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft));98 return false;99 }100 101 {102 const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt);103 const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft);104 const int vocab_diff = n_vocab_tgt > n_vocab_dft105 ? n_vocab_tgt - n_vocab_dft106 : n_vocab_dft - n_vocab_tgt;107 108 if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {109 SPC_DBG("draft model vocab must closely match target model to use speculation but "110 "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",111 n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);112 return false;113 }114 115 for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) {116 const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i);117 const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);118 119 if (std::strcmp(token_text_tgt, token_text_dft) != 0) {120 SPC_DBG("draft model vocab must match target model to use speculation but "121 "token %d content differs - target '%s', draft '%s'\n", i,122 common_token_to_piece(vocab_tgt, i).c_str(),123 common_token_to_piece(vocab_dft, i).c_str());124 return false;125 }126 }127 }128 129 return true;130}131 132using common_speculative_draft_params_vec = std::vector<common_speculative_draft_params>;133 134// state of an implementation of speculative decoding135//136// each implementation has a unique type and a state that is implementation-specific137// in a subclass of common_speculative_impl138struct common_speculative_impl {139 const common_speculative_type type;140 141 uint32_t n_seq;142 int32_t n_max; // maximum draft length after implementation-specific limits143 144 size_t n_call_begin = 0; // number of times this implementation was called for refresh.145 size_t n_call_draft = 0; // number of times this implementation was called for generation.146 size_t n_call_accept = 0; // number of times this implementation was called for accumulation.147 148 size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation.149 size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model.150 size_t n_gen_tokens = 0; // number of tokens generated by this implementation.151 size_t n_acc_tokens = 0; // number of tokens accepted by the target model.152 153 std::vector<size_t> n_acc_tokens_per_pos; // number of tokens accepted per draft position.154 155 // TODO: track performance of most recent calls156 const bool gen_perf = true; // whether to generate performance stats.157 158 int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds.159 int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds.160 int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds.161 162 common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {}163 164 virtual ~common_speculative_impl() = default;165 166 virtual void begin(llama_seq_id seq_id, const llama_tokens & prompt) = 0;167 168 virtual bool process(const llama_batch & batch) = 0;169 170 virtual void draft(common_speculative_draft_params_vec & dparams) = 0;171 172 virtual void accept(llama_seq_id seq_id, uint16_t n_accepted, bool is_other) = 0;173 174 // (optional) serialize/restore per-seq internal state (e.g. eagle3's deferred boundary).175 virtual bool get_state(llama_seq_id /*seq_id*/, std::vector<uint8_t> & /*data*/) const { return false; }176 virtual void set_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*data*/) {}177};178 179struct common_speculative_impl_draft_simple : public common_speculative_impl {180 common_params_speculative_draft params;181 182 llama_batch batch;183 184 std::vector<common_sampler_ptr> smpls;185 186 common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq)187 : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq, params.draft.n_max)188 , params(params.draft)189 {190 auto * ctx_dft = this->params.ctx_dft;191 auto * ctx_tgt = this->params.ctx_tgt;192 193 if (!ctx_dft) {194 throw std::runtime_error("draft-simple requires a draft context");195 }196 197 SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n");198 SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min);199 SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",200 this->params.n_gpu_layers,201 ggml_type_name(this->params.cache_type_k),202 ggml_type_name(this->params.cache_type_v),203 ctx_tgt ? "yes" : "no",204 ctx_dft ? "yes" : "no",205 common_speculative_get_devices_str(this->params.devices).c_str());206 207 batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1);208 209 // TODO: optimize or pass from outside?210 // {211 // common_params_sampling params;212 // params.no_perf = false;213 //214 // params.top_k = 40;215 // params.top_p = 0.9;216 //217 // params.samplers = {218 // COMMON_SAMPLER_TYPE_TOP_K,219 // COMMON_SAMPLER_TYPE_TOP_P,220 // COMMON_SAMPLER_TYPE_INFILL,221 // };222 //223 // result->smpl = common_sampler_init(llama_get_model(ctx_dft), params);224 // }225 226 smpls.resize(n_seq);227 for (auto & smpl : smpls) {228 common_params_sampling params;229 params.no_perf = false;230 params.top_k = 10;231 params.samplers = {232 COMMON_SAMPLER_TYPE_TOP_K,233 };234 235 smpl.reset(common_sampler_init(llama_get_model(ctx_dft), params));236 }237 238 const bool vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft));239 SPC_DBG("vocab_cmpt = %d\n", vocab_cmpt);240 241 if (!vocab_cmpt) {242 SPC_ERR("%s", "the target and draft vocabs are not compatible\n");243 244 throw std::runtime_error("draft model vocab type must match target model to use speculation");245 }246 247 if (n_seq != llama_n_seq_max(ctx_dft)) {248 SPC_ERR("n_seq mismatch: %d != %d\n", n_seq, llama_n_seq_max(ctx_dft));249 250 throw std::runtime_error("the draft model number of sequences is incompatible with the speculative n_seq");251 }252 }253 254 ~common_speculative_impl_draft_simple() override {255 llama_batch_free(batch);256 }257 258 void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {259 // noop260 }261 262 bool process(const llama_batch & batch) override {263 auto * ctx_dft = params.ctx_dft;264 265 llama_batch batch_dft = batch;266 batch_dft.logits = nullptr;267 268 const int ret = llama_decode(ctx_dft, batch_dft);269 270 if (ret != 0) {271 SPC_ERR("failed to decode draft batch, ret = %d\n", ret);272 273 return false;274 }275 276 return true;277 }278 279 void draft(common_speculative_draft_params_vec & dparams) override {280 auto & ctx_dft = params.ctx_dft;281 282 common_batch_clear(batch);283 284 // keep track of which sequences are still drafting285 int n_drafting = 0;286 std::vector<bool> drafting(n_seq);287 288 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {289 auto & dp = dparams[seq_id];290 291 if (!dp.drafting) {292 continue;293 }294 295 n_drafting++;296 drafting[seq_id] = true;297 common_sampler_reset(smpls[seq_id].get());298 299 common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true);300 }301 302 int ret = llama_decode(ctx_dft, batch);303 if (ret != 0) {304 SPC_ERR("llama_decode returned %d\n", ret);305 return;306 }307 308 int i = 0;309 310 while (n_drafting > 0) {311 int i_batch = 0;312 313 common_batch_clear(batch);314 315 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {316 if (!drafting[seq_id]) {317 continue;318 }319 320 auto * smpl = smpls[seq_id].get();321 322 common_sampler_sample(smpl, ctx_dft, i_batch, true);323 ++i_batch;324 325 const auto * cur_p = common_sampler_get_candidates(smpl, true);326 327 for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {328 SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",329 seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,330 common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());331 }332 333 // add drafted token for each sequence334 const llama_token id = cur_p->data[0].id;335 336 // only collect very high-confidence draft tokens337 if (cur_p->data[0].p < params.p_min) {338 drafting[seq_id] = false;339 n_drafting--;340 341 continue;342 }343 344 common_sampler_accept(smpl, id, true);345 346 auto & dp = dparams.at(seq_id);347 auto & result = *dp.result;348 349 result.push_back(id);350 351 if ((params.n_max <= (int) result.size()) ||352 (dp.n_max > 0 && dp.n_max <= (int) result.size())) {353 drafting[seq_id] = false;354 n_drafting--;355 continue;356 }357 358 common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true);359 }360 361 if (batch.n_tokens == 0) {362 break;363 }364 365 // evaluate the drafted tokens on the draft model366 ret = llama_decode(ctx_dft, batch);367 if (ret != 0) {368 SPC_ERR("llama_decode[%d] returned %d\n", i, ret);369 break;370 }371 372 ++i;373 }374 375 for (auto & dp : dparams) {376 if (!dp.drafting) {377 continue;378 }379 380 if (dp.result->size() < (size_t) params.n_min) {381 dp.result->clear();382 }383 }384 }385 386 void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {387 // noop388 }389};390 391 392// EAGLE3 speculative decoding state393//394// Input of draft decoder: (This is different compared to MTP)395// At "pos P", the decoder takes input pair (t_{P+1}, g_P), with RoPE at P.396// - t_{P+1} = token at sequence pos P+1 (the *next* token after P)397// - g_P = encoder output = projection of target's extracted hidden states at P398//399// Deferred boundary (MTP doesn't have this issue):400// Within a single process() call with n_tokens, we can only write decoder KV for401// training pos 0..n_tokens-2. The last training pos (n_tokens-1) needs t_{n_tokens}402// which lies *outside* this batch — it is the token target will sample next or the first token from next ubatch.403// So the last training pos of each process() call is *deferred* to whichever next call has404// the missing token in hand:405// - multi-ubatch prefill: the next process()'s first token completes the pair406// (handled by the per-seq "cross-ubatch bridge")407// - single-ubatch prefill / after verify: draft()'s seed step uses "dp.id_last"408// (target's freshest sample) to complete the pair409//410// Per-seq carry-over state:411// pending_g_last [n_embd_dec] ┐ the deferred boundary's (g, pos). Set by412// pending_pos_last llama_pos ┘ process() at end of ubatch (= last row);413// rebased by accept() to first-non-accepted pos.414// verify_g [N × n_embd_dec] snapshot of process()'s encoder output;415// verify_pos_first llama_pos consumed by accept() to recover the right416// verify_g_rows int32_t pending_g_last row for any n_accepted value.417//418// Performance is overall good but there is waste in verify cycle:419// process() runs encoder + decoder on the *full* verify batch including rows for420// rejected drafts. The KV at those positions is then dropped.421//422// TODO: Not sure if we need optimization for this waste?423// If so we may need hybrid stash:424// in verify mode, have process() only stash features and let draft() seed run425// encoder+decoder on n_accepted+1 rows).426struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {427 common_params_speculative_draft params;428 llama_batch batch;429 430 std::vector<common_sampler_ptr> smpls;431 432 // backend sampler chain per seq, attached to ctx_dft433 std::vector<llama_sampler *> backend_chains;434 435 int32_t n_embd_dec = 0; // draft hidden size436 int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size437 int32_t n_embd_tgt = 0; // target model hidden size438 int32_t n_layer_tgt = 0; // target model layer count439 440 const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices441 uint32_t target_layer_ids_n = 0;442 443 // [per-seq] deferred boundary state444 std::vector<std::vector<float>> pending_g_last;445 std::vector<llama_pos> pending_pos_last;446 447 // [per-seq] snapshot of the most recent process()'s encoder output448 std::vector<std::vector<float>> verify_g; // [n_seq][n_rows * n_embd_dec]449 std::vector<llama_pos> verify_pos_first; // [n_seq] — pos of verify_g[seq][0]450 std::vector<int32_t> verify_g_rows; // [n_seq] — number of rows451 452 // scratch buffer for concatenated target features [n_tokens, n_embd_enc]453 std::vector<float> features_buf;454 std::vector<float> g_embd_buf;455 456 common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq)457 : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq, params.draft.n_max)458 , params(params.draft)459 {460 SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n");461 SPC_TRC("- n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling);462 463 auto * ctx_tgt = this->params.ctx_tgt;464 auto * ctx_dft = this->params.ctx_dft;465 GGML_ASSERT(ctx_tgt && ctx_dft && "EAGLE3 requires ctx_tgt and ctx_dft to be set");466 467 const llama_model * model_dft = llama_get_model(ctx_dft);468 const llama_model * model_tgt = llama_get_model(ctx_tgt);469 470 target_layer_ids = llama_model_target_layer_ids (model_dft);471 target_layer_ids_n = llama_model_target_layer_ids_n(model_dft);472 if (target_layer_ids_n != 3) {473 throw std::runtime_error("draft model is not eagle3 (expected 3 extract layers, got " +474 std::to_string(target_layer_ids_n) + ")");475 }476 477 n_embd_tgt = llama_model_n_embd(model_tgt);478 n_embd_dec = llama_model_n_embd(model_dft);479 n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;480 n_layer_tgt = llama_model_n_layer(model_tgt);481 482 const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);483 batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);484 // llama_batch_init allocates only one of token/embd; eagle3 decoder needs both.485 // TODO: fix, how to call without malloc486 batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b);487 488 smpls.resize(n_seq);489 for (auto & s : smpls) {490 common_params_sampling sparams;491 sparams.no_perf = false;492 sparams.top_k = 10;493 sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };494 s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams));495 }496 497 // offload draft sampling to the backend498 backend_chains.assign(n_seq, nullptr);499 if (this->params.backend_sampling) {500 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {501 llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());502 llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));503 504 if (!llama_set_sampler(ctx_dft, seq_id, chain)) {505 SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id);506 llama_sampler_free(chain);507 chain = nullptr;508 }509 backend_chains[seq_id] = chain;510 }511 }512 513 // turn on extraction of the target layers' hidden states514 for (uint32_t k = 0; k < target_layer_ids_n; ++k) {515 if (target_layer_ids[k] < n_layer_tgt) {516 llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);517 } else if (target_layer_ids[k] == n_layer_tgt) {518 llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);519 } else {520 GGML_ABORT("EAGLE3: target layer id %d exceeds target n_layer %d", target_layer_ids[k], n_layer_tgt);521 }522 }523 524 // turn on extraction of the draft model's pre-norm hidden state525 // (used both for the encoder output g_embd and the decoder pre-norm output).526 llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);527 528 pending_g_last.assign(n_seq, std::vector<float>(n_embd_dec, 0.0f));529 pending_pos_last.assign(n_seq, -1);530 531 verify_g.assign(n_seq, std::vector<float>());532 verify_pos_first.assign(n_seq, -1);533 verify_g_rows.assign(n_seq, 0);534 }535 536 ~common_speculative_impl_draft_eagle3() override {537 auto * ctx_dft = this->params.ctx_dft;538 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {539 if (backend_chains[seq_id] == nullptr) {540 continue;541 }542 if (ctx_dft) {543 llama_set_sampler(ctx_dft, seq_id, nullptr);544 }545 llama_sampler_free(backend_chains[seq_id]);546 }547 backend_chains.clear();548 549 if (batch.token != nullptr) {550 free(batch.token);551 batch.token = nullptr;552 }553 llama_batch_free(batch);554 }555 556 void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {557 const int32_t N = (int32_t) prompt.size();558 if (N <= 0) {559 return;560 }561 // expected state after prefill: ctx_dft has pos 0..N-2 (last position is deferred to562 // draft()'s seed step). Warn only if more than one position is missing.563 auto * ctx_dft = this->params.ctx_dft;564 const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);565 if (pos_max < N - 2) {566 SPC_WRN("ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. "567 "Drafts may degrade.\n",568 (int) pos_max, N - 2);569 }570 }571 572 bool process(const llama_batch & batch_in) override {573 if (batch_in.n_tokens <= 0) {574 return true;575 }576 577 if (batch_in.token == nullptr || batch_in.embd != nullptr) {578 return true;579 }580 581 const int32_t n_tokens = batch_in.n_tokens;582 583 // i_batch_beg[seq] / i_batch_end[seq]: inclusive batch indices of this seq's584 // first/last token in batch_in. Assumes per-seq tokens are contiguous within585 // the ubatch (server's default ordering).586 std::vector<int32_t> i_batch_beg(n_seq, -1);587 std::vector<int32_t> i_batch_end(n_seq, -1);588 for (int k = 0; k < n_tokens; ++k) {589 GGML_ASSERT(batch_in.n_seq_id[k] == 1);590 const llama_seq_id seq_id = batch_in.seq_id[k][0];591 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {592 continue;593 }594 i_batch_end[seq_id] = k;595 if (i_batch_beg[seq_id] < 0) {596 i_batch_beg[seq_id] = k;597 }598 }599 600 auto * ctx_tgt = this->params.ctx_tgt;601 auto * ctx_dft = this->params.ctx_dft;602 603 // Interleave each extract_layer's hidden state into a contiguous buffer of604 // shape [n_tokens, target_layer_ids_n * n_embd_tgt]. Then run EAGLE3 encoder605 // to get one g_embd row per token.606 features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f);607 608 for (uint32_t k = 0; k < target_layer_ids_n; ++k) {609 const float * layer = target_layer_ids[k] < n_layer_tgt610 ? llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k])611 : llama_get_embeddings_nextn(ctx_tgt);612 if (!layer) {613 GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]);614 }615 for (int32_t i = 0; i < n_tokens; ++i) {616 float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;617 const float * src = layer + (size_t) i * n_embd_tgt;618 std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));619 }620 }621 622 g_embd_buf.resize((size_t) n_tokens * n_embd_dec);623 624 // llama_encode() requires the full encoder batch to fit in n_ubatch.625 // Allow batch > ubatch: eagle3's per-token encoder can be chunked safely.626 const int32_t n_ubatch_dft = (int32_t) llama_n_ubatch(ctx_dft);627 for (int32_t i = 0; i < n_tokens; i += n_ubatch_dft) {628 const int32_t n_chunk = std::min(n_ubatch_dft, n_tokens - i);629 630 llama_batch enc_batch = {631 /*.n_tokens =*/ n_chunk,632 /*.token =*/ nullptr,633 /*.embd =*/ features_buf.data() + (size_t) i * n_embd_enc,634 /*.pos =*/ nullptr,635 /*.n_seq_id =*/ nullptr,636 /*.seq_id =*/ nullptr,637 /*.logits =*/ nullptr,638 };639 const int32_t rc = llama_encode(ctx_dft, enc_batch);640 if (rc != 0) {641 SPC_ERR("llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",642 rc, (int) n_chunk, (int) i);643 return false;644 }645 646 // g_embd has shape [n_chunk, n_embd_dec] in ctx_dft's pre-norm embeddings buffer.647 const float * g_embd_chunk = llama_get_embeddings_nextn(ctx_dft);648 GGML_ASSERT(g_embd_chunk && "EAGLE3 encoder produced no output.");649 std::memcpy(g_embd_buf.data() + (size_t) i * n_embd_dec,650 g_embd_chunk,651 (size_t) n_chunk * n_embd_dec * sizeof(float));652 }653 654 const float * g_embd = g_embd_buf.data();655 656 const size_t row_bytes = (size_t) n_embd_dec * sizeof(float);657 658 // EAGLE3 decoder input convention: at memory pos P the input pair is659 // (token[P+1], g_embd[P]). This shifts the token index "left by one" relative to g_embd.660 //661 // Per seq, in order:662 // (a) cross-ubatch bridge — when applicable, write the previously-deferred663 // pos using this ubatch's first token + pending_g_last.664 // (b) main write loop — for k in [beg, end-1], write (token[k+1], g_embd[k])665 // at pos[k]. The last training pos (k=end) is left unwritten = new666 // deferred boundary, completed by the next process() or draft() call.667 // (c) refresh deferred state — stash this ubatch's full g_embd into verify_g,668 // update pending_g_last / pending_pos_last to the last row.669 common_batch_clear(batch);670 671 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {672 const int32_t beg = i_batch_beg[seq_id];673 const int32_t end = i_batch_end[seq_id];674 if (beg < 0 || end < 0) {675 continue;676 }677 678 // cross-ubatch bridge — complete the prior ubatch's deferred boundary.679 // Fires iff all three preconditions hold:680 // 1) pending_pos_last >= 0681 // 2) pending_pos_last + 1 == pos[beg]682 // 3) pending_pos_last > dft_pos_max // TODO: is this check needed?683 const llama_pos pending_pos = pending_pos_last[seq_id];684 if (pending_pos >= 0 && pending_pos + 1 == batch_in.pos[beg]) {685 const llama_pos dft_pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);686 if (pending_pos > dft_pos_max) {687 common_batch_add(batch, batch_in.token[beg], pending_pos, { seq_id }, /*logits=*/ false);688 std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,689 pending_g_last[seq_id].data(), row_bytes);690 }691 }692 693 for (int32_t k = beg; k < end; ++k) {694 common_batch_add(batch, batch_in.token[k + 1], batch_in.pos[k], { seq_id }, /*logits=*/ false);695 std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,696 g_embd + (size_t) k * n_embd_dec, row_bytes);697 }698 699 // refresh deferred state700 const int32_t n_rows = end - beg + 1;701 verify_pos_first[seq_id] = batch_in.pos[beg];702 pending_pos_last[seq_id] = batch_in.pos[end];703 verify_g_rows[seq_id] = n_rows;704 verify_g[seq_id].resize((size_t) n_rows * n_embd_dec, 0.0f);705 std::memcpy(verify_g[seq_id].data(), g_embd + (size_t) beg * n_embd_dec, row_bytes * n_rows);706 std::memcpy(pending_g_last[seq_id].data(), g_embd + (size_t) end * n_embd_dec, row_bytes);707 }708 709 if (batch.n_tokens > 0) {710 const int32_t rc = llama_decode(ctx_dft, batch);711 if (rc != 0) {712 SPC_ERR("llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",713 rc, (int) batch.n_tokens, (int) batch_in.pos[0]);714 return false;715 }716 }717 718 return true;719 }720 721 void draft(common_speculative_draft_params_vec & dparams) override {722 auto & ctx_dft = params.ctx_dft;723 724 common_batch_clear(batch);725 726 // keep track of which sequences are still drafting727 int n_drafting = 0;728 std::vector<bool> drafting(n_seq);729 730 const size_t row_bytes = (size_t) n_embd_dec * sizeof(float);731 732 // Complete the deferred boundary pair (dp.id_last, pending_g_last) at memory733 // pos pending_pos_last. dp.id_last is target's freshest sample (= corrected734 // token after verify, or first generated token after prefill), matching the735 // EAGLE3 input convention (token[P+1], g_embd[P]) at pos P.736 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {737 auto & dp = dparams[seq_id];738 739 if (!dp.drafting) {740 continue;741 }742 if (pending_pos_last[seq_id] < 0) {743 continue;744 }745 746 n_drafting++;747 drafting[seq_id] = true;748 common_sampler_reset(smpls[seq_id].get());749 750 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, pending_pos_last[seq_id], -1);751 752 common_batch_add(batch, dp.id_last, pending_pos_last[seq_id], { seq_id }, true);753 std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,754 pending_g_last[seq_id].data(),755 row_bytes);756 }757 758 if (batch.n_tokens == 0) {759 return;760 }761 762 int ret = llama_decode(ctx_dft, batch);763 if (ret != 0) {764 SPC_ERR("llama_decode returned %d\n", ret);765 return;766 }767 768 int i = 0;769 770 while (n_drafting > 0) {771 int i_batch = 0;772 773 common_batch_clear(batch);774 775 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {776 if (!drafting[seq_id]) {777 continue;778 }779 780 auto * smpl = smpls[seq_id].get();781 782 common_sampler_sample(smpl, ctx_dft, i_batch, true);783 // pre-norm hidden state of this position becomes g_embd for the next step784 const float * prenorm = llama_get_embeddings_nextn_ith(ctx_dft, i_batch);785 ++i_batch;786 787 const auto * cur_p = common_sampler_get_candidates(smpl, true);788 789 for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {790 SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",791 seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,792 common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());793 }794 795 const llama_token id = cur_p->data[0].id;796 797 // only collect very high-confidence draft tokens798 // (configurable via --spec-draft-p-min, set to 0.0 to disable early-stop)799 if (cur_p->data[0].p < params.p_min) {800 drafting[seq_id] = false;801 n_drafting--;802 803 continue;804 }805 806 common_sampler_accept(smpl, id, true);807 808 auto & dp = dparams.at(seq_id);809 auto & result = *dp.result;810 811 result.push_back(id);812 813 if (params.n_max <= (int) result.size()) {814 drafting[seq_id] = false;815 n_drafting--;816 continue;817 }818 819 common_batch_add(batch, id, pending_pos_last[seq_id] + (i + 1), { seq_id }, true);820 std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec, prenorm, row_bytes);821 }822 823 if (batch.n_tokens == 0) {824 break;825 }826 827 ret = llama_decode(ctx_dft, batch);828 if (ret != 0) {829 SPC_ERR("llama_decode[%d] returned %d\n", i, ret);830 break;831 }832 833 ++i;834 }835 836 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {837 auto & dp = dparams[seq_id];838 if (!dp.drafting) {839 continue;840 }841 842 if (dp.result->size() < (size_t) params.n_min) {843 dp.result->clear();844 }845 }846 }847 848 void accept(llama_seq_id seq_id, uint16_t n_accepted, bool /*is_other*/) override {849 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {850 return;851 }852 853 const int32_t n_rows = verify_g_rows[seq_id];854 if (n_rows <= 0) {855 return;856 }857 858 const int32_t i_g = std::min<int32_t>(n_accepted, n_rows - 1);859 pending_pos_last[seq_id] = verify_pos_first[seq_id] + i_g;860 std::memcpy(pending_g_last[seq_id].data(),861 verify_g[seq_id].data() + (size_t) i_g * n_embd_dec,862 (size_t) n_embd_dec * sizeof(float));863 }864 865 // we only need to stash the deferred boundary's g_embd row for recurrent/hybrid targets:866 // their single-position checkpoints drop it on restore867 bool need_boundary_stash() const {868 const llama_model * model_tgt = llama_get_model(params.ctx_tgt);869 return llama_model_is_recurrent(model_tgt) || llama_model_is_hybrid(model_tgt);870 }871 872 bool get_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {873 if (!need_boundary_stash()) {874 return false;875 }876 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || pending_pos_last[seq_id] < 0) {877 return false;878 }879 880 const llama_pos pos = pending_pos_last[seq_id];881 const std::vector<float> & g = pending_g_last[seq_id];882 883 data.resize(sizeof(llama_pos) + g.size() * sizeof(float));884 std::memcpy(data.data(), &pos, sizeof(llama_pos));885 std::memcpy(data.data() + sizeof(llama_pos), g.data(), g.size() * sizeof(float));886 return true;887 }888 889 void set_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {890 if (!need_boundary_stash()) {891 return;892 }893 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {894 return;895 }896 if (data.size() != sizeof(llama_pos) + (size_t) n_embd_dec * sizeof(float)) {897 return;898 }899 900 llama_pos pos = -1;901 std::memcpy(&pos, data.data(), sizeof(llama_pos));902 903 pending_pos_last[seq_id] = pos;904 pending_g_last[seq_id].resize(n_embd_dec);905 std::memcpy(pending_g_last[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd_dec * sizeof(float));906 }907};908 909// DFlash: block-diffusion drafting with a draft-side KV cache injection910struct common_speculative_impl_draft_dflash : public common_speculative_impl {911 common_params_speculative_draft params;912 913 llama_batch batch; // noise tokens914 llama_batch batch_inject; // target features for KV cache injection915 916 std::vector<common_sampler_ptr> smpls;917 918 // backend sampler chain per seq, attached to ctx_dft919 std::vector<llama_sampler *> backend_chains;920 921 int32_t n_embd_dec = 0; // draft hidden size922 int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size923 int32_t n_embd_tgt = 0; // target model hidden size924 925 int32_t block_size = 0;926 llama_token mask_token_id = 0;927 928 bool is_dflash2 = false;929 bool is_mrope = false;930 int32_t selector_top_k = 0;931 932 // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout933 const bool is_dspark;934 935 // dspark speculators936 bool sample_from_anchor = true;937 938 // block-internal attention939 bool causal_attn = false;940 941 const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices942 uint32_t target_layer_ids_n = 0;943 944 common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,945 common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)946 : common_speculative_impl(type, n_seq, params.draft.n_max)947 , params(params.draft)948 , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)949 {950 auto * ctx_tgt = this->params.ctx_tgt;951 auto * ctx_dft = this->params.ctx_dft;952 GGML_ASSERT(ctx_tgt && ctx_dft && "DFlash requires ctx_tgt and ctx_dft to be set");953 954 const llama_model * model_dft = llama_get_model(ctx_dft);955 const llama_model * model_tgt = llama_get_model(ctx_tgt);956 957 target_layer_ids = llama_model_target_layer_ids (model_dft);958 target_layer_ids_n = llama_model_target_layer_ids_n(model_dft);959 GGML_ASSERT(target_layer_ids_n > 0 && "DFlash model has no target_layer_ids");960 961 n_embd_tgt = llama_model_n_embd(model_tgt);962 n_embd_dec = llama_model_n_embd(model_dft);963 n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;964 965 // read the trained block size from the dflash.block_size metadata key966 block_size = 16;967 {968 char buf[32] = {};969 if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) {970 block_size = std::atoi(buf);971 }972 if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {973 sample_from_anchor = std::strcmp(buf, "true") == 0;974 }975 if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) {976 causal_attn = std::strcmp(buf, "true") == 0;977 }978 }979 980 selector_top_k = llama_model_dflash_selector_top_k(model_dft);981 is_dflash2 = selector_top_k > 0;982 mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));983 984 if (is_dspark && this->params.p_min > 0.0f) {985 char buf[16] = {};986 const bool has_conf =987 llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||988 std::strcmp(buf, "true") == 0;989 if (!has_conf) {990 throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");991 }992 }993 994 LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());995 LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);996 LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__,997 block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false");998 999 // DFlash input is [id_last, <mask> * (block_size-1)]: in-place denoising yields at most1000 // block_size-1 draft tokens, anchor-first DSpark yields a full block_size draft tokens1001 const int32_t n_draft_max = is_dspark && sample_from_anchor ? block_size : block_size - 1;1002 if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) {1003 LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n",1004 __func__, this->params.n_max, this->params.n_min, block_size, n_draft_max);1005 this->params.n_max = std::min(this->params.n_max, n_draft_max);1006 this->params.n_min = std::min(this->params.n_min, n_draft_max);1007 }1008 this->n_max = this->params.n_max;1009 1010 batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq);1011 batch_inject = llama_batch_init(llama_n_ubatch(ctx_dft), n_embd_enc, n_seq);1012 1013 // embd batches on an M-RoPE draft need 4 position rows per token1014 is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;1015 if (is_mrope) {1016 free(batch_inject.pos);1017 batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft));1018 }1019 1020 smpls.resize(n_seq);1021 for (auto & s : smpls) {1022 common_params_sampling sparams;1023 sparams.no_perf = false;1024 sparams.top_k = 10;1025 sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };1026 s.reset(common_sampler_init(model_dft, sparams));1027 }1028 1029 // offload draft sampling to the backend1030 backend_chains.assign(n_seq, nullptr);1031 if (this->params.backend_sampling && !is_dflash2) {1032 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {1033 llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());1034 llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));1035 1036 if (!llama_set_sampler(ctx_dft, seq_id, chain)) {1037 SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id);1038 llama_sampler_free(chain);1039 chain = nullptr;1040 }1041 backend_chains[seq_id] = chain;1042 }1043 }1044 1045 // turn on extraction of the target layers' input embeddings1046 for (uint32_t k = 0; k < target_layer_ids_n; ++k) {1047 llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);1048 }1049 1050 // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.1051 llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);1052 llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise1053 }1054 1055 ~common_speculative_impl_draft_dflash() override {1056 auto * ctx_dft = this->params.ctx_dft;1057 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {1058 if (backend_chains[seq_id] == nullptr) {1059 continue;1060 }1061 if (ctx_dft) {1062 llama_set_sampler(ctx_dft, seq_id, nullptr);1063 }1064 llama_sampler_free(backend_chains[seq_id]);1065 }1066 backend_chains.clear();1067 1068 llama_batch_free(batch);1069 llama_batch_free(batch_inject);1070 }1071 1072 void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {1073 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {1074 return;1075 }1076 1077 const int32_t N = (int32_t) prompt.size();1078 if (N <= 0) {1079 return;1080 }1081 1082 const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(params.ctx_dft), seq_id);1083 if (pos_max < N - 1) {1084 LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - process() did not run on every prefill ubatch. "1085 "Drafts may degrade.\n",1086 __func__, (int) pos_max, N - 1);1087 }1088 }1089 1090 bool process(const llama_batch & batch_in) override {1091 if (batch_in.n_tokens <= 0) {1092 return true;1093 }1094 1095 // Target prefill may contain token IDs or multimodal embeddings. Both1096 // produce the target-layer features used to seed the draft KV cache, so1097 // embeddings are injected too, except the pinned ones skipped below.1098 // TODO: revisit after https://github.com/ggml-org/llama.cpp/pull/24669 is merged1099 const bool has_tokens = batch_in.token != nullptr;1100 const bool has_embeddings = batch_in.embd != nullptr;1101 if (has_tokens == has_embeddings) {1102 return true;1103 }1104 1105 const int32_t n_tokens = batch_in.n_tokens;1106 1107 // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch)1108 std::vector<int32_t> i_batch_beg(n_seq, -1);1109 std::vector<int32_t> i_batch_end(n_seq, -1);1110 for (int32_t k = 0; k < n_tokens; ++k) {1111 GGML_ASSERT(batch_in.n_seq_id[k] == 1);1112 const llama_seq_id seq_id = batch_in.seq_id[k][0];1113 if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {1114 continue;1115 }1116 i_batch_end[seq_id] = k;1117 if (i_batch_beg[seq_id] < 0) {1118 i_batch_beg[seq_id] = k;1119 }1120 }1121 1122 auto * ctx_tgt = this->params.ctx_tgt;1123 auto * ctx_dft = this->params.ctx_dft;1124 1125 const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft);1126 1127 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {1128 if (i_batch_beg[seq_id] < 0) {1129 continue;1130 }1131 const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;1132 1133 // an M-RoPE image pins all its rows to one position, so a windowed draft1134 // cache cannot free cells for it - skip it, the draft can jump over the gap1135 const bool pos_pinned = batch_in.pos[i_batch_beg[seq_id]] == batch_in.pos[i_batch_end[seq_id]];1136 if (has_embeddings && n_rows > 1 && pos_pinned) {1137 continue;1138 }1139 1140 for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {1141 const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);1142 1143 // gather target features per extract layer; the fused decode encodes and1144 // injects them into the K/V cache at the target positions1145 batch_inject.n_tokens = n_chunk;1146 for (uint32_t k = 0; k < target_layer_ids_n; ++k) {1147 const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);1148 if (!layer) {1149 GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);1150 }1151 for (int32_t i = 0; i < n_chunk; ++i) {1152 float * dst = batch_inject.embd + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;1153 const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;1154 std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));1155 }1156 }1157 1158 for (int32_t i = 0; i < n_chunk; ++i) {1159 const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];1160 batch_inject.pos[i] = p;1161 if (is_mrope) {1162 batch_inject.pos[1 * n_chunk + i] = p;1163 batch_inject.pos[2 * n_chunk + i] = p;1164 batch_inject.pos[3 * n_chunk + i] = 0;1165 }1166 batch_inject.n_seq_id[i] = 1;1167 batch_inject.seq_id[i][0] = seq_id;1168 batch_inject.logits[i] = false;1169 }1170 const int32_t rc = llama_decode(ctx_dft, batch_inject);1171 if (rc != 0) {1172 LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",1173 __func__, rc, (int) n_chunk, (int) offset);1174 return false;1175 }1176 }1177 }1178 1179 return true;1180 }1181 1182 void draft(common_speculative_draft_params_vec & dparams) override {1183 auto & ctx_dft = params.ctx_dft;1184 1185 common_batch_clear(batch);1186 1187 // build one batch holding every drafting sequence's noise block into a single decode)1188 // record where each block starts and its size1189 std::vector<int32_t> i_block_beg(n_seq, -1);1190 std::vector<int32_t> n_block (n_seq, 0);1191 1192 for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {1193 auto & dp = dparams[seq_id];1194 if (!dp.drafting) {1195 continue;1196 }1197 1198 common_sampler_reset(smpls[seq_id].get());1199 1200 const int32_t n = (int32_t) dp.pos0;