diff --git a/common/common.cpp b/common/common.cpp index 341f8e8bb9c..7a4e12e30f7 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1503,7 +1503,8 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode } if (llama_model_has_encoder(model)) { - llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size())); + common_batch batch = common_batch_get_one(lctx, tmp); + llama_process(lctx, LLAMA_PROCESS_TYPE_ENCODE, batch.get()); llama_token decoder_start_token_id = llama_model_decoder_start_token(model); if (decoder_start_token_id == LLAMA_TOKEN_NULL) { decoder_start_token_id = bos; @@ -1512,7 +1513,9 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode tmp.push_back(decoder_start_token_id); } if (llama_model_has_decoder(model)) { - llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch))); + tmp.resize(std::min(tmp.size(), (size_t) params.n_batch)); + common_batch batch = common_batch_get_one(lctx, tmp); + llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); } llama_memory_clear(llama_get_memory(lctx), true); llama_synchronize(lctx); @@ -1571,9 +1574,13 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { tmp.push_back(0); tmp.push_back(0); - int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size())); + int ret; + { + common_batch batch = common_batch_get_one(ctx, tmp); + ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); + } if (ret != 0) { - COM_ERR("llama_decode() failed: %d\n", ret); + COM_ERR("llama_process() failed: %d\n", ret); res = COMMON_CONTEXT_SEQ_RM_TYPE_NO; goto done; } @@ -2163,31 +2170,58 @@ float lr_opt::get_lr(float epoch) const { } bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos) { - llama_batch batch = llama_batch_get_one(&last_token, 1); - batch.pos = &pos; - if (llama_decode(ctx, batch)) { + common_batch batch(ctx); + batch.add(last_token, pos, 0, true); + + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) { LOG_ERR("%s: failed to replay last token\n", __func__); return false; } return true; } -llama_batch_ext_ptr common_batch_ext_get_one(llama_context * ctx, const llama_tokens & tokens) { - llama_batch_ext_ptr batch(llama_batch_ext_init(ctx)); +void common_batch::clear() { + tokens.clear(); + llama_batch_ext_clear(batch.get()); +} + +int32_t common_batch::add(llama_token id, llama_pos pos, llama_seq_id seq_id, bool output) { + const int32_t idx = llama_batch_ext_add_token(batch.get(), seq_id, id); + if (idx < 0) { + return idx; + } + llama_batch_ext_set_pos(batch.get(), idx, &pos); + if (output) { + llama_batch_ext_set_output_logits(batch.get(), idx, true); + } + tokens.push_back({ id, pos, seq_id, output }); + return idx; +} + +bool common_batch::set_output(int32_t idx, bool value) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + tokens[idx].output = value; + return llama_batch_ext_set_output_logits(batch.get(), idx, value); +} + +bool common_batch::set_embd(int32_t idx, llama_embd embd) { + return llama_batch_ext_set_embd_token(batch.get(), idx, embd); +} + +common_batch common_batch_get_one(llama_context * ctx, const llama_tokens & tokens) { + common_batch batch(ctx); auto mem = llama_get_memory(ctx); - llama_pos pos = mem ? llama_memory_seq_pos_max(mem, 0) + 1 : 0; + llama_pos pos = llama_memory_seq_pos_max(mem, 0) + 1; // -1 + 1 == 0 when the memory is empty for (size_t i = 0; i < tokens.size(); ++i) { - const int32_t idx = llama_batch_ext_add_token(batch.get(), 0, tokens[i]); - llama_batch_ext_set_pos(batch.get(), idx, &pos); + const bool output = i == tokens.size() - 1; + batch.add(tokens[i], pos, 0, output); pos++; } - if (!tokens.empty()) { - llama_batch_ext_set_output_logits(batch.get(), (int32_t) tokens.size() - 1, true); - } - return batch; } @@ -2215,7 +2249,7 @@ bool common_prompt_batch_decode( // memory, so we can't just remove the last token from the memory and replay the last token which // is the reason for this logic. llama_tokens prefix_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_tokens_before_last); - llama_batch_ext_ptr batch_prefix = common_batch_ext_get_one(ctx, prefix_tokens); + common_batch batch_prefix = common_batch_get_one(ctx, prefix_tokens); if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_prefix.get())) { COM_ERR("%s", "failed to eval\n"); return false; @@ -2225,10 +2259,8 @@ bool common_prompt_batch_decode( llama_state_save_file(ctx, state_path.data(), all_tokens.data(), all_tokens.size()); COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size()); - llama_token last_token = all_tokens.back(); - llama_batch_ext_ptr batch_last = common_batch_ext_get_one(ctx, { last_token }); - llama_pos pos = n_past; - llama_batch_ext_set_pos(batch_last.get(), 0, &pos); + common_batch batch_last(ctx); + batch_last.add(all_tokens.back(), n_past, 0, true); if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_last.get())) { COM_ERR("%s", "failed to eval last token\n"); @@ -2237,7 +2269,7 @@ bool common_prompt_batch_decode( n_past++; } else { llama_tokens new_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_new); - llama_batch_ext_ptr batch = common_batch_ext_get_one(ctx, new_tokens); + common_batch batch = common_batch_get_one(ctx, new_tokens); if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) { COM_ERR("%s", "failed to eval\n"); return false; diff --git a/common/common.h b/common/common.h index 00f99933f12..513970c8a7a 100644 --- a/common/common.h +++ b/common/common.h @@ -1003,9 +1003,39 @@ void common_batch_add( const std::vector & seq_ids, bool logits); +// wrapper around llama_batch_ext that provide getter functions for downstream code +struct common_batch { + struct token { + llama_token id; + llama_pos pos; + llama_seq_id seq_id; + bool output; + }; + + std::vector tokens; // mirror of the entries, tokens[i] describes batch index i + llama_batch_ext_ptr batch; + + common_batch() = default; + common_batch(struct llama_context * ctx) : batch(llama_batch_ext_init(ctx)) {} + + llama_batch_ext * get() const { return batch.get(); } + + void clear(); + + // returns the batch index (>= 0), or a negative error from llama_batch_ext_add_token() + int32_t add(llama_token id, llama_pos pos, llama_seq_id seq_id, bool output); + + bool set_output(int32_t idx, bool value); + + // attach a token embedding to the entry at idx, can only be set once per entry + bool set_embd(int32_t idx, llama_embd embd); + + int32_t size() const { return (int32_t) tokens.size(); } +}; + // create a single-sequence batch from a list of tokens // last token always have output_logits set to true -llama_batch_ext_ptr common_batch_ext_get_one(struct llama_context * ctx, const llama_tokens & tokens); +common_batch common_batch_get_one(struct llama_context * ctx, const llama_tokens & tokens); // decodes a single batch of tokens for a prompt and manages session tokens // diff --git a/common/speculative.cpp b/common/speculative.cpp index cd2dfc760b8..f873d1a6e9d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -177,7 +177,7 @@ struct common_speculative_impl { struct common_speculative_impl_draft_simple : public common_speculative_impl { common_params_speculative_draft params; - llama_batch batch; + common_batch batch; std::vector smpls; @@ -202,7 +202,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { 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); + batch = common_batch(ctx_dft); // TODO: optimize or pass from outside? // { @@ -249,10 +249,6 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { } } - ~common_speculative_impl_draft_simple() override { - llama_batch_free(batch); - } - void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override { // noop } @@ -277,7 +273,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { void draft(common_speculative_draft_params_vec & dparams) override { auto & ctx_dft = params.ctx_dft; - common_batch_clear(batch); + batch.clear(); // keep track of which sequences are still drafting int n_drafting = 0; @@ -294,12 +290,12 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { drafting[seq_id] = true; common_sampler_reset(smpls[seq_id].get()); - common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); + batch.add(dp.id_last, dp.n_past, seq_id, true); } - int ret = llama_decode(ctx_dft, batch); + int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { - SPC_ERR("llama_decode returned %d\n", ret); + SPC_ERR("llama_process returned %d\n", ret); return; } @@ -308,7 +304,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { while (n_drafting > 0) { int i_batch = 0; - common_batch_clear(batch); + batch.clear(); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { if (!drafting[seq_id]) { @@ -353,17 +349,17 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { continue; } - common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); + batch.add(id, dp.n_past + i + 1, seq_id, true); } - if (batch.n_tokens == 0) { + if (batch.size() == 0) { break; } // evaluate the drafted tokens on the draft model - ret = llama_decode(ctx_dft, batch); + ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { - SPC_ERR("llama_decode[%d] returned %d\n", i, ret); + SPC_ERR("llama_process[%d] returned %d\n", i, ret); break; } @@ -543,12 +539,6 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { 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 { @@ -908,7 +898,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_params_speculative_draft params; - llama_batch batch; // noise tokens + common_batch batch; // noise tokens llama_batch batch_inject; // target features for KV cache injection std::vector smpls; @@ -977,7 +967,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { this->params.n_min = std::min(this->params.n_min, n_draft_max); } - batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); + batch = common_batch(ctx_dft); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); smpls.resize(n_seq); @@ -1027,7 +1017,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } backend_chains.clear(); - llama_batch_free(batch); llama_batch_free(batch_inject); } @@ -1156,7 +1145,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { void draft(common_speculative_draft_params_vec & dparams) override { auto & ctx_dft = params.ctx_dft; - common_batch_clear(batch); + batch.clear(); // build one batch holding every drafting sequence's noise block into a single decode) // record where each block starts and its size @@ -1176,21 +1165,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n_draft = params.n_max; const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1); - i_block_beg[seq_id] = batch.n_tokens; + i_block_beg[seq_id] = batch.size(); n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { - common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, true); } } - if (batch.n_tokens == 0) { + if (batch.size() == 0) { return; } // decode all sequence's noise block in a single batch - int ret = llama_decode(ctx_dft, batch); + int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { - LOG_WRN("%s: llama_decode returned %d\n", __func__, ret); + LOG_WRN("%s: llama_process returned %d\n", __func__, ret); return; } @@ -1274,7 +1263,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) - llama_batch batch; + common_batch batch; std::vector smpls; @@ -1330,11 +1319,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { 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); + batch = common_batch(ctx_dft); smpls.resize(n_seq); for (auto & s : smpls) { @@ -1398,12 +1383,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { 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 { @@ -1460,33 +1439,26 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // 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); - } + batch.clear(); - // shift the tgt embeddings to the right by one position + // pair each token with the tgt embedding shifted right by one position, and + // the first token of each sequence with the pending embedding from a previous run // 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 // TODO:this is generally true, but would be nice to assert it - { - const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt); - std::memcpy(batch.embd + (size_t) 1 * n_embd, h_tgt, row_bytes * (n_tokens-1)); - } + const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt); - // fill the pending embeddings from a previous run - auto set_h = [&](int idx, const float * h_row) { - std::memcpy(batch.embd + (size_t) idx * n_embd, h_row, row_bytes); - }; + for (int k = 0; k < n_tokens; ++k) { + const llama_seq_id seq_id = batch_in.seq_id[k][0]; - for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { - if (i_batch_beg[seq_id] < 0) { - continue; - } + const int32_t idx = batch.add(batch_in.token[k], batch_in.pos[k], seq_id, false); + + const float * h_row = k == i_batch_beg[seq_id] + ? pending_h[seq_id].data() + : h_tgt + (size_t) (k - 1) * n_embd; - set_h(i_batch_beg[seq_id], pending_h[seq_id].data()); + batch.set_embd(idx, { h_row, 1, (size_t) n_embd }); } auto * mem_dft = llama_get_memory(ctx_dft); @@ -1504,9 +1476,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_set_nextn_layer_offset(ctx_dft, head); } - const int32_t rc = llama_decode(ctx_dft, batch); + const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (rc != 0) { - SPC_ERR("llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n", + SPC_ERR("llama_process(ctx_dft) head=%d failed rc=%d (pos=%d)\n", head, (int) rc, (int) batch_in.pos[0]); ok = false; break; @@ -1545,14 +1517,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { void draft(common_speculative_draft_params_vec & dparams) override { auto & ctx_dft = params.ctx_dft; - common_batch_clear(batch); + batch.clear(); // keep track of which sequences are still drafting int n_drafting = 0; std::vector drafting(n_seq); - const size_t row_bytes = (size_t) n_embd * sizeof(float); - for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { auto & dp = dparams[seq_id]; @@ -1564,10 +1534,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { drafting[seq_id] = true; common_sampler_reset(smpls[seq_id].get()); - common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); - std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[seq_id].data(), row_bytes); + const int32_t idx = batch.add(dp.id_last, dp.n_past, seq_id, true); + batch.set_embd(idx, { pending_h[seq_id].data(), 1, (size_t) n_embd }); - i_last[seq_id] = batch.n_tokens - 1; + i_last[seq_id] = idx; if (chain_heads) { chain_h[seq_id].assign(pending_h[seq_id].begin(), pending_h[seq_id].end()); @@ -1593,16 +1563,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_set_nextn_layer_offset(ctx_dft, i); } - int ret = llama_decode(ctx_dft, batch); + int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { - SPC_ERR("llama_decode[%d] returned %d\n", i, ret); + SPC_ERR("llama_process[%d] returned %d\n", i, ret); break; } // rebuild the batch for the next step: the growing-KV paths re-add only the // new token (the KV already holds the prefix), while chained heads re-add the // whole prefix at the next head. dropped sequences are simply not re-added. - common_batch_clear(batch); + batch.clear(); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { if (!drafting[seq_id]) { @@ -1653,24 +1623,24 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const int n_rows = (int) result.size() + 1; // id_last + tokens drafted so far for (int t = 0; t < n_rows; ++t) { const llama_token tok = (t == 0) ? dp.id_last : result[t - 1]; - common_batch_add(batch, tok, dp.n_past + t, { seq_id }, t == n_rows - 1); - std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, - chain_h[seq_id].data() + (size_t) t * n_embd, row_bytes); + const int32_t idx = batch.add(tok, dp.n_past + t, seq_id, t == n_rows - 1); + batch.set_embd(idx, { chain_h[seq_id].data() + (size_t) t * n_embd, 1, (size_t) n_embd }); + i_last[seq_id] = idx; } } else if (is_mem_shared) { // note: with shared memory (e.g. Gemma4 assistants) we use the same position for all draft tokens // ref: https://github.com/huggingface/transformers/blob/effde20942e3f82a1b97449f60b3a48c5ff96145/docs/source/en/model_doc/gemma4_assistant.md?plain=1#L36-L37 - common_batch_add(batch, id, dp.n_past, { seq_id }, true); - std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); + const int32_t idx = batch.add(id, dp.n_past, seq_id, true); + batch.set_embd(idx, { h_row, 1, (size_t) n_embd }); + i_last[seq_id] = idx; } else { - common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); - std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); + const int32_t idx = batch.add(id, dp.n_past + i + 1, seq_id, true); + batch.set_embd(idx, { h_row, 1, (size_t) n_embd }); + i_last[seq_id] = idx; } - - i_last[seq_id] = batch.n_tokens - 1; } - if (batch.n_tokens == 0) { + if (batch.size() == 0) { break; }