llama: add llama_batch_ext

[ngxson]

Overview

Supersede #11875

Status: early WIP, for discussion only

Additional information

Demo usage:

auto * batch = llama_batch_ext_init(ctx);
int32_t last_idx = 0;
for (auto token_id : input_tokens) {
  llama_batch_token t{
    token_id,
    nullptr, // embd
    nullptr, // embd_nextn
    nullptr, // pos
    0 // seq_id
  };
  last_idx = llama_batch_ext_add_token(batch, t);
}

llama_batch_ext_set_output(batch, last_idx, true);
llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch); // process the prompt

while (true) {
  float * logits = llama_batch_ext_get_logits(batch, last_idx);

  // Sample the next token from the logits
  // optionally check for stop condition
  llama_token next_token_id = sample_next_token(logits);

  // Process the sampled token
  llama_batch_token t{
    next_token_id,
    nullptr, // embd
    nullptr, // embd_nextn
    nullptr, // pos
    0 // seq_id
  };
  int32_t idx = llama_batch_ext_add_token(batch, t);
  llama_batch_ext_set_output(batch, idx, true);
  int32_t result = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch); // process the next token
  if (result != 0) {
    break; // stop if there is an error or end of sequence
  }
}

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: TBD

[ggerganov]

To keep it simpler, we can do the output-related logic in a next PR. I.e. in the first PR, we just introduce the llama_batch_ext and use it to pass the inputs, but we leave the llama_context to handle the output buffers and embedding extractions as it is. Then in the next PR, we will move all the output logic to the batch.

[ggerganov]

The comment is incorrect here - it's not the last added token, but the idx token.

[ggerganov]

We still want to support multiple sequence ids per token.

[ngxson]

hmm tbh I don't quite like passing pointer-to-pointer llama_seq_id **, as it makes the caller do more works. so I'm wondering if we should redesign it to get rid of struct llama_batch_token. my idea now is to have 2 categories of calls:

_add call that returns batch index:

• llama_batch_ext_add_token --> add by token ID
• llama_batch_ext_add_embd --> add by embeddings

then an array of _set that adds more info to the returned batch index:

• llama_batch_ext_set_embd_nextn(int32_t idx, float * embd)
• llama_batch_ext_set_seq_id(int32_t idx, llama_seq_id * seq_id, size_t n_seq) --> can set to multiple sequences
• llama_batch_ext_set_pos
• llama_batch_ext_set_output

also, do you think _set_output should be a boolean, or it should be a bit field, for example LLAMA_OUTPUT_NEXTN | LLAMA_OUTPUT_EMBD ?

[ggerganov]

so I'm wondering if we should redesign it to get rid of struct llama_batch_token

Yes, API based on the entry index seems quite generic and clean.

llama_batch_ext_add_token --> add by token ID
llama_batch_ext_add_embd --> add by embeddings

I think you can simplify by having single llama_batch_ext_add instead of differentiating token/embd.

also, do you think _set_output should be a boolean, or it should be a bit field, for example LLAMA_OUTPUT_NEXTN | LLAMA_OUTPUT_EMBD ?

I think we probably need per-entry llama_batch_ext_set_output(batch, idx, value);. And then the contents of the outputs likely not have to be per-entry, but for the entire batch:

llama_batch_ext_output_embd      (batch, value); 
llama_batch_ext_output_embd_nextn(batch, value, masked);
llama_batch_ext_output_layer_inp (batch, value); 
...

But for now these can remain controlled by the llama_context for now because the llama_context currently will own the output buffers, not the batch.

[ngxson]

I think you can simplify by having single llama_batch_ext_add instead of differentiating token/embd.

actually there will be 3 choices for that:

1. llama_batch_ext_add() that simply returns an idx, and need a separated _set_token(id) or _set_embd(float * embd)
2. llama_batch_ext_add(id), if id == LLAMA_TOKEN_NULL then _set_embd(float *) is required
3. llama_batch_ext_add(id, float * embd) where either one of two can be set

however, since each input entry in the batch requires at least token ID or token embd to be consider "valid", I think having explicit _add_token and _add_embd will be a better design overall. WDYT ?

[ggerganov]

We can have all three:

int32_t llama_batch_ext_add      (llama_batch_ext * batch);
int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_token id);
int32_t llama_batch_ext_add_embd (llama_batch_ext * batch, llama_embd embd);

[ngxson]

note that llama_batch_ext_add() is currently not useful as-is because there is no calls to attach either token or embd to a idx in batch --> may need to define how it can be used in the future

[ngxson]

@ggerganov Ok so I've implemented the first working PoC of the new API:

• llama_process accepts llama_batch_ext as input
• llama_batch_allocr takes llama_batch_ext as input (no more llama_batch here)
• To convert llama_batch to the _ext version, we have a small llama_batch_compat that acts as a RAII compat layer

I think llama_batch_allocr need a rework too, but probably better to do it with the _get_output API

[ggerganov]

nit: this seems a bit more consistent:

Suggested change

	LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_token id, llama_seq_id seq_id);
	LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, float * embd, llama_seq_id seq_id);
	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, float * embd);

[ngxson]

