What's the difference between batch-size and ubatch-size?

[Sunt-ing]

Thanks for your help!

[phymbert]

Hello, good question!

• --batch-size size of the logits and embeddings buffer, which limits the maximum batch size passed to llama_decode. For the server, this is the maximum number of tokens per iteration during continuous batching

• --ubatch-size physical maximum batch size for computation

So batch size is at the application level, while ubatch size is at the device level. batch_size >= ubatch_size.

You can find some references here:

• llama : add pipeline parallelism support #6017
• https://github.com/ggerganov/llama.cpp/blob/master/examples/server/README.md

[Sunt-ing]

Thanks for your provided info! So it looks like this is for pipeline parallelization. For example, if I use batch-size=2048 and ubatch-size=512, then it's a 4-stage pipeline parallelization. Do I understand correctly?

Also, I find that in the main example, the default batch-size is 512, while in the server doc it's 2048. Is it correct?

[phymbert]

The default values are here:

llama.cpp/common/common.h

Line 57 in 557410b

int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS)

[Sunt-ing]

Thanks for your careful and patient explanation!

[Sunt-ing]

And I think maybe we can update the doc as its default value is outdated: https://github.com/ggerganov/llama.cpp/blob/master/examples/main/README.md#batch-size

[phymbert]

Yes, feel free to open a PR

[dlyz]

https://github.com/ggerganov/llama.cpp/blob/b0d943de179ad5dbd83d51f327fb566066f4ccda/llama.cpp#L11432C1-L11434C14

@phymbert
Am I missing something or it may not work correctly from the API standpoint, when n_batch > n_ubatch and some of the first ubatches has been already committed to the kv-cache and decoded, but no slot found for the last ubatch for example?
What I mean is after this case state of the kv-cache will be kind of corrupted, original batch is partially processed (and we don't know what part of it is) so there is no way to recover from it on the fly (for example do the defrag and try again) and the whole cache should be cleaned. Or, again, I am missing something.

[phymbert]

@slaren probably has a better view

[slaren]

This is something that can happen. It should be possible to detect when this happens by using llama_kv_cache_seq_pos_max, but none of the examples do this. Ideally, I think this should be avoided entirely by reserving the KV space for the entire batch before processing anything, but I am not sure if this is possible with the current functions.