batch : rework llama_batch_allocr #14153

ggerganov · 2025-06-12T14:36:14Z

ggerganov · 2025-06-12T14:42:14Z

src/llama-context.cpp

+                    for (uint32_t i = 0; i < n_tokens; i++) {
                        const llama_seq_id seq_id = ubatch.seq_id[i][0];


@compilade Regarding this comment from earlier, how does this sequence traversal work correctly when the ubatch is created with split_simple()?

AFAIU the original meaning of ubatch.seq_id[i][j] was "the jth sequence of the ith token". With split_equal(), this now changes to "the ith sequence and j == 0". What is not clear to me is if I used split_simple() how could the sequence traversal be correct?

llama.cpp/src/llama-context.cpp

Lines 1146 to 1152 in f164ba9

for (uint32_t s = 0; s < ubatch.n_seqs; ++s) {

const llama_seq_id seq_id = ubatch.seq_id[s][0];

if (embd_seq_out.find(seq_id) != embd_seq_out.end()) {

continue;

}

embd_seq_out[seq_id].resize(n_embd);

I am planning to rework this in some way, so any suggestions how to improve this logic are welcome.

@ggerganov
With split_simple(), an invariant is that ubatch.n_seqs == n_tokens, and ubatch.n_seq_tokens == 1, because the sequences are not aggregated.

llama.cpp/src/llama-batch.cpp

Line 141 in a592c13

ubatch.n_seqs += ubatch.equal_seqs ? 1 : length; // virtual sequences for simple splits

This makes traversal which would work correctly with split_equal also be correct with split_simple, even though the seq_ids are definitely repeated (when ubatch.equal_seqs == false, ubatch.n_seqs doesn't really map to distinct sequences).

I'm not sure how to make it more obvious while still sharing the same traversal code.

Thanks, I understand that this traversal over the tokens is correct for both split strategies:

llama.cpp/src/llama-kv-cache-unified.cpp

Lines 816 to 822 in 4c07964

for (uint32_t s = 0; s < n_seqs; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[s][0];

for (uint32_t j = 0; j < n_seq_tokens; ++j) {

const uint32_t idx = s*n_seq_tokens + j;

const llama_pos p1 = ubatch->pos[idx];

However, if I want to traverse over the unique sequence ids in the ubatch, or traverse over all sequence ids to which a token in the ubatch is assigned, there is no way to do it correctly for both splits. Is this correct?

For example, in the snippet above, if I wanted to get the list of all sequence ids of token idx, there is no way to do it without checking the ubatch.equal_seqs. Correct?

if I want to traverse over the unique sequence ids in the ubatch

Yes, traversing unique seq_ids with simple splits (when ubatch.equal_seqs == false) is a bit more complicated, because they are not aggregated (simple splits are plain slices of the user-provided batch).

traverse over all sequence ids to which a token in the ubatch is assigned

This is easier, though, and possible by traversing ubatch.seq_id[s][_] with ubatch.n_seq_id[s]. For example:

llama.cpp/src/llama-kv-cache-recurrent.cpp

Lines 446 to 449 in 26ff368

for (uint32_t s = 0; s < n_seqs; ++s) {

const uint32_t n_seq_id = ubatch.n_seq_id[s];

for (uint32_t j = 0; j < n_seq_id; ++j) {

const llama_seq_id seq_id = ubatch.seq_id[s][j];

llama.cpp/src/llama-kv-cache-unified.cpp

Lines 816 to 822 in 4c07964

for (uint32_t s = 0; s < n_seqs; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[s][0];

for (uint32_t j = 0; j < n_seq_tokens; ++j) {

const uint32_t idx = s*n_seq_tokens + j;

const llama_pos p1 = ubatch->pos[idx];

[...]

For example, in the snippet above, if I wanted to get the list of all sequence ids of token idx

In that snippet, seq_id would need to be defined later:

for (uint32_t s = 0; s < n_seqs; ++s) { for (uint32_t j = 0; j < n_seq_tokens; ++j) { const uint32_t idx = s*n_seq_tokens + j; const llama_pos p1 = ubatch->pos[idx]; for (uint32_t k = 0; k < ubatch.n_seq_id[s]; ++k) { const llama_seq_id seq_id = ubatch->seq_id[s][k];

Although depending on what you need it's also possible to swap the two inner loops:

for (uint32_t s = 0; s < n_seqs; ++s) { for (uint32_t k = 0; k < ubatch.n_seq_id[s]; ++k) { const llama_seq_id seq_id = ubatch->seq_id[s][k]; for (uint32_t j = 0; j < n_seq_tokens; ++j) { const uint32_t idx = s*n_seq_tokens + j; const llama_pos p1 = ubatch->pos[idx];

In this situation, you would not need to check ubatch.equal_seqs unless unique sequences are required.

Ok, thank you. I think I understand now.

ggerganov · 2025-06-12T14:50:07Z

src/llama-batch.h

-    int32_t      *  n_seq_id; // [n_seqs] // TODO: remove, should belong to only 1 sequence
-    llama_seq_id ** seq_id;   // [n_seqs] // TODO: become llama_seq_id * seq_id;


Decided against these TODOs because multiple sequences per input token actually has some useful properties that cannot be achieved otherwise (for example see the hellaswag usage). Instead, will add logic to guarantee that the provided ids are valid, utilizing the memory's seq_pos_min() and seq_pos_max() methods.

batch : rework llama_batch_allocr

f164ba9

ggml-ci

ggerganov commented Jun 12, 2025

View reviewed changes

ggerganov added 4 commits June 12, 2025 19:56

cont : move validation inside class

99be6b7

ggml-ci

cont : move output counting to class

75a0d52

ggml-ci

cont : minor

4c07964

ggml-ci

batch : add TODOs

13b0247

ggml-ci

ggerganov marked this pull request as ready for review June 13, 2025 08:02

ggerganov merged commit 60c6663 into master Jun 13, 2025
55 checks passed

ggerganov deleted the gg/batch-allocr-rework branch June 13, 2025 10:48

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

batch : rework llama_batch_allocr #14153

batch : rework llama_batch_allocr #14153

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

batch : rework llama_batch_allocr #14153

batch : rework llama_batch_allocr #14153

Uh oh!

Conversation

Uh oh!

Uh oh!

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!

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!

Uh oh!

Uh oh!