@ggerganov I made a bit of progress here, I migrated common_prompt_batch_decode to llama_batch_ext to demo how it can be used. Will continue working on this next week.

I'm noting down things that might be missing (probably follow-up PRs), but feel free to review this PR as-is:

• Positional tracking is a bit messy. This might be problematic if we create one batch then decode twice. Example: batch [t0:p0, t1:p1], when decode second time we expect position to increased [t0:p2, t1:p3], but that is not currently the case --> user need to create a new batch so that position can increase
• Currently, adding both text tokens and embd to the batch is UB --> maybe add a simple assert for now and fix it in the future

[ngxson]

gentle ping @ggerganov if you have bandwidth to review this in the next few days. I'm hopping to finalize it this month, probably shipped with a core libllama PR first (this one) then a follow-up to migrate it everywhere else in the code base

[ggerganov]

Yes, sorry for the delay. Will prioritize this.

[ggerganov]

Currently, adding both text tokens and embd to the batch is UB --> maybe add a simple assert for now and fix it in the future

Yes, the assertion should guard that all entries in the batch have the same content types.

For example, if we say that atm we support 3 types of content for each batch entry: token id (A), token embd (B), mtmd embd (C)

# NOT allowed: batch with mixed content types
   0 1 2 3 4 5 6 ...
A: x x x x . . . ...
B: . . . . x x x ...
C: . . . . . . . ...

# NOT allowed
   0 1 2 3 4 5 6 ...
A: x x x x x x x ...
B: . . . . x x x ...
C: . . . . . . . ...


# allowed: all entries have the same content types
   0 1 2 3 4 5 6 ...
A: x x x x x x x ...
B: x x x x x x x ...
C: . . . . . . . ...

# allowed
   0 1 2 3 4 5 6 ...
A: . . . . . . . ...
B: x x x x x x x ...
C: x x x x x x x ...

# allowed
   0 1 2 3 4 5 6 ...
A: x x x x x x x ...
B: . . . . . . . ...
C: . . . . . . . ...

# allowed
   0 1 2 3 4 5 6 ...
A: x x x x x x x ...
B: . . . . . . . ...
C: x x x x x x x ...

---

One of the main use cases that we have to validate with this refactor is to be able to pass more than one type of embeddings for a batch entry. For example:

• multi-modal embeddings
• target-model embeddings

The use case is for multi-modal speculative decoding, for example here:

llama.cpp/common/speculative.cpp

Lines 1402 to 1411 in 92d1bb0

	bool process(const llama_batch & batch_in) override {
	if (batch_in.n_tokens <= 0) {
	return true;
	}
	// TODO: how to make it work with vision tokens?
	if (batch_in.token == nullptr || batch_in.embd != nullptr) {
	return true;
	}

llama.cpp/src/llama-graph.cpp

Lines 105 to 112 in 92d1bb0

	// TODO: extend llama_ubatch to differentiate between token embeddings and hidden states
	// for now, we assume that the hidden state is always provided as an embedding
	// ref: https://github.com/ggml-org/llama.cpp/pull/23643
	if (ubatch->embd) {
	GGML_ASSERT(n_embd == h->ne[0]);
	ggml_backend_tensor_set(h, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h));
	}

To do that properly, I am thinking we need to formalize the "embeddings" object a bit better. Likely introduce struct llama_embd which carries some notion of the embeddings types. Currently we are aware of the following embeddings types:

• token embeddings
• mtmd embeddings (these come from a vision encoder and can have different size from the regular token embeddings)
• target-model embeddings (during spec decoding, we can extract embeddings from the target model, often referred to as "activations")

We will probably need to support more types in the future, so the API should be able to extend easily.

Also, one shortcoming of the current way we treat embeddings is that they don't carry information about their size (i.e. we pass a raw float pointer), so we need to infer their size from other things which is a bit fragile. Therefore the llama_embd should also carry this information too.

[ggerganov]

This state is used to auto-generate the next position when adding entries to the batch. However, the logic assumes append-only usage. For example, set_token_pos() can desync the contents of this array (e.g. by overriding the current max pos with a smaller pos).

Correct position tracking can be done like we do it in llama_kv_cells:

llama.cpp/src/llama-kv-cells.h

Lines 491 to 499 in 92d1bb0

	// the set seq_pos[s][p] tells us how many times the position p is currently present for sequence s
	// if the position p is not present, seq_pos[s][p] is not set
	// this way seq_pos[s].begin() and seq_pos[s].rbegin() give us the min/max positions currently in the cache
	//
	// note that we cannot a use an std::set because in some cases a position can occur more than once for the same seq:
	// - during performing a cache reuse via (rm + add)
	// - some vision models have input embeddings with repeating positions
	//
	std::map<llama_pos, int> seq_pos[LLAMA_MAX_SEQ];

But this logic might be too heavy to use here.

I think for now we don't need to keep track of the maximum pos as part of the llama_batch_ext state. In the llama_batch_compat we can have a local state tracking the max position, starting from the memory's max pos (similar to the logic in master). And the llama_batch_ext_add APIs will add entries with undefined positions - the user would have to set the correct position explicitly.