llama: add llama_batch_ext by ngxson · Pull Request #24669 · ggml-org/llama.cpp

ngxson · 2026-06-15T21:15:21Z

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 · 2026-06-16T11:30:50Z

+    // Set output = true for the last added token in the batch
+    // Returns the batch index (>= 0)
+    LLAMA_API bool llama_batch_ext_set_output(
+            struct llama_batch_ext * batch,
+                            int32_t  idx,
+                               bool  output_last);


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

ggerganov · 2026-06-16T11:36:43Z

+        float         * embd_nextn; // used by nextn layers
+        llama_pos     * pos;        // if nullptr, the position will be automatically assigned
+                                    // for M-RoPE models, embedding tokens must have multiple positions per token; text token only requires one single position per token
+        llama_seq_id    seq_id;


We still want to support multiple sequence ids per token.

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 ?

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.

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:

llama_batch_ext_add() that simply returns an idx, and need a separated _set_token(id) or _set_embd(float * embd)

llama_batch_ext_add(id), if id == LLAMA_TOKEN_NULL then _set_embd(float *) is required

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 ?

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);

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 · 2026-06-17T09:36:07Z

@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 · 2026-06-21T12:54:30Z

+    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);


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 · 2026-07-13T17:03:11Z

@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 · 2026-08-07T14:05:25Z

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 · 2026-08-07T14:20:42Z

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 · 2026-08-10T07:38:58Z

+    llama_token n_vocab; // max token ID that we accept
+    size_t n_pos_per_embd;
+
+    std::vector<llama_pos> pos_max; // keep track of the current position


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.

Removed pos_max in 3b86110

# Conflicts: # src/llama-context.cpp

ngxson · 2026-08-13T15:39:00Z

Re. your comment about embd API, I agree with the llama_embd

Just one thing I want to check quickly with you though: which naming version should we take:

embd_token from last-layer state embd from MTP target model and embd_mtmd for mtmd input embd
embd_state for MTP and embd_token for mtmd?

Edit: personally I prefer the later one, because it's better reflected in the current llama-graph: the build_inp_embd() currently means "token embd"

ggerganov · 2026-08-13T17:29:02Z

embd_state for MTP and embd_token for mtmd?

Sounds good. Just not sure about the Qwen3 VL mtmd embeddings - I think they had some "deepstack" associated with them. Do we need to differentiate them from embd_token?

ngxson · 2026-08-13T17:42:49Z

Deepstack is currently considered as part of the token embd (what's why we have the llama_n_embd_inp())

But I think we can reuse embd_state for that, with the "state" here means "extra hidden state from last stage"

For example:

Input state of MTP is the tokens + the embd of N last layers (normally N=1, but IIRC gemma 4 has N=3)
Input state of qwen3 VL is the vision tokens + the embd from N layers of the vision encoder

(wip) add llama_batch_ext

0dfa5e1

ngxson requested a review from ggerganov as a code owner June 15, 2026 21:15

wip

c7c5468

ngxson marked this pull request as draft June 15, 2026 21:15

ngxson mentioned this pull request Jun 15, 2026

(superseded) llama : add llama_batch_ext #11875

Closed

ngxson assigned ggerganov Jun 15, 2026

ggerganov reviewed Jun 16, 2026

View reviewed changes

ngxson added 2 commits June 16, 2026 22:29

updated design

c0fb071

updated impl

3132371

ngxson mentioned this pull request Jun 20, 2026

server: refactor batch construction #24843

Merged

3 tasks

ggerganov reviewed Jun 21, 2026

View reviewed changes

ggerganov mentioned this pull request Jul 7, 2026

speculative: fix MTP draft crash on vision inputs #25144

Draft

ngxson added 5 commits July 13, 2026 17:33

Merge branch 'master' into xsn/llama_batch_ext

32343b2

change signature

bf372b3

unused var

901ed68

demo common_prompt_batch_decode

231af77

fix pos

4cd8c26

tmp disable test-batch-alloc

e4c474f

github-actions Bot added the testing Everything test related label Jul 13, 2026

fix compat

4ba39e5

ruixiang63 mentioned this pull request Jul 23, 2026

[Speculative decoding] feat: add DFlash support #22105

Merged

ggerganov reviewed Aug 10, 2026

View reviewed changes

ngxson mentioned this pull request Aug 10, 2026

model: Muse Glimmer Support #26841

Merged

ngxson added 2 commits August 13, 2026 01:24

Merge remote-tracking branch 'upstream/master' into xsn/llama_batch_ext

7db50e9

# Conflicts: # src/llama-context.cpp

nits: add const

7d626f5

no more pos_max

3b86110

		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);

	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;
	}

	// 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_tokensn_embdggml_element_size(h));
	}

	// 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];

Uh oh!

Conversation

ngxson commented Jun 15, 2026

Overview

Additional information

Requirements

Uh oh!

ggerganov left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngxson Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngxson commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngxson commented Jul 13, 2026

Uh oh!

ngxson commented Aug 7, 2026

Uh oh!

ggerganov commented Aug 7, 2026

Uh oh!

ggerganov left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngxson commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ggerganov commented Aug 13, 2026

Uh oh!

ngxson commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ggerganov left a comment •

edited

Loading

ngxson Jun 16, 2026 •

edited

Loading

ngxson commented Jun 17, 2026 •

edited

Loading

ggerganov left a comment •

edited

Loading

ngxson commented Aug 13, 2026 •

edited

Loading