diff --git a/common/common.cpp b/common/common.cpp index d8319cd9ac..a316b2ae17 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2198,9 +2198,28 @@ bool common_replay_last_token(struct llama_context * ctx, llama_token last_token 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)); + + auto mem = llama_get_memory(ctx); + llama_pos pos = mem ? llama_memory_seq_pos_max(mem, 0) + 1 : 0; + + 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); + pos++; + } + + if (!tokens.empty()) { + llama_batch_ext_set_output_logits(batch.get(), (int32_t) tokens.size() - 1, true); + } + + return batch; +} + bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, @@ -2221,7 +2240,9 @@ bool common_prompt_batch_decode( // Memory implementations in recurrent/hybrid models don't support removing tokens from their // 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. - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_tokens_before_last))) { + 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); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_prefix.get())) { COM_ERR("%s", "failed to eval\n"); return false; } @@ -2231,17 +2252,19 @@ bool common_prompt_batch_decode( 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 batch = llama_batch_get_one(&last_token, 1); - int32_t pos = n_past; - batch.pos = &pos; + 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); - if (llama_decode(ctx, batch)) { + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_last.get())) { COM_ERR("%s", "failed to eval last token\n"); return false; } n_past++; } else { - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_new))) { + 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); + 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 7afc266acd..39f855c4f2 100644 --- a/common/common.h +++ b/common/common.h @@ -1021,6 +1021,10 @@ void common_batch_add( const std::vector & seq_ids, bool logits); +// 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); + // 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 @@ -1028,7 +1032,7 @@ void common_batch_add( // tokens from memory, so this approach works across all model architectures. bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, diff --git a/include/llama-cpp.h b/include/llama-cpp.h index 8f6368177d..880a6a5fae 100644 --- a/include/llama-cpp.h +++ b/include/llama-cpp.h @@ -24,7 +24,12 @@ struct llama_adapter_lora_deleter { void operator()(llama_adapter_lora * adapter) { llama_adapter_lora_free(adapter); } }; +struct llama_batch_ext_deleter { + void operator()(llama_batch_ext * batch) { llama_batch_ext_free(batch); } +}; + typedef std::unique_ptr llama_model_ptr; typedef std::unique_ptr llama_context_ptr; typedef std::unique_ptr llama_sampler_ptr; typedef std::unique_ptr llama_adapter_lora_ptr; +typedef std::unique_ptr llama_batch_ext_ptr; diff --git a/include/llama.h b/include/llama.h index 31bbf8b0d9..1805ed0559 100644 --- a/include/llama.h +++ b/include/llama.h @@ -293,6 +293,11 @@ extern "C" { LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA, }; + enum llama_process_type { + LLAMA_PROCESS_TYPE_ENCODE, + LLAMA_PROCESS_TYPE_DECODE, + }; + struct llama_model_kv_override { enum llama_model_kv_override_type tag; @@ -999,6 +1004,91 @@ extern "C" { struct llama_context * ctx, struct llama_batch batch); + // + // Extended batch API + // + + struct llama_batch_ext; + + struct llama_embd { + const float * data; + size_t n_rows; // number of embedding rows in data + size_t n_embd; // size of one row + }; + + LLAMA_API struct llama_batch_ext * llama_batch_ext_init (struct llama_context * ctx); + LLAMA_API void llama_batch_ext_free (struct llama_batch_ext * batch); + LLAMA_API void llama_batch_ext_clear(struct llama_batch_ext * batch); + + // Add an input token to the batch, with default values: + // id = LLAMA_TOKEN_NULL + // embd = nullptr + // pos = not set, the caller must set it with llama_batch_ext_set_pos() + // Returns the batch index (>= 0) + // On error: + // -1: batch is full + // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) + // -3: invalid sequence id + LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); + + // Add an input token to the batch, with a specified token ID or token embedding + LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); + LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, llama_seq_id seq_id, struct llama_embd embd); + + // Add the token at index idx in the batch to another sequence id. The position will stays the same. + // Note: this should be called before other _set() functions + LLAMA_API bool llama_batch_ext_add_seq( + struct llama_batch_ext * batch, + int32_t idx, + llama_seq_id seq_id); + + // Set the token embedding for the token at index idx in the batch + // use it after llama_batch_ext_add_token() to have an entry with both a token id and an embedding + LLAMA_API bool llama_batch_ext_set_embd_token( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); + + // Set the "state" embedding for the token at index idx in the batch + // "state" here means extra hidden state carried over from a previous stage, e.g.: + // - MTP: state from N layers of the target model + // - Qwen3 VL (deepstack): state from N layers of the vision encoder + LLAMA_API bool llama_batch_ext_set_embd_state( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); + + // Set if output embeddings should be available for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output logits + LLAMA_API bool llama_batch_ext_set_output_embd( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + + // Set output logits for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output embd + LLAMA_API bool llama_batch_ext_set_output_logits( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + + // Set custom position for the token at index idx in the batch + // For M-RoPE models: + // - Embedding tokens must have multiple positions per token + // - Text token only requires one single position per token + LLAMA_API bool llama_batch_ext_set_pos( + struct llama_batch_ext * batch, + int32_t idx, + const llama_pos * pos); + + // TODO: implement get_embeddings() and get_logits() for llama_batch_ext + + // Return values are the same as llama_decode() + LLAMA_API int32_t llama_process( + struct llama_context * ctx, + enum llama_process_type type, + struct llama_batch_ext * batch); + // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) // n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f4..89a1f3f37c 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -3,6 +3,9 @@ #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h" +#include "llama-hparams.h" +#include "llama-model.h" +#include "llama-context.h" #include #include @@ -23,43 +26,120 @@ llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd } bool llama_batch_allocr::init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all) { clear(); - batch = batch_inp; + this->vocab = &vocab; + this->n_embd = batch_inp.n_embd > 0 ? batch_inp.n_embd : batch_inp.n_embd_inp; + this->n_seq_max = batch_inp.n_seq_max; - this->vocab = &vocab; + const int32_t n_tok = (int32_t) batch_inp.tokens.size(); - GGML_ASSERT(batch.n_tokens > 0); + GGML_ASSERT(n_tok > 0); - // - // validate input batch - // - - if (n_seq_max > LLAMA_MAX_SEQ) { + if ((uint32_t) n_seq_max > LLAMA_MAX_SEQ) { LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ); return false; } - if (batch.token) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.token[i] < 0 || (uint32_t) batch.token[i] >= vocab.n_tokens()) { - LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, batch.token[i]); + const llama_memory_i * mem = batch_inp.mem; + + // + // determine the content types of the batch + // an entry can carry a token id, a token embedding, or both (e.g. MTP hook batches) + // all entries must carry the same combination + // + + const bool has_token = batch_inp.tokens[0].id != LLAMA_TOKEN_NULL; + const bool has_embd = batch_inp.tokens[0].has_embd; + + for (int32_t i = 1; i < n_tok; ++i) { + if ((batch_inp.tokens[i].id != LLAMA_TOKEN_NULL) != has_token || + batch_inp.tokens[i].has_embd != has_embd) { + LLAMA_LOG_ERROR("%s: all entries in the batch must have the same content types\n", __func__); + return false; + } + } + + if (!has_token && !has_embd) { + LLAMA_LOG_ERROR("%s: batch has neither token ids nor embeddings\n", __func__); + return false; + } + + // + // build flat token/embd array + // + + if (has_token) { + token_vec.resize(n_tok); + for (int32_t i = 0; i < n_tok; ++i) { + const llama_token id = batch_inp.tokens[i].id; + if (id < 0 || id >= batch_inp.n_vocab) { + LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, id); return false; } + token_vec[i] = id; + } + } + + if (has_embd) { + embd_vec = batch_inp.embd; + } + + // + // build flat pos array + // token batch: pos[i] = tokens[i].pos[0] + // embedding batch: pos[j*n_tok + i] = tokens[i].pos[j] (section-major) + // + + { + const int32_t n_pos_total = has_token ? n_tok : n_tok * (int32_t) n_pos_per_embd; + pos.resize(n_pos_total); + if (has_token) { + for (int32_t i = 0; i < n_tok; ++i) { + pos[i] = batch_inp.tokens[i].pos[0]; + } + } else { + for (int32_t i = 0; i < n_tok; ++i) { + for (uint32_t j = 0; j < n_pos_per_embd; ++j) { + pos[(int32_t) j * n_tok + i] = batch_inp.tokens[i].pos[j]; + } + } } } - if (batch.seq_id) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) { - LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, batch.seq_id[i][s], (llama_seq_id) n_seq_max); + // + // build n_seq_id / seq_id arrays + // + + n_seq_id.resize(n_tok); + seq_id.resize(n_tok + 1); + seq_id[n_tok] = nullptr; + + { + size_t total = 0; + for (int32_t i = 0; i < n_tok; ++i) { + total += batch_inp.tokens[i].seq_ids.size(); + } + seq_id_data.reserve(total); + + for (int32_t i = 0; i < n_tok; ++i) { + for (auto sid : batch_inp.tokens[i].seq_ids) { + seq_id_data.push_back(sid); + } + } + + size_t off = 0; + for (int32_t i = 0; i < n_tok; ++i) { + n_seq_id[i] = (int32_t) batch_inp.tokens[i].seq_ids.size(); + seq_id[i] = seq_id_data.data() + off; + off += n_seq_id[i]; + + for (int32_t s = 0; s < n_seq_id[i]; ++s) { + if (seq_id[i][s] < 0 || seq_id[i][s] >= (llama_seq_id) n_seq_max) { + LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, seq_id[i][s], (llama_seq_id) n_seq_max); return false; } } @@ -67,91 +147,43 @@ bool llama_batch_allocr::init( } // - // auto-generate missing fields + // build output/logits array // - if (!batch.n_seq_id) { - n_seq_id.resize(batch.n_tokens); - for (int32_t i = 0; i < batch.n_tokens; i++) { - n_seq_id[i] = seq_id_0.size(); - } - batch.n_seq_id = n_seq_id.data(); - } - - if (!batch.seq_id) { - seq_id.resize(batch.n_tokens + 1); - seq_id[batch.n_tokens] = NULL; - for (int32_t i = 0; i < batch.n_tokens; i++) { - seq_id[i] = seq_id_0.data(); - } - batch.seq_id = seq_id.data(); - } - - if (!batch.pos) { - pos.resize(batch.n_tokens); - - // initialize the starting position for each sequence based on the positions in the memory - llama_pos p0[LLAMA_MAX_SEQ]; - for (uint32_t s = 0; s < n_seq_max; ++s) { - if (!memory) { - // if no memory -> start from 0 - p0[s] = 0; - } else { - p0[s] = memory->seq_pos_max(s) + 1; - } + { + output.resize(n_tok, 0); + for (int32_t i = 0; i < n_tok; ++i) { + output[i] = batch_inp.tokens[i].output ? 1 : 0; } - for (int32_t i = 0; i < batch.n_tokens; i++) { - const llama_seq_id seq_id = batch.seq_id[i][0]; - - pos[i] = p0[seq_id]; - - // update the starting position for all sequences that are assigned to the this token - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - const llama_seq_id seq_id = batch.seq_id[i][s]; - - p0[seq_id] = pos[i] + 1; - } - } - - batch.pos = pos.data(); - } - - if (!batch.logits) { if (output_all) { - // return the output for all tokens - output.resize(batch.n_tokens, true); - } else { - // return the output only for the last token - output.resize(batch.n_tokens, false); - output[output.size() - 1] = true; - } - - batch.logits = output.data(); - } else if (output_all) { - bool warn = false; - - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.logits[i] == 0) { - warn = true; + bool warn = false; + for (int32_t i = 0; i < n_tok; ++i) { + if (!output[i]) { warn = true; break; } + } + if (warn) { + LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); + std::fill(output.begin(), output.end(), 1); } } - - if (warn) { - LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); - - output.resize(batch.n_tokens, true); - batch.logits = output.data(); - } } + // + // set up the internal llama_batch to point to our owned arrays + // + + batch.n_tokens = n_tok; + batch.token = has_token ? token_vec.data() : nullptr; + batch.embd = has_embd ? embd_vec.data() : nullptr; + batch.pos = pos.data(); + batch.n_seq_id = n_seq_id.data(); + batch.seq_id = seq_id.data(); + batch.logits = output.data(); + // // compute stats // - this->n_embd = n_embd; - this->n_seq_max = n_seq_max; - // count the outputs in this batch for (int32_t i = 0; i < batch.n_tokens; ++i) { n_outputs += batch.logits[i] != 0; @@ -259,7 +291,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (batch.token) { if (p0 >= 0 && p0 >= seq_pos_min(s)) { @@ -292,7 +324,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (p0 >= 0) { bool ok = true; @@ -320,12 +352,12 @@ bool llama_batch_allocr::init( } } - if (memory) { + if (mem) { for (uint32_t s0 = 0; s0 < n_seq_max; ++s0) { for (uint32_t s1 = 0; s1 < n_seq_max; ++s1) { if (seq_cpl[s0][s1]) { - if (memory->seq_pos_min(s0) != memory->seq_pos_min(s1) || - memory->seq_pos_max(s0) != memory->seq_pos_max(s1)) { + if (mem->seq_pos_min(s0) != mem->seq_pos_min(s1) || + mem->seq_pos_max(s0) != mem->seq_pos_max(s1)) { LLAMA_LOG_ERROR("%s: sequence %d is coupled to %d in the input batch, but have divereged\n", __func__, s0, s1); return false; } @@ -725,11 +757,14 @@ void llama_batch_allocr::clear() { batch = {}; - pos .clear(); - n_seq_id .clear(); - seq_id .clear(); - seq_id_unq.clear(); - output .clear(); + token_vec .clear(); + embd_vec .clear(); + seq_id_data .clear(); + pos .clear(); + n_seq_id .clear(); + seq_id .clear(); + seq_id_unq .clear(); + output .clear(); for (auto & cur : seq_pos) { cur.clear(); @@ -985,3 +1020,305 @@ void llama_batch_free(struct llama_batch batch) { } if (batch.logits) free(batch.logits); } + + +// llama_batch_ext + +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams) { + if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) { + return hparams.n_embd_out(); + } + if (arch == LLM_ARCH_DFLASH) { + return hparams.n_embd_inp_enc(); + } + return hparams.n_embd_inp(); +} + +llama_batch_ext::llama_batch_ext(llama_context * ctx) : + n_tokens_max(llama_n_batch(ctx)), + n_embd_inp(llama_batch_ext_select_n_embd_inp(ctx->get_cparams().ctx_type, llama_get_model(ctx)->arch, llama_get_model(ctx)->hparams)), + n_embd_inp_enc(llama_get_model(ctx)->hparams.n_embd_inp_enc()), + n_seq_max(llama_n_seq_max(ctx)), + mem(llama_get_memory(ctx)), + n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), + n_pos_per_embd(llama_get_model(ctx)->hparams.n_pos_per_embd()) { + clear(); +} + +llama_batch_ext::llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd) : + n_tokens_max(n_tokens_max), + n_embd_inp(n_embd_inp), + n_embd_inp_enc(n_embd_inp_enc), + n_seq_max(n_seq_max), + mem(mem), + n_vocab(n_vocab), + n_pos_per_embd(n_pos_per_embd) { + clear(); +} + +void llama_batch_ext::clear() { + tokens.clear(); + embd .clear(); + n_embd = 0; +} + +int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { + if (tokens.size() >= n_tokens_max) { + return -1; // size limit reached + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return -3; // invalid sequence id + } + + // position is left undefined; call set_token_pos() before decoding + token t; + t.seq_ids.insert(seq_id); + + tokens.push_back(t); + + return (int32_t)(tokens.size() - 1); +} + +bool llama_batch_ext::add_seq(int32_t idx, llama_seq_id seq_id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return false; + } + + token & t = tokens[idx]; + + t.seq_ids.insert(seq_id); + + return true; +} + +bool llama_batch_ext::set_token_id(int32_t idx, llama_token id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (id < 0 || id >= n_vocab) { + return false; + } + tokens[idx].id = id; + return true; +} + +bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!embd_in.data) { + return false; + } + + const size_t n_total = embd_in.n_rows * embd_in.n_embd; + if (n_embd == 0) { + if (n_total != n_embd_inp && n_total != n_embd_inp_enc) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu or %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd_inp, n_embd_inp_enc); + return false; + } + n_embd = n_total; + } else if (n_total != n_embd) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd); + return false; + } + + token & t = tokens[idx]; + + if (t.has_embd) { + LLAMA_LOG_ERROR("%s: embedding for token %d is already set\n", __func__, idx); + return false; + } + + t.has_embd = true; + t.embd_off = embd.size(); + embd.insert(embd.end(), embd_in.data, embd_in.data + n_total); + + return true; +} + +bool llama_batch_ext::set_token_pos(int32_t idx, const llama_pos * pos_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!pos_in) { + return false; + } + + token & t = tokens[idx]; + + size_t n_pos = t.id != LLAMA_TOKEN_NULL ? 1 : n_pos_per_embd; + for (size_t i = 0; i < n_pos; ++i) { + t.pos[i] = pos_in[i]; + } + + return true; +} + +bool llama_batch_ext::set_output(int32_t idx, bool output_last) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + tokens[idx].output = output_last; + return true; +} + +// llama_batch_ext C API + +llama_batch_ext * llama_batch_ext_init(llama_context * ctx) { + return new llama_batch_ext(ctx); +} + +void llama_batch_ext_free(llama_batch_ext * batch) { + delete batch; +} + +void llama_batch_ext_clear(llama_batch_ext * batch) { + batch->clear(); +} + +int32_t llama_batch_ext_add(llama_batch_ext * batch, llama_seq_id seq_id) { + return batch->add_token(seq_id); +} + +int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, llama_token id) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_id(idx, id)) { + return -2; + } + return idx; +} + +int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, llama_embd embd) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_embd(idx, embd)) { + return -2; + } + return idx; +} + +bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id seq_id) { + return batch->add_seq(idx, seq_id); +} + +bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, const llama_pos * pos) { + return batch->set_token_pos(idx, pos); +} + +bool llama_batch_ext_set_embd_token(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + return batch->set_token_embd(idx, embd); +} + +bool llama_batch_ext_set_embd_state(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + // TODO + GGML_UNUSED(batch); + GGML_UNUSED(idx); + GGML_UNUSED(embd); + return false; +} + +bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + +bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + +// llama_batch_compat + +void llama_batch_compat::init(llama_batch_ext & dst, const llama_batch & batch_inp, size_t n_embd_row) { + llama_batch_ext * batch_ext = &dst; + + if (n_embd_row == 0) { + n_embd_row = batch_ext->n_embd_inp; + } + + // a batch can carry both, for example the MTP hook batches + const bool has_token = batch_inp.token != nullptr; + const bool has_embd = batch_inp.embd != nullptr; + + static const llama_seq_id default_seq_id = 0; + static const int32_t default_n_seq_id = 1; + + // auto-generates positions locally when batch_inp.pos is null, continuing from memory + std::vector pos_next(batch_ext->n_seq_max); + for (llama_seq_id s = 0; s < (llama_seq_id) batch_ext->n_seq_max; ++s) { + pos_next[s] = llama_memory_seq_pos_max(batch_ext->mem, s) + 1; // assume next pos + } + + for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { + const int32_t n_sid = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : default_n_seq_id; + const llama_seq_id * sids = batch_inp.seq_id ? batch_inp.seq_id[i] : &default_seq_id; + + llama_batch_ext::token t; + + // seq_ids + for (int32_t s = 0; s < n_sid; ++s) { + t.seq_ids.insert(sids[s]); + } + + // position(s) + if (batch_inp.pos) { + if (has_token) { + // token batch: one position per token + t.pos[0] = batch_inp.pos[i]; + } else { + // embedding batch (M-RoPE): section-major layout pos[j*n_tokens + i] + for (uint32_t j = 0; j < batch_ext->n_pos_per_embd; ++j) { + t.pos[j] = batch_inp.pos[(int32_t) j * batch_inp.n_tokens + i]; + } + } + } else { + // auto-generate position from the first seq_id + t.pos[0] = pos_next[sids[0]]++; + } + + // token id and/or embeddings + if (has_token) { + t.id = batch_inp.token[i]; + } + + if (has_embd) { + t.has_embd = true; + t.embd_off = batch_ext->embd.size(); + const float * src = batch_inp.embd + (size_t) i * n_embd_row; + batch_ext->embd.insert(batch_ext->embd.end(), src, src + n_embd_row); + batch_ext->n_embd = n_embd_row; + } + + // output flag + // if no logits array is given, default to only the last token being an output + t.output = batch_inp.logits + ? (batch_inp.logits[i] != 0) + : (i == batch_inp.n_tokens - 1); + + batch_ext->tokens.push_back(t); + } +} + +llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row) { + batch_ext = new llama_batch_ext(ctx); + init(*batch_ext, batch_inp, n_embd_row); +} + +llama_batch_compat::~llama_batch_compat() { + delete batch_ext; +} diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a..201d48cce1 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -2,6 +2,7 @@ #include "llama.h" +#include "llama-arch.h" #include "llama-cparams.h" #include @@ -10,6 +11,7 @@ #include #include #include +#include // keep this struct lightweight struct llama_ubatch { @@ -68,19 +70,71 @@ struct llama_ubatch { std::shared_ptr data; }; +struct llama_hparams; + +// MTP hook batches carry the target model's hidden state (n_embd_out size). +// DFlash batches carry the fused target features at the encoder input width (n_embd_inp_enc size). +// Normal batches carry token embeddings (n_embd_inp size). +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams); + +struct llama_batch_ext { + const size_t n_tokens_max; // max number of tokens that can be stored in the batch + const size_t n_embd_inp; // decoder embd row width + const size_t n_embd_inp_enc; // encoder embd row width (e.g. eagle3/dflash extracted features) + const llama_seq_id n_seq_max; // max number of sequences + llama_memory_i * mem; // memory for position inference + const llama_token n_vocab; // max token ID that we accept + const size_t n_pos_per_embd; + + // actual embd row width of this batch, set by the first set_token_embd() + // must be either n_embd_inp or n_embd_inp_enc; encode/decode verify it against the graph input + size_t n_embd = 0; + + struct token { + llama_token id = LLAMA_TOKEN_NULL; + bool has_embd = false; // whether embd_off is set + size_t embd_off = 0; // index offset in the embd array + bool output = false; // TODO: have dedicated output flags + std::unordered_set seq_ids; + std::array pos = {0, 0, 0, 0}; + }; + std::vector tokens; + std::vector embd; + + llama_batch_ext(llama_context * ctx); + + // build without a llama_context, used by tests + llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd); + + void clear(); + + // add an entry with an undefined position + // the caller must set it explicitly via set_token_pos() + int32_t add_token(llama_seq_id seq_id); + + bool add_seq(int32_t idx, llama_seq_id seq_id); + bool set_token_id(int32_t idx, llama_token id); + bool set_token_embd(int32_t idx, llama_embd embd_in); + bool set_token_pos(int32_t idx, const llama_pos * pos_in); + bool set_output(int32_t idx, bool output_last); +}; + // a helper for sanitizing, fulfilling and splitting a batch class llama_batch_allocr { public: llama_batch_allocr(uint32_t n_pos_per_embd); - // sanitize and auto-gen missing data in the input batch - // memory is optional. if provided will be used to check for sequence continuity and to determine the positions + // convert a llama_batch_ext to internal llama_batch and sanitize it bool init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all); const llama_batch & get_batch() const; @@ -137,7 +191,9 @@ private: uint32_t n_seq_max; uint32_t n_outputs; - std::array seq_id_0 = {{ 0 }}; // default sequence id + std::vector token_vec; // owned token IDs built from llama_batch_ext + std::vector embd_vec; // owned embeddings built from llama_batch_ext + std::vector seq_id_data; // flat storage for seq_id pointers below std::vector pos; std::vector n_seq_id; @@ -172,3 +228,16 @@ private: int debug; }; + +// RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext +struct llama_batch_compat { + llama_batch_ext * batch_ext; + + // n_embd_row is the embd row width of batch_inp, 0 = use the decoder width + llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0); + ~llama_batch_compat(); + + // fill an existing llama_batch_ext from a llama_batch (old API) + // note: this is called directly by the tests, skipping llama_context creation + static void init(llama_batch_ext & batch_ext, const llama_batch & batch_inp, size_t n_embd_row = 0); +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index fcd4dfb13b..da9c558d28 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1463,24 +1463,25 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return res; } -int llama_context::encode(const llama_batch & batch_inp) { - // MTP hook batches carry both token (next-token id) and embd (h_nextn row), - // so accept either present rather than requiring exactly one. - GGML_ASSERT(batch_inp.token || batch_inp.embd); - - if (batch_inp.n_tokens == 0) { +int llama_context::encode(const llama_batch_ext & batch_inp) { + if (batch_inp.tokens.empty()) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; } const auto & hparams = model.hparams; + if (batch_inp.n_embd > 0 && batch_inp.n_embd != hparams.n_embd_inp_enc()) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the encoder input %u\n", + __func__, batch_inp.n_embd, hparams.n_embd_inp_enc()); + return -1; + } + // eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim - const int64_t n_embd = hparams.n_embd_inp_enc(); const int64_t n_vocab = model.vocab.n_tokens(); - // note: during encode, we always pass the full sequence starting from pos = 0 - if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { + // note: during encode, we always output all tokens and skip position continuity checks (output_all=true) + if (!balloc->init(batch_inp, model.vocab, true)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -1701,29 +1702,27 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map 0 && batch_inp.n_embd != batch_inp.n_embd_inp) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the decoder input %zu\n", + __func__, batch_inp.n_embd, batch_inp.n_embd_inp); + return -1; + } + const auto & vocab = model.vocab; const auto & hparams = model.hparams; const int64_t n_vocab = vocab.n_tokens(); - const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd; - // DFlash embd batches carry the fused target features at the encoder input width - const bool dflash_embd = model.arch == LLM_ARCH_DFLASH && batch_inp.embd; - const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : dflash_embd ? hparams.n_embd_inp_enc() : hparams.n_embd_inp(); // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; @@ -1731,20 +1730,17 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; - // embedding contexts output every token even when batch.logits is not set - if (has_samplers && (output_all || batch_inp.logits)) { + // TODO: avoid this workaround in the future + // embedding contexts output every token even when no token is explicitly marked as output + if (has_samplers) { std::vector seq_output_count(n_seq_max, 0); - for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { - if (!output_all && batch_inp.logits[i] == 0) { + for (const auto & tok : batch_inp.tokens) { + if (!output_all && !tok.output) { continue; } - const int ns = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : 1; - - for (int32_t s = 0; s < ns; ++s) { - const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; - + for (auto seq_id : tok.seq_ids) { if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { continue; } @@ -1762,7 +1758,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } } - if (!balloc->init(batch_inp, vocab, memory.get(), n_embd, n_seq_max, output_all)) { + if (!balloc->init(batch_inp, vocab, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -3557,9 +3553,13 @@ void llama_context::opt_epoch_iter( batch.logits [pos_batch] = true; } - if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); - return; + // TODO: use llama_batch_ext here + { + llama_batch_compat compat(this, batch); + if (!balloc->init(*compat.batch_ext, model.vocab, true)) { + LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + return; + } } const uint32_t n_tokens_all = balloc->get_n_tokens(); @@ -4310,6 +4310,18 @@ size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, lla } } +// compat: llama_batch -> llama_batch_ext -> encode/decode + +int llama_context::encode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp, model.hparams.n_embd_inp_enc()); + return encode(*compat.batch_ext); +} + +int llama_context::decode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp); + return decode(*compat.batch_ext); +} + /// int32_t llama_encode( @@ -4399,6 +4411,14 @@ void llama_opt_epoch( callback_eval); } +int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { + switch (type) { + case LLAMA_PROCESS_TYPE_ENCODE: return ctx->encode(*batch); + case LLAMA_PROCESS_TYPE_DECODE: return ctx->decode(*batch); + } + return -1; +} + // // ext // diff --git a/src/llama-context.h b/src/llama-context.h index 77ef92fc6a..b403b099b7 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -141,6 +141,10 @@ struct llama_context { llama_memory_context_i * mctx, ggml_status & ret); + int encode(const llama_batch_ext & batch_inp); + int decode(const llama_batch_ext & batch_inp); + + // compat version int encode(const llama_batch & batch_inp); int decode(const llama_batch & batch_inp); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 34b3c68801..b83f45ba9e 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -283,7 +283,8 @@ bool llama_hparams::is_ple(uint32_t il) const { } uint32_t llama_hparams::n_pos_per_embd() const { - return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1; + return (rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE) + ? GGML_MROPE_SECTIONS : 1; } bool llama_hparams::is_swa(uint32_t il) const { diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index 66d29d6f51..ad186c6936 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -3,6 +3,8 @@ #include "llama.h" #include "../src/llama-batch.h" +#include "../src/llama-arch.h" +#include "../src/llama-hparams.h" #include "../src/llama-memory.h" #include "../src/llama-vocab.h" @@ -47,49 +49,55 @@ struct mock_memory : public llama_memory_i { void state_read (llama_io_read_i &, llama_seq_id, llama_state_seq_flags) override { GGML_ASSERT(false && "not implemented"); } }; -// builds embedding batches - an empty llama_vocab rejects all token ids, so -// the tests use embeddings everywhere except the token validation tests +// builds a llama_batch_ext without a llama_context +// n_vocab = 0 by default, so every token id is invalid and the tests use embeddings unless stated otherwise struct batch_builder { - uint32_t n_embd; + const uint32_t n_embd; - std::vector embd; - std::vector pos; - std::vector n_seq_id; - std::vector logits; + llama_batch_ext b; - std::vector> seq; - std::vector seq_ptr; + batch_builder( + uint32_t n_embd = 2, + llama_memory_i * mem = nullptr, + llama_seq_id n_seq_max = 4, + uint32_t n_pos_per_embd = 1, + llama_token n_vocab = 0, + uint32_t n_embd_inp_enc = 0) + : n_embd(n_embd), + b(/*n_tokens_max*/ 64, n_embd, n_embd_inp_enc > 0 ? n_embd_inp_enc : n_embd, n_seq_max, mem, n_vocab, n_pos_per_embd) {} - batch_builder(uint32_t n_embd = 2) : n_embd(n_embd) {} - - // embd values are 100*i + k so that ubatch contents can be traced back to batch indices - void add(llama_pos p, std::initializer_list seq_ids, bool output) { - const int32_t i = (int32_t) seq.size(); - for (uint32_t k = 0; k < n_embd; ++k) { - embd.push_back(100.0f*i + k); + // one embedding row for batch index i, values 100*i + k so ubatch contents can be traced back + std::vector row(int32_t i, uint32_t width) const { + std::vector r(width); + for (uint32_t k = 0; k < width; ++k) { + r[k] = 100.0f*i + k; } - pos.push_back(p); - n_seq_id.push_back((int32_t) seq_ids.size()); - seq.emplace_back(seq_ids); - logits.push_back(output ? 1 : 0); + return r; } - llama_batch make(bool with_pos = true, bool with_seq = true, bool with_logits = true) { - seq_ptr.clear(); - for (auto & s : seq) { - seq_ptr.push_back(s.data()); + // embedding entry with full M-RoPE positions + int32_t add_embd(const llama_pos * pos, std::initializer_list seq_ids, bool output, uint32_t width = 0) { + width = width > 0 ? width : n_embd; + + auto it = seq_ids.begin(); + const int32_t idx = b.add_token(*it); + GGML_ASSERT(idx >= 0); + for (++it; it != seq_ids.end(); ++it) { + GGML_ASSERT(b.add_seq(idx, *it)); } - seq_ptr.push_back(nullptr); - llama_batch res = {}; - res.n_tokens = (int32_t) seq.size(); - res.embd = embd.data(); - res.pos = with_pos ? pos.data() : nullptr; - res.n_seq_id = with_seq ? n_seq_id.data() : nullptr; - res.seq_id = with_seq ? seq_ptr.data() : nullptr; - res.logits = with_logits ? logits.data() : nullptr; + const auto r = row(idx, width); + GGML_ASSERT(b.set_token_embd(idx, { r.data(), 1, width })); + GGML_ASSERT(b.set_token_pos(idx, pos)); + GGML_ASSERT(b.set_output(idx, output)); - return res; + return idx; + } + + // embedding entry with a single sequential position + int32_t add(llama_pos p, std::initializer_list seq_ids, bool output) { + const llama_pos pos[GGML_MROPE_SECTIONS] = { p, 0, 0, 0 }; + return add_embd(pos, seq_ids, output); } }; @@ -97,22 +105,31 @@ static void test_init(testing & t) { llama_vocab vocab; t.test("rejects_n_seq_max_too_large", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, LLAMA_MAX_SEQ + 1); bb.add(0, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, LLAMA_MAX_SEQ + 1, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_token", [&](testing & t) { - llama_token tok = 0; // empty vocab -> every token id is out of range - llama_batch batch = llama_batch_get_one(&tok, 1); + // n_vocab = 0 -> every token id is out of range + // set_token_id() refuses such ids, so the token is poked directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); llama_batch_allocr ba(1); - t.assert_true("token id >= n_tokens", !ba.init(batch, vocab, nullptr, 0, 1, false)); - tok = -1; - t.assert_true("negative token id", !ba.init(batch, vocab, nullptr, 0, 1, false)); + t.assert_true("set_token_id refuses out of range id", !bb.b.set_token_id(idx, 0)); + + bb.b.tokens[idx].id = 0; + t.assert_true("token id >= n_vocab", !ba.init(bb.b, vocab, false)); + + bb.b.tokens[idx].id = -1; + t.assert_true("negative token id", !ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_seq_id", [&](testing & t) { @@ -120,33 +137,44 @@ static void test_init(testing & t) { { batch_builder bb; - bb.add(0, {4}, true); - t.assert_true("seq_id >= n_seq_max", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true("add_token refuses seq_id >= n_seq_max", bb.b.add_token(4) == -3); + t.assert_true("add_token refuses negative seq_id", bb.b.add_token(-1) == -3); + } + { + // poke the seq_ids directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { 4 }; + t.assert_true("seq_id >= n_seq_max", !ba.init(bb.b, vocab, false)); } { batch_builder bb; - bb.add(0, {-1}, true); - t.assert_true("negative seq_id", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { -1 }; + t.assert_true("negative seq_id", !ba.init(bb.b, vocab, false)); } }); - t.test("autofill_defaults", [&](testing & t) { + t.test("copies_pos_seq_output", [&](testing & t) { batch_builder bb; for (int i = 0; i < 4; ++i) { - bb.add(0, {0}, false); + bb.add(i, {0}, i == 3); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, false, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); const llama_batch & batch = ba.get_batch(); t.assert_equal(4u, ba.get_n_tokens()); + t.assert_true("embedding batch", batch.embd != nullptr); + t.assert_true("no token ids", batch.token == nullptr); for (int i = 0; i < 4; ++i) { - t.assert_equal("pos defaults to 0..n-1", i, batch.pos[i]); - t.assert_equal("n_seq_id defaults to 1", 1, batch.n_seq_id[i]); - t.assert_equal("seq_id defaults to 0", 0, batch.seq_id[i][0]); + t.assert_equal(i, batch.pos[i]); + t.assert_equal(1, batch.n_seq_id[i]); + t.assert_equal(0, batch.seq_id[i][0]); + t.assert_equal(100.0f*i, batch.embd[i*bb.n_embd]); } t.assert_equal("only the last token is an output", 1u, ba.get_n_outputs()); @@ -165,7 +193,7 @@ static void test_init(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, true)); + t.assert_true(ba.init(bb.b, vocab, true)); t.assert_equal(4u, ba.get_n_outputs()); }); @@ -176,7 +204,7 @@ static void test_init(testing & t) { bb.add(2, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal(2u, ba.get_n_outputs()); llama_ubatch ub = ba.split_simple(10); @@ -191,17 +219,17 @@ static void test_init(testing & t) { t.assert_equal(2, out_ids[1]); }); - t.test("pos_from_memory", [&](testing & t) { + t.test("pos_after_memory", [&](testing & t) { mock_memory mem; mem.ranges[0] = {0, 9}; - batch_builder bb; + batch_builder bb(2, &mem); for (int i = 0; i < 3; ++i) { - bb.add(0, {0}, false); + bb.add(10 + i, {0}, false); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, true, false), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal("pos continues after memory", 10, ba.seq_pos_min(0)); t.assert_equal(12, ba.seq_pos_max(0)); @@ -214,22 +242,22 @@ static void test_init(testing & t) { llama_batch_allocr ba(1); { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(10, {0}, false); bb.add(11, {0}, true); - t.assert_true("pos_max + 1 is accepted", ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("pos_max + 1 is accepted", ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(11, {0}, false); bb.add(12, {0}, true); - t.assert_true("gap after memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("gap after memory is rejected", !ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(9, {0}, false); bb.add(10, {0}, true); - t.assert_true("overlap with memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("overlap with memory is rejected", !ba.init(bb.b, vocab, false)); } }); @@ -240,7 +268,7 @@ static void test_init(testing & t) { bb.add(3, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_decreasing_positions", [&](testing & t) { @@ -253,7 +281,7 @@ static void test_init(testing & t) { // seq 0 sees positions 4,5,6,3 in batch order -> the trailing 3 decreases llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("allows_equal_positions_in_seq", [&](testing & t) { @@ -263,23 +291,143 @@ static void test_init(testing & t) { bb.add(1, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); }); - t.test("rejects_coupled_diverged_seqs", [&](testing & t) { - batch_builder bb; - bb.add(6, {0, 1}, true); - llama_batch_allocr ba(1); mock_memory mem; mem.ranges[0] = {0, 5}; mem.ranges[1] = {2, 5}; // same pos_max, different pos_min -> diverged - t.assert_true(!ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(!ba.init(bb.b, vocab, false)); + } mem.ranges[1] = {0, 5}; - t.assert_true(ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(ba.init(bb.b, vocab, false)); + } + }); +} + +static void test_content_types(testing & t) { + llama_vocab vocab; + + t.test("token_and_embd_together", [&](testing & t) { + // e.g. MTP hook batches: a token id and its embedding on the same entry + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + const int32_t idx = bb.b.add_token(0); + t.assert_true(bb.b.set_token_id(idx, 3)); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + const llama_batch & batch = ba.get_batch(); + t.assert_true("token ids are kept", batch.token != nullptr); + t.assert_true("embeddings are kept", batch.embd != nullptr); + t.assert_equal(3, batch.token[0]); + t.assert_equal(0.0f, batch.embd[0]); + t.assert_equal(1.0f, batch.embd[1]); + + llama_ubatch ub = ba.split_simple(1); + t.assert_true(ub.token != nullptr && ub.embd != nullptr); + t.assert_equal(3, ub.token[0]); + }); + + t.test("rejects_mixed_content_types", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + // entry 0: token only, entry 1: token + embd + const llama_pos p0 = 0; + const llama_pos p1 = 1; + + int32_t i0 = bb.b.add_token(0); + bb.b.set_token_id(i0, 1); + bb.b.set_token_pos(i0, &p0); + + int32_t i1 = bb.b.add_token(0); + bb.b.set_token_id(i1, 2); + const auto r = bb.row(i1, bb.n_embd); + bb.b.set_token_embd(i1, { r.data(), 1, bb.n_embd }); + bb.b.set_token_pos(i1, &p1); + bb.b.set_output(i1, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_neither_token_nor_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_embd_size_mismatch", [&](testing & t) { + batch_builder bb; // n_embd = 2, n_embd_inp_enc = 2 + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 8); + + t.assert_true("too small", !bb.b.set_token_embd(idx, { r.data(), 1, 1 })); + t.assert_true("too large", !bb.b.set_token_embd(idx, { r.data(), 1, 3 })); + t.assert_true("zero rows", !bb.b.set_token_embd(idx, { r.data(), 0, 2 })); + t.assert_true("null data", !bb.b.set_token_embd(idx, { nullptr, 1, 2 })); + t.assert_true("same total via a different split is accepted", bb.b.set_token_embd(idx, { r.data(), 2, 1 })); + }); + + t.test("rejects_double_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + }); + + t.test("encoder_width", [&](testing & t) { + // e.g. eagle3/dflash: extracted features are wider than the decoder input + const uint32_t n_embd_enc = 6; + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + + const llama_pos p0 = 0; + const llama_pos p1 = 1; + bb.add_embd(&p0, {0}, false, n_embd_enc); + bb.add_embd(&p1, {0}, true, n_embd_enc); + + t.assert_equal("batch width follows the first embedding", (size_t) n_embd_enc, bb.b.n_embd); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + // the ubatch uses the encoder stride: token 1 starts at offset n_embd_enc + llama_ubatch ub = ba.split_simple(2); + t.assert_equal(2u, ub.n_tokens); + t.assert_equal(100.0f, ub.embd[n_embd_enc]); + t.assert_equal(105.0f, ub.embd[n_embd_enc + 5]); + }); + + t.test("rejects_mixing_widths", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, 0, /*n_embd_inp_enc*/ 6); + + const llama_pos p0 = 0; + bb.add_embd(&p0, {0}, false, 2); // first entry fixes the batch width to 2 + + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 6); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, 6 })); }); } @@ -293,7 +441,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); @@ -336,7 +484,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); while (ba.split_simple(1).n_tokens > 0) { } @@ -359,7 +507,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(8, false, 0); t.assert_true(ub.equal_seqs()); @@ -395,7 +543,7 @@ static void test_split(testing & t) { bb.add(1, {0, 1}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, true, 0); t.assert_equal("sequential split rejects coupled seqs", 0u, ub.n_tokens); @@ -417,7 +565,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); for (llama_seq_id s = 0; s < 3; ++s) { llama_ubatch ub = ba.split_seq(8); @@ -459,14 +607,14 @@ static void test_keep_tail(testing & t) { } ++s; } - return bb.make(); }; t.test("noop_when_seqs_complete", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 2}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 2}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("both seqs fit whole", 4u, ub.n_tokens); @@ -478,9 +626,10 @@ static void test_keep_tail(testing & t) { t.test("defers_seq_with_short_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq: seq 0 completes, seq 1 would be left // with 1 < n_keep_tail remaining, so it is deferred entirely @@ -504,9 +653,10 @@ static void test_keep_tail(testing & t) { t.test("completes_first_seq_when_all_violate", [&](testing & t) { batch_builder bb; + make_batch(bb, {3, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq, leaving both with 1 < n_keep_tail remaining; // seq 0 still fits in n_ubatch, so it is extended to completion and emitted alone @@ -528,9 +678,10 @@ static void test_keep_tail(testing & t) { t.test("truncates_to_preserve_tail", [&](testing & t) { batch_builder bb; + make_batch(bb, {5}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {5}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // 4 tokens would leave a remainder of 1, and the seq does not fit in n_ubatch, // so the ubatch is truncated until n_keep_tail tokens remain @@ -551,9 +702,10 @@ static void test_keep_tail(testing & t) { t.test("keeps_full_ubatch_with_sufficient_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {6}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {6}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("remainder >= n_keep_tail, no truncation", 4u, ub.n_tokens); @@ -567,10 +719,11 @@ static void test_keep_tail(testing & t) { }); t.test("multi_seq_prefix_kept", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, 6); + make_batch(bb, {3, 4}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 4}), vocab, nullptr, bb.n_embd, 6, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 3 tokens per seq: seq 0 completes, seq 1 has 1 < n_keep_tail // remaining and is deferred even though its tokens were already gathered @@ -591,32 +744,26 @@ static void test_mrope(testing & t) { llama_vocab vocab; t.test("pos_layout_and_split", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; - batch_builder bb(n_embd); - bb.add(10, {0}, false); - bb.add(11, {0}, true); + batch_builder bb(n_embd, nullptr, 4, n_pos); - // M-RoPE positions for embeddings are laid out [n_pos][n_tokens] - std::vector pos = { - 10, 11, // temporal - 5, 6, // y - 7, 8, // x - 0, 0, - }; - - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); + // M-RoPE positions per embedding: [temporal, y, x, other] + const llama_pos pos0[n_pos] = { 10, 5, 7, 0 }; + const llama_pos pos1[n_pos] = { 11, 6, 8, 0 }; + bb.add_embd(pos0, {0}, false); + bb.add_embd(pos1, {0}, true); llama_batch_allocr ba(n_pos); - t.assert_true(ba.init(batch, vocab, nullptr, n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); t.assert_equal(n_pos, ub.n_pos); t.assert_true(ub.is_pos_2d()); + // the ubatch stores positions section-major: [n_pos][n_tokens] const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; for (int i = 0; i < 8; ++i) { t.assert_equal(expected[i], ub.pos[i]); @@ -624,7 +771,7 @@ static void test_mrope(testing & t) { }); t.test("pos_jump_allowed", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; mock_memory mem; @@ -633,15 +780,12 @@ static void test_mrope(testing & t) { llama_batch_allocr ba(n_pos); auto try_pos = [&](llama_pos p0) { - batch_builder bb(n_embd); - bb.add(p0, {0}, true); + batch_builder bb(n_embd, &mem, 4, n_pos); - std::vector pos = {p0, 1, 1, 0}; + const llama_pos pos[n_pos] = { p0, 1, 1, 0 }; + bb.add_embd(pos, {0}, true); - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); - - return ba.init(batch, vocab, &mem, n_embd, 4, false); + return ba.init(bb.b, vocab, false); }; t.assert_true("gap after memory is allowed", try_pos(15)); @@ -650,6 +794,254 @@ static void test_mrope(testing & t) { }); } +// conversion from the old llama_batch API (llama_batch_compat::init) +static void test_compat(testing & t) { + llama_vocab vocab; + + t.test("token_batch_explicit_fields", [&](testing & t) { + llama_token token[3] = { 5, 6, 7 }; + llama_pos pos[3] = { 3, 4, 5 }; + int32_t n_seq_id[3] = { 1, 1, 2 }; + llama_seq_id s0[1] = { 1 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id s2[2] = { 1, 2 }; + llama_seq_id * seq_id[4] = { s0, s1, s2, nullptr }; + int8_t logits[3] = { 0, 1, 0 }; + + llama_batch lb = {}; + lb.n_tokens = 3; + lb.token = token; + lb.pos = pos; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + lb.logits = logits; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + t.assert_true("no embeddings", bb.b.embd.empty() && bb.b.n_embd == 0); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_equal(pos[i], bb.b.tokens[i].pos[0]); + t.assert_true(!bb.b.tokens[i].has_embd); + t.assert_equal(logits[i] != 0, bb.b.tokens[i].output); + } + t.assert_equal((size_t) 1, bb.b.tokens[0].seq_ids.size()); + t.assert_true(bb.b.tokens[0].seq_ids.count(1) == 1); + t.assert_equal((size_t) 2, bb.b.tokens[2].seq_ids.size()); + t.assert_true(bb.b.tokens[2].seq_ids.count(1) == 1 && bb.b.tokens[2].seq_ids.count(2) == 1); + + // round trip through the allocator + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true(batch.token != nullptr && batch.embd == nullptr); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], batch.token[i]); + t.assert_equal(pos[i], batch.pos[i]); + } + t.assert_equal(1u, ba.get_n_outputs()); + }); + + t.test("defaults_for_null_fields", [&](testing & t) { + // llama_batch_get_one: only token and n_tokens are set + mock_memory mem; + mem.ranges[0] = {0, 9}; + + llama_token token[3] = { 5, 6, 7 }; + llama_batch lb = llama_batch_get_one(token, 3); + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + for (int i = 0; i < 3; ++i) { + t.assert_equal("pos continues after memory", 10 + i, bb.b.tokens[i].pos[0]); + t.assert_equal("seq_id defaults to 0", (size_t) 1, bb.b.tokens[i].seq_ids.size()); + t.assert_true(bb.b.tokens[i].seq_ids.count(0) == 1); + } + t.assert_true("only the last token is an output", !bb.b.tokens[0].output && !bb.b.tokens[1].output && bb.b.tokens[2].output); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + t.assert_equal(10, ba.seq_pos_min(0)); + t.assert_equal(12, ba.seq_pos_max(0)); + }); + + t.test("auto_pos_starts_at_zero_without_memory", [&](testing & t) { + llama_token token[2] = { 5, 6 }; + llama_batch lb = llama_batch_get_one(token, 2); + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal(0, bb.b.tokens[0].pos[0]); + t.assert_equal(1, bb.b.tokens[1].pos[0]); + }); + + t.test("auto_pos_is_tracked_per_seq", [&](testing & t) { + mock_memory mem; + mem.ranges[0] = {0, 9}; // seq 1 is empty + + llama_token token[4] = { 5, 6, 7, 8 }; + int32_t n_seq_id[4] = { 1, 1, 1, 1 }; + llama_seq_id s0[1] = { 0 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id * seq_id[5] = { s0, s1, s0, s1, nullptr }; + + llama_batch lb = {}; + lb.n_tokens = 4; + lb.token = token; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal("seq 0 continues after memory", 10, bb.b.tokens[0].pos[0]); + t.assert_equal("seq 1 starts from 0", 0, bb.b.tokens[1].pos[0]); + t.assert_equal(11, bb.b.tokens[2].pos[0]); + t.assert_equal( 1, bb.b.tokens[3].pos[0]); + }); + + t.test("embd_batch_with_mrope_positions", [&](testing & t) { + const uint32_t n_pos = 4; + const uint32_t n_embd = 2; + + float embd[2*n_embd] = { 0, 1, 100, 101 }; + // section-major layout: pos[j*n_tokens + i] + llama_pos pos[n_pos*2] = { + 10, 11, // temporal + 5, 6, // y + 7, 8, // x + 0, 0, + }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(n_embd, nullptr, 4, n_pos); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 2, bb.b.tokens.size()); + t.assert_equal("batch width", (size_t) n_embd, bb.b.n_embd); + for (int i = 0; i < 2; ++i) { + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal(LLAMA_TOKEN_NULL, bb.b.tokens[i].id); + t.assert_equal((size_t) i*n_embd, bb.b.tokens[i].embd_off); + for (uint32_t j = 0; j < n_pos; ++j) { + t.assert_equal(pos[j*2 + i], bb.b.tokens[i].pos[j]); + } + } + t.assert_equal(100.0f, bb.b.embd[2]); + t.assert_equal(101.0f, bb.b.embd[3]); + + llama_batch_allocr ba(n_pos); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; + for (int i = 0; i < 8; ++i) { + t.assert_equal(expected[i], ub.pos[i]); + } + }); + + t.test("token_and_embd_both_set", [&](testing & t) { + // e.g. MTP hook batches + llama_token token[2] = { 5, 6 }; + float embd[4] = { 0, 1, 100, 101 }; + llama_pos pos[2] = { 3, 4 }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.token = token; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + for (int i = 0; i < 2; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal("one position per token", pos[i], bb.b.tokens[i].pos[0]); + } + t.assert_equal(100.0f, bb.b.embd[2]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true("both kept", batch.token != nullptr && batch.embd != nullptr); + }); + + t.test("embd_row_width_override", [&](testing & t) { + // encoder input (e.g. eagle3/dflash) is wider than the decoder input + const uint32_t n_embd_enc = 6; + float embd[2*n_embd_enc]; + for (int i = 0; i < 2*6; ++i) { + embd[i] = (float) i; + } + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + llama_batch_compat::init(bb.b, lb, n_embd_enc); + + t.assert_equal((size_t) n_embd_enc, bb.b.n_embd); + t.assert_equal((size_t) 2*n_embd_enc, bb.b.embd.size()); + t.assert_equal((size_t) n_embd_enc, bb.b.tokens[1].embd_off); + t.assert_equal(6.0f, bb.b.embd[n_embd_enc]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + t.assert_equal("ubatch uses the encoder stride", 6.0f, ub.embd[n_embd_enc]); + }); +} + +static void test_mtp_embd_width(testing & t) { + t.test("mtp_uses_n_embd_out", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; + hparams.n_deepstack_layers = 2; // makes n_embd_inp() = 64 + 64*2 = 192 + hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96 + + t.assert_equal("default context uses n_embd_inp (deepstack-aware)", + (size_t) 192, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + + t.assert_equal("MTP context uses n_embd_out instead (target-model hidden state width)", + (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams)); + }); + + t.test("mtp_falls_back_to_n_embd_when_no_override", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; // no deepstack, no n_embd_out_impl override + + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams)); + }); + + t.test("dflash_uses_n_embd_inp_enc", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; + hparams.n_embd_inp_enc_impl = 128; // makes n_embd_inp_enc() = 128 + hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96 + + t.assert_equal("DFlash uses the encoder input width", + (size_t) 128, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_DFLASH, hparams)); + + t.assert_equal("other archs ignore n_embd_inp_enc", + (size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + + t.assert_equal("MTP takes precedence over DFlash", + (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_DFLASH, hparams)); + }); +} + int main(int argc, char ** argv) { testing t; @@ -665,10 +1057,13 @@ int main(int argc, char ** argv) { t.set_filter(argv[1]); } - t.test("init", test_init); - t.test("split", test_split); - t.test("keep_tail", test_keep_tail); - t.test("mrope", test_mrope); + t.test("init", test_init); + t.test("content_types", test_content_types); + t.test("compat", test_compat); + t.test("split", test_split); + t.test("keep_tail", test_keep_tail); + t.test("mrope", test_mrope); + t.test("mtp_embd_width", test_mtp_embd_width); return t.summary(); }