diff --git a/common/common.cpp b/common/common.cpp index 05b14cc72..10712d389 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1745,33 +1745,6 @@ void common_threadpools::init(llama_context * ctx, const common_params & params) llama_attach_threadpool(ctx, threadpool, threadpool_batch); } -// -// Batch utils -// - -void common_batch_clear(struct llama_batch & batch) { - batch.n_tokens = 0; -} - -void common_batch_add( - struct llama_batch & batch, - llama_token id, - llama_pos pos, - const std::vector & seq_ids, - bool logits) { - GGML_ASSERT(batch.seq_id[batch.n_tokens] && "llama_batch size exceeded"); - - batch.token [batch.n_tokens] = id; - batch.pos [batch.n_tokens] = pos; - batch.n_seq_id[batch.n_tokens] = seq_ids.size(); - for (size_t i = 0; i < seq_ids.size(); ++i) { - batch.seq_id[batch.n_tokens][i] = seq_ids[i]; - } - batch.logits [batch.n_tokens] = logits; - - batch.n_tokens++; -} - // // Vocab utils // @@ -2125,35 +2098,41 @@ common_batch::common_batch(llama_context * ctx) : 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) { - GGML_ABORT("%s: failed to add token %d to the batch (error %d, n_tokens = %d)\n", __func__, id, idx, size()); + tokens.push_back({ id, { pos, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 }, {} }); + return size() - 1; +} + +int32_t common_batch::add(llama_token id, llama_pos pos, const std::vector & seq_ids, bool output) { + GGML_ASSERT(!seq_ids.empty()); + + const int32_t idx = add(id, pos, seq_ids[0], output); + for (size_t s = 1; s < seq_ids.size(); ++s) { + add_seq(idx, seq_ids[s]); } - 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, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 } }); return idx; } +bool common_batch::add_seq(int32_t idx, llama_seq_id seq_id) { + if (idx < 0 || idx >= size()) { + return false; + } + tokens[idx].seq_ids_extra.push_back(seq_id); + return true; +} + bool common_batch::set_output(int32_t idx, bool value) { - if (idx < 0 || idx >= (int32_t) tokens.size()) { + if (idx < 0 || idx >= size()) { return false; } tokens[idx].output = value; - return llama_batch_ext_set_output_logits(batch.get(), idx, value); + return true; } bool common_batch::set_embd(int32_t idx, llama_embd embd) { - if (idx < 0 || idx >= (int32_t) tokens.size()) { - return false; - } - if (!llama_batch_ext_set_embd_token(batch.get(), idx, embd)) { + if (idx < 0 || idx >= size() || tokens[idx].embd.data != nullptr) { return false; } tokens[idx].embd = embd; @@ -2161,83 +2140,64 @@ bool common_batch::set_embd(int32_t idx, llama_embd embd) { } int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output) { - const int32_t idx = llama_batch_ext_add_embd(batch.get(), seq_id, embd); - if (idx < 0) { - GGML_ABORT("%s: failed to add embedding to the batch (error %d, n_tokens = %d)\n", __func__, idx, size()); - } - llama_batch_ext_set_pos(batch.get(), idx, pos); - if (output) { - llama_batch_ext_set_output_logits(batch.get(), idx, true); - } - token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, embd }; + token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, embd, {} }; for (int32_t j = 0; j < n_pos; ++j) { t.pos[j] = pos[j]; } tokens.push_back(t); - return idx; + return size() - 1; } -common_batch common_batch_from_llama_batch(llama_context * ctx, const llama_batch & batch) { - common_batch res(ctx); +llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) { + GGML_ASSERT(batch && "common_batch was not initialized with a context"); + GGML_ASSERT(off >= 0 && n >= 0 && off + n <= size()); - const bool has_token = batch.token != nullptr; - const bool has_embd = batch.embd != nullptr; + llama_batch_ext * res = batch.get(); + llama_batch_ext_clear(res); - const size_t n_embd = llama_model_n_embd_inp(llama_get_model(ctx)); - - // positions continue from the memory when none are given - auto * mem = llama_get_memory(ctx); - std::vector pos_next(llama_n_seq_max(ctx)); - for (llama_seq_id s = 0; s < (llama_seq_id) pos_next.size(); ++s) { - pos_next[s] = llama_memory_seq_pos_max(mem, s) + 1; - } - - for (int32_t i = 0; i < batch.n_tokens; ++i) { - const int32_t n_sid = batch.n_seq_id ? batch.n_seq_id[i] : 1; - const llama_seq_id seq_id = batch.seq_id ? batch.seq_id[i][0] : 0; - - llama_pos pos[GGML_MROPE_SECTIONS] = { 0, 0, 0, 0 }; - if (!batch.pos) { - pos[0] = pos_next[seq_id]++; - } else if (has_token) { - pos[0] = batch.pos[i]; - } else { - // embedding batch: section-major layout pos[j*n_tokens + i] - for (int32_t j = 0; j < res.n_pos; ++j) { - pos[j] = batch.pos[j * batch.n_tokens + i]; - } - } - - const bool output = batch.logits ? batch.logits[i] != 0 : i == batch.n_tokens - 1; - - const llama_embd embd = { has_embd ? batch.embd + (size_t) i * n_embd : nullptr, 1, n_embd }; + for (int32_t i = off; i < off + n; ++i) { + const token & t = tokens[i]; int32_t idx; - if (has_token) { - idx = res.add(batch.token[i], pos[0], seq_id, output); - if (has_embd) { - res.set_embd(idx, embd); + if (t.id != LLAMA_TOKEN_NULL) { + idx = llama_batch_ext_add_token(res, t.seq_id, t.id); + if (idx < 0) { + GGML_ABORT("%s: failed to add token %d at index %d (error %d, n = %d)\n", __func__, t.id, i, idx, n); + } + llama_batch_ext_set_pos(res, idx, t.pos.data()); + if (t.embd.data && !llama_batch_ext_set_embd_token(res, idx, t.embd)) { + GGML_ABORT("%s: failed to set the embedding of token %d at index %d\n", __func__, t.id, i); } } else { - idx = res.add_embd(embd, pos, seq_id, output); + idx = llama_batch_ext_add_embd(res, t.seq_id, t.embd); + if (idx < 0) { + GGML_ABORT("%s: failed to add embedding at index %d (error %d, n = %d)\n", __func__, i, idx, n); + } + llama_batch_ext_set_pos(res, idx, t.pos.data()); } + GGML_ASSERT(idx == i - off); - for (int32_t s = 1; s < n_sid; ++s) { - llama_batch_ext_add_seq(res.get(), idx, batch.seq_id[i][s]); + for (const llama_seq_id seq_id : t.seq_ids_extra) { + if (!llama_batch_ext_add_seq(res, idx, seq_id)) { + GGML_ABORT("%s: failed to add seq %d to the entry at index %d\n", __func__, seq_id, i); + } + } + if (t.output) { + llama_batch_ext_set_output_logits(res, idx, true); } } return res; } -common_batch common_batch_get_one(llama_context * ctx, const llama_tokens & tokens) { +common_batch common_batch_get_one(llama_context * ctx, const llama_token * tokens, int32_t n_tokens) { common_batch batch(ctx); auto mem = llama_get_memory(ctx); 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 bool output = i == tokens.size() - 1; + for (int32_t i = 0; i < n_tokens; ++i) { + const bool output = i == n_tokens - 1; batch.add(tokens[i], pos, 0, output); pos++; } @@ -2245,6 +2205,10 @@ common_batch common_batch_get_one(llama_context * ctx, const llama_tokens & toke return batch; } +common_batch common_batch_get_one(llama_context * ctx, const llama_tokens & tokens) { + return common_batch_get_one(ctx, tokens.data(), (int32_t) tokens.size()); +} + bool common_prompt_batch_decode( struct llama_context * ctx, const llama_tokens & all_tokens, diff --git a/common/common.h b/common/common.h index f5ae1a15c..fd222ba03 100644 --- a/common/common.h +++ b/common/common.h @@ -1032,23 +1032,16 @@ struct common_memory { // Batch utils // -void common_batch_clear(struct llama_batch & batch); - -void common_batch_add( - struct llama_batch & batch, - llama_token id, - llama_pos pos, - const std::vector & seq_ids, - bool logits); - // wrapper around llama_batch_ext that provide getter functions for downstream code +// entries can exceed n_batch, use get_sub_batch() to decode them in chunks struct common_batch { struct token { llama_token id; std::array pos; // only pos[0] is used for text tokens - llama_seq_id seq_id; + llama_seq_id seq_id; // the first sequence id, see add_seq() bool output; llama_embd embd; // non-owning view of the data passed to add_embd()/set_embd(), data == NULL if none + std::vector seq_ids_extra; // see add_seq() }; std::vector tokens; // mirror of the entries, tokens[i] describes batch index i @@ -1059,7 +1052,10 @@ struct common_batch { common_batch() = default; common_batch(struct llama_context * ctx); - llama_batch_ext * get() const { return batch.get(); } + llama_batch_ext * get() { return get_sub_batch(0, size()); } + + // render entries [off, off + n) into batch, the result is overwritten by the next call + llama_batch_ext * get_sub_batch(int32_t off, int32_t n); // content type of the batch, all entries carry the same combination bool has_token() const { return !tokens.empty() && tokens[0].id != LLAMA_TOKEN_NULL; } @@ -1067,15 +1063,21 @@ struct common_batch { void clear(); - // returns the batch index (>= 0), aborts if the entry cannot be added (batch full, invalid token or seq id) + // returns the batch index int32_t add(llama_token id, llama_pos pos, llama_seq_id seq_id, bool output); + // same, with the entry shared by all seq_ids (must not be empty) + int32_t add(llama_token id, llama_pos pos, const std::vector & seq_ids, bool output); + + // add the entry at idx to another sequence, tokens[idx].seq_id keeps the first one + bool add_seq(int32_t idx, llama_seq_id seq_id); + 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); - // add an embedding-only entry (no token id), aborts like add() on failure + // add an embedding-only entry (no token id) // pos points to n_pos positions int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output); @@ -1083,13 +1085,10 @@ struct common_batch { }; // create a single-sequence batch from a list of tokens -// last token always have output_logits set to true +// positions continue from the memory, last token always have output_logits set to true +common_batch common_batch_get_one(struct llama_context * ctx, const llama_token * tokens, int32_t n_tokens); common_batch common_batch_get_one(struct llama_context * ctx, const llama_tokens & tokens); -// convert a legacy llama_batch, applying its defaults: seq 0, positions continue from memory, last token is output -// the embd rows are read at the model input width -common_batch common_batch_from_llama_batch(struct llama_context * ctx, const llama_batch & batch); - // decodes a single batch of tokens for a prompt and manages session tokens // // Note: We save state before the last token so that we can replay it to ensure diff --git a/common/speculative.cpp b/common/speculative.cpp index 3180c3c85..a648197f5 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2163,9 +2163,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { struct common_speculative { common_speculative_draft_params_vec dparams; - // the target context, used to convert legacy llama_batch inputs - llama_context * ctx_tgt = nullptr; - // list of implementations to use and their states std::vector> impls; @@ -2711,7 +2708,6 @@ common_speculative * common_speculative_init(common_params_speculative & params, common_speculative_ptr result(new common_speculative { /* .dparams = */ common_speculative_draft_params_vec(n_seq), - /* .ctx_tgt = */ params.draft.ctx_tgt, /* .impls = */ std::move(impls), /* .impl_last = */ std::vector(n_seq, nullptr), /* .synth_probs = */ {}, @@ -2774,17 +2770,6 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co } } -bool common_speculative_process(common_speculative * spec, const llama_batch & batch) { - if (spec == nullptr) { - return true; - } - - // ngram-only setups have no target context, they do not read the batch anyway - const common_batch tmp = spec->ctx_tgt ? common_batch_from_llama_batch(spec->ctx_tgt, batch) : common_batch(); - - return common_speculative_process(spec, tmp); -} - bool common_speculative_process(common_speculative * spec, const common_batch & batch) { bool result = true; diff --git a/common/speculative.h b/common/speculative.h index 211fcdabd..d46b21eb7 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -79,9 +79,6 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co // process the batch and update the internal state of the speculative context bool common_speculative_process(common_speculative * spec, const common_batch & batch); -// legacy llama_batch input, converted with common_batch_from_llama_batch() -bool common_speculative_process(common_speculative * spec, const llama_batch & batch); - // generate drafts for the sequences specified with `common_speculative_get_draft_params` void common_speculative_draft(common_speculative * spec); diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 97d6b6944..0ff57d953 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -1,5 +1,7 @@ #include "diffusion.h" +#include "common.h" + #include "log.h" #include @@ -144,8 +146,7 @@ void diffusion_generate(llama_context * ctx, struct llama_sampler * dist_sampler = llama_sampler_init_dist(params.seed); - llama_batch batch = llama_batch_init(params.max_length, 0, 1); - batch.n_tokens = params.max_length; + common_batch batch(ctx); // Pre-allocate buffers for CFG if needed int32_t logits_size = n_vocab * params.max_length; @@ -202,18 +203,15 @@ void diffusion_generate(llama_context * ctx, } // Setup batch + batch.clear(); for (int32_t i = 0; i < params.max_length; i++) { - batch.token[i] = output_tokens[i]; - batch.pos[i] = i; - batch.n_seq_id[i] = 1; - batch.seq_id[i][0] = 0; - batch.logits[i] = 1; + batch.add(output_tokens[i], i, 0, true); } float * logits = nullptr; if (params.cfg_scale > 0.0f) { - int ret = llama_decode(ctx, batch); + int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { LOG_ERR("Failed to generate conditional"); break; @@ -227,10 +225,11 @@ void diffusion_generate(llama_context * ctx, un_x_buffer[i] = params.mask_token_id; } + batch.clear(); for (int32_t i = 0; i < params.max_length; i++) { - batch.token[i] = un_x_buffer[i]; + batch.add(un_x_buffer[i], i, 0, true); } - ret = llama_decode(ctx, batch); + ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { LOG_ERR("Failed to generate unconditional"); break; @@ -244,7 +243,7 @@ void diffusion_generate(llama_context * ctx, } logits = cond_logits_buffer.data(); } else { - int ret = llama_decode(ctx, batch); + int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); if (ret != 0) { LOG_ERR("%s: failed to decode at step %d, ret = %d\n", __func__, global_step, ret); break; @@ -400,7 +399,6 @@ void diffusion_generate(llama_context * ctx, total_time / 1000.0 / params.steps, total_sampling_time / 1000.0 / params.steps); - llama_batch_free(batch); llama_sampler_free(sampler); llama_sampler_free(dist_sampler); diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp index f6a20ef9d..a59f04cae 100644 --- a/examples/embedding/embedding.cpp +++ b/examples/embedding/embedding.cpp @@ -27,27 +27,27 @@ static std::vector split_lines(const std::string & s, const std::st return lines; } -static void batch_add_seq(llama_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { +static void batch_add_seq(common_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { size_t n_tokens = tokens.size(); for (size_t i = 0; i < n_tokens; i++) { - common_batch_add(batch, tokens[i], i, { seq_id }, true); + batch.add(tokens[i], i, seq_id, true); } } -static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) { +static void batch_decode(llama_context * ctx, common_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) { const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); // clear previous kv_cache values (irrelevant for embeddings) llama_memory_clear(llama_get_memory(ctx), true); // run model - LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq); - if (llama_decode(ctx, batch) < 0) { + LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.size(), n_seq); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) < 0) { LOG_ERR("%s : failed to process\n", __func__); } - for (int i = 0; i < batch.n_tokens; i++) { - if (!batch.logits[i]) { + for (int i = 0; i < batch.size(); i++) { + if (!batch.tokens[i].output) { continue; } @@ -61,8 +61,8 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu GGML_ASSERT(embd != NULL && "failed to get token embeddings"); } else { // try to get sequence embeddings - supported only when pooling_type is not NONE - embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]); - embd_pos = batch.seq_id[i][0]; + embd = llama_get_embeddings_seq(ctx, batch.tokens[i].seq_id); + embd_pos = batch.tokens[i].seq_id; GGML_ASSERT(embd != NULL && "failed to get sequence embeddings"); } @@ -242,7 +242,7 @@ int main(int argc, char ** argv) { // initialize batch const int n_prompts = prompts.size(); - struct llama_batch batch = llama_batch_init(n_batch, 0, 1); + common_batch batch(ctx); // count number of embeddings int n_embd_count = 0; @@ -269,12 +269,12 @@ int main(int argc, char ** argv) { const uint64_t n_toks = inp.size(); // encode if at capacity - if (batch.n_tokens + n_toks > n_batch || s >= n_seq_max) { + if (batch.size() + n_toks > n_batch || s >= n_seq_max) { float * out = emb + e * n_embd_out; batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize); - e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s; + e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.size() : s; s = 0; - common_batch_clear(batch); + batch.clear(); } // add to batch @@ -407,7 +407,6 @@ int main(int argc, char ** argv) { llama_perf_context_print(ctx); // clean up - llama_batch_free(batch); llama_backend_free(); return 0; diff --git a/kcpp_src/gpttype_adapter.cpp b/kcpp_src/gpttype_adapter.cpp index ba6724ecf..926180527 100644 --- a/kcpp_src/gpttype_adapter.cpp +++ b/kcpp_src/gpttype_adapter.cpp @@ -1082,7 +1082,18 @@ static int32_t kcpp_decode_main_and_spec(llama_context * main_ctx, llama_batch b { llama_memory_seq_rm(llama_get_memory(draft_ctx), batch.seq_id[0][0], batch.pos[0], -1); } - if(!common_speculative_process(draft_spec, batch)) + // Callers provide text batches with explicit positions, sequence IDs and logits flags. + GGML_ASSERT(batch.token && !batch.embd); + common_batch spec_batch; + for(int32_t i = 0; i < batch.n_tokens; ++i) + { + const int32_t idx = spec_batch.add(batch.token[i], batch.pos[i], batch.seq_id[i][0], batch.logits[i] != 0); + for(int32_t s = 1; s < batch.n_seq_id[i]; ++s) + { + spec_batch.add_seq(idx, batch.seq_id[i][s]); + } + } + if(!common_speculative_process(draft_spec, spec_batch)) { kcpp_flush_log_output(); printf("\nERROR: Speculative state update failed!\n"); @@ -4810,7 +4821,7 @@ static bool batch_claim_waiting_locked() static void batch_worker_loop() { const int batch_cap = std::max(1, kcpp_data ? kcpp_data->n_batch : 512); - llama_batch batch = llama_batch_init(batch_cap, 0, 1); + common_batch batch(llama_ctx_v4); while(true) { std::vector decode_ids; @@ -4828,10 +4839,10 @@ static void batch_worker_loop() continue; } batch_claim_waiting_locked(); - common_batch_clear(batch); + batch.clear(); for(auto & req_ptr : batch_requests) { - if(!req_ptr || !batch_is_live_state(req_ptr->state) || req_ptr->slot < 0 || batch.n_tokens >= batch_cap) + if(!req_ptr || !batch_is_live_state(req_ptr->state) || req_ptr->slot < 0 || batch.size() >= batch_cap) { continue; } @@ -4845,15 +4856,15 @@ static void batch_worker_loop() } if(req.state == BatchState::PREFILL) { - while(req.prompt_pos < (int) req.prompt_tokens.size() && batch.n_tokens < batch_cap) + while(req.prompt_pos < (int) req.prompt_tokens.size() && batch.size() < batch_cap) { bool is_last = req.prompt_pos == (int) req.prompt_tokens.size() - 1; if(is_last) { - req.i_batch = batch.n_tokens; + req.i_batch = batch.size(); req.i_batch_is_prefill = true; } - common_batch_add(batch, req.prompt_tokens[req.prompt_pos], req.n_past, { req.slot }, is_last); + batch.add(req.prompt_tokens[req.prompt_pos], req.n_past, req.slot, is_last); req.prompt_pos++; req.n_past++; } @@ -4864,14 +4875,14 @@ static void batch_worker_loop() } else if(req.state == BatchState::GENERATING && req.has_pending) { - req.i_batch = batch.n_tokens; + req.i_batch = batch.size(); req.i_batch_is_prefill = false; - common_batch_add(batch, req.pending_token, req.n_past, { req.slot }, true); + batch.add(req.pending_token, req.n_past, req.slot, true); req.n_past++; req.has_pending = false; } } - if(batch.n_tokens == 0) + if(batch.size() == 0) { continue; } @@ -4884,7 +4895,7 @@ static void batch_worker_loop() } } - int decode_status = llama_decode(llama_ctx_v4, batch); + int decode_status = llama_process(llama_ctx_v4, LLAMA_PROCESS_TYPE_DECODE, batch.get()); auto decode_finish_time = std::chrono::steady_clock::now(); std::lock_guard lock(batch_mutex); @@ -4941,7 +4952,6 @@ static void batch_worker_loop() req->i_batch = -1; } } - llama_batch_free(batch); } static void batch_start_worker_locked() diff --git a/otherarch/embeddings_adapter.cpp b/otherarch/embeddings_adapter.cpp index f9d322c8a..0261d0ce2 100644 --- a/otherarch/embeddings_adapter.cpp +++ b/otherarch/embeddings_adapter.cpp @@ -26,14 +26,14 @@ bool embeddings_debug = false; static int max_batchsize = 512; static std::string last_output = ""; -static void batch_add_seq(llama_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { +static void batch_add_seq(common_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { size_t n_tokens = tokens.size(); for (size_t i = 0; i < n_tokens; i++) { - common_batch_add(batch, tokens[i], i, { seq_id }, true); + batch.add(tokens[i], i, seq_id, true); } } -static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) { +static void batch_decode(llama_context * ctx, common_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) { const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); const struct llama_model * model = llama_get_model(ctx); @@ -43,16 +43,16 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu // run model if(embeddings_debug) { - printf("\n%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq); + printf("\n%s: n_tokens = %d, n_seq = %d\n", __func__, batch.size(), n_seq); } // run model - if (llama_decode(ctx, batch) < 0) { + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) < 0) { printf("%s : failed to process\n", __func__); } - for (int i = 0; i < batch.n_tokens; i++) { - if (!batch.logits[i]) { + for (int i = 0; i < batch.size(); i++) { + if (!batch.tokens[i].output) { continue; } const float * embd = nullptr; @@ -67,8 +67,8 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu } } else { // try to get sequence embeddings - supported only when pooling_type is not NONE - embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]); - embd_pos = batch.seq_id[i][0]; + embd = llama_get_embeddings_seq(ctx, batch.tokens[i].seq_id); + embd_pos = batch.tokens[i].seq_id; if(embd == NULL) { printf("\nfailed to get sequence embeddings\n"); @@ -220,7 +220,7 @@ embeddings_generation_outputs embeddingstype_generate(const embeddings_generatio // initialize batch const int n_prompts = 1; const enum llama_pooling_type pooling_type = llama_pooling_type(embeddings_ctx); - struct llama_batch batch = llama_batch_init(n_batch, 0, 1); + common_batch batch(embeddings_ctx); // count number of embeddings int n_embd_count = 0; @@ -247,12 +247,12 @@ embeddings_generation_outputs embeddingstype_generate(const embeddings_generatio auto & inp = prompt_inputs[k]; const uint64_t n_toks = inp.size(); // encode if at capacity - if (batch.n_tokens + n_toks > n_batch) { + if (batch.size() + n_toks > n_batch) { float * out = emb + e * n_embd; batch_decode(embeddings_ctx, batch, out, s, n_embd, embd_normalize); - e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s; + e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.size() : s; s = 0; - common_batch_clear(batch); + batch.clear(); } // add to batch batch_add_seq(batch, inp, s); @@ -274,9 +274,6 @@ embeddings_generation_outputs embeddingstype_generate(const embeddings_generatio outputarray += "]"; last_output = outputarray; - // clean up - llama_batch_free(batch); - timetaken = timer_check(); printf("\nText Embeddings Generated %d values in %.2fs.\n",(int) n_embd,timetaken); diff --git a/tools/completion/completion.cpp b/tools/completion/completion.cpp index 83da454ce..1b1fb1f65 100644 --- a/tools/completion/completion.cpp +++ b/tools/completion/completion.cpp @@ -526,10 +526,9 @@ int llama_completion(int argc, char ** argv) { } if (llama_model_has_encoder(model)) { - int enc_input_size = embd_inp.size(); - llama_token * enc_input_buf = embd_inp.data(); + common_batch batch = common_batch_get_one(ctx, embd_inp); - if (llama_encode(ctx, llama_batch_get_one(enc_input_buf, enc_input_size))) { + if (llama_process(ctx, LLAMA_PROCESS_TYPE_ENCODE, batch.get())) { LOG_ERR("%s : failed to eval\n", __func__); return 1; }