upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-29 01:17:36 -05:00
Code Issues Packages Projects Releases Wiki Activity

llama: add llama_batch_ext (#24669)

Browse Source 
* (wip) add llama_batch_ext

* wip

* updated design

* updated impl

* change signature

* unused var

* demo common_prompt_batch_decode

* fix pos

* tmp disable test-batch-alloc

* fix compat

* nits: add const

* no more pos_max

* add comment about llama_batch_ext_set_embd_state

* handle n_embd_out properly

* rename api --> embd_token

* llama_embd

* stub llama_batch_ext_set_embd_state

* support both token + embd + state in batch

* llama_batch_ext_add_embd

* upstream some changes

* nits

* fix test-batch-alloc

* add test for compat
This commit is contained in:
Xuan-Son Nguyen
2026-09-24 16:25:07 +02:00
committed by GitHub
parent 308883b335
commit fc343a84bb
10 changed files with 1210 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-context.cpp
+52 -32
View File
@@ -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<llama_s
return false; // all sequences use backend sampling
}
int llama_context::decode(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);
int llama_context::decode(const llama_batch_ext & batch_inp) {
if (!memory) {
LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__);
return encode(batch_inp);
}
if (batch_inp.n_tokens == 0) {
if (batch_inp.tokens.empty()) {
LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__);
return -1;
}
if (batch_inp.n_embd > 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<int32_t> 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
//