llama : auto-batch preparation

[ggerganov]

target #13746

The memory implementations now implement their respective optimal strategy for splitting the input batch into ubatches (see llama_memory_i::init_batch()). For example, the iSWA KV cache currently attempts to do a simple split of the input batch. In the future, it will be updated to try different splitting strategies if the simple split fails. For example, the batch splitting logic from llama-server should be done as a fallback. But to implement this, we first need to refactor the llama_sbatch and llama_ubatch implementation.

outdated

This change adds logic to llama_decode() for retrying with smaller n_ubatch if we fail to fit the input batch. This logic is usually implemented in the user code, but it now comes integrated into libllama.

Note that the user code can still continue to do it's own batching of the input. It's just no longer really needed.

As an example, the llama-server is updated to no longer perform this process manually. The rest of the examples will be updated in a follow-up PRs.

This PR will be merged after #13746

Next PRs

• Remove the notion of n_batch. We can now always work with the full n_ctx and simplify the logic of splitting the input into n_batch sizes. (low prio) (not ready for this yet).

[compilade]

Remove the notion of n_batch. We can now always work with the full n_ctx and simplify the logic of splitting the input into n_batch sizes

Depending on how this will be implemented, this can lead to much bigger buffers than necessary for very big n_ctx (e.g 1M tokens for a recurrent model (so the KV cache doesn't scale with n_ctx)). Batches help keeping that small.

[ggerganov]

Remove the notion of n_batch. We can now always work with the full n_ctx and simplify the logic of splitting the input into n_batch sizes

Depending on how this will be implemented, this can lead to much bigger buffers than necessary for very big n_ctx (e.g 1M tokens for a recurrent model (so the KV cache doesn't scale with n_ctx)). Batches help keeping that small.

Yes, in some situations n_batch still makes sense in the user code. But for the llama_context I think this parameter is redundant. Referring to this n_batch:

llama.cpp/src/llama-cparams.h

Lines 9 to 12 in 2921e85

	struct llama_cparams {
	uint32_t n_ctx; // context size used during inference
	uint32_t n_batch;
	uint32_t n_ubatch;

[slaren]

n_batch still puts a limit the size of the output buffer. That's why it was kept when n_ubatch was added, the original implementation removed or ignored n_batch. However since then the output buffers are dynamically reallocated as needed, so it is not as much of a problem anymore.

[ggerganov]

I realized that halving the n_ubatch size and retrying (as proposed originally in this PR) is not equivalent to what the server batching actually does. The halving of n_ubatch strategy is not very smart because there are some degenerate cases that can affect the performance significantly. For example, a large batch with mixed sequences:

# prompt input for sequence 0 and generation tokens for sequences 1 and 2:
00000000[large prompt for seq 0]00012

In some cases, when the SWA cache is full, but has room because of masked tokens still present in the cells, the logic will end up reducing n_ubatch down to 2 in order to be able to fit the multi-sequence tail of the input batch. But processing thousands of tokens with n_ubatch == 2 is very slow.

So what we want to do instead is split into the following ubatches for example:

# ubatch 0
00000000[large prompt for seq 0]00
# ubatch 1
012

This is the split_equal strategy. So the goal would be to implement such kind of logics for each init_batch(). This is just a small preparation for this.