llama : add support for batched inference

[ggerganov]

We want to be able to generate multiple sequences sharing the same context (a.k.a. prompt) in parallel.

Demonstrated in one of the examples by @xaedes :

https://github.com/ggerganov/llama.cpp/blob/eff86d4f1334c08300d3cb1110dbac3c8e26286c/examples/baby-llama/baby-llama.cpp#L785-L794

Should become part of the official llama.cpp API

ref: #2789

Implementation details

Regarding the API for the batched inference functionality, one way is to add a function:

// TODO: better name?
void llama_context_set_parallel(struct llama_context * ctx, int n_batches);

This would reallocate the kv_self cache to fit n_batches batches.

During llama_eval, we do what we normally do, with the extra step of batching the input as demonstrated in the example. We can probably avoid changing the eval API by adding the implicit assumption that tokens will contain the tokens for n_batches batches:

https://github.com/ggerganov/llama.cpp/blob/dd0dc366dab10e8df28d3924e7f313b5c695e908/llama.h#L315-L320

In the end, we just need to update the API for accessing the logits of all the batches, or once again - without changing the API, have an implicit assumption that the results will be for n_batches batches:

https://github.com/ggerganov/llama.cpp/blob/dd0dc366dab10e8df28d3924e7f313b5c695e908/llama.h#L341

---

So on first thought, we would just need a single new function added to llama.h - llama_context_set_parallel().
I think this should be enough, but I could be missing something

One optimization to consider is if we can avoid having separate KV caches for the common prefix of the parallel runs. The straightforward implementation would create a copy of this for each batch, while in theory we need just one. Not sure how complicated it would be to handle this. Might need to implement Paged Attention, which is probably a job for another time

[spion]

Would it be possible to design this API in such a way ot support the following use case?

In some cases, its useful to do constrained evaluation of logits based on a union of possible text values, then pick the sum { logits } (i.e. product(probabilities)) that gives the most probable outcome overall.

E.g. template (using MS guidance)

{{#select 'armor'}}leather{{or}}chainmail{{or}}plate{{/select}}

To definitely make the best choice, we'd need to calculate the probability of all 3 token sequences. Its easy if all the choices map to a single token, but with multiple tokens we'd need not just parallel generation but parallel logit evaluation of multiple possible paths.

If we go greedy, we might get suboptimal results in cases multiple choices start with the same logit.

Appologies if I misunderstood the feature in advance. 😅

[AdityaSher]

is there any issue / improvement with respecting to supporting multiple requests in the server implementation similiarly?

[ggerganov]

@AdityaSher The server serving multiple requests is tangentially related to this functionality. The problem there would be to have a logic that batches the different requests together - but this is high-level logic not related to the llama.cpp library

@spion This sounds related to the new beam-search functionality, although I'm not sure I fully understand the use case that you have in mind.

[spion]

@ggerganov Sort of, but its slightly more constrained, I think. There is a set number of well defined potential paths that need to be evaluated in a manner similar to how beam search does it, but with the constraint that most followup tokens are impossible and don't need to be considered at all (as they would not fit the constraints of the specified union)

[ggerganov]

@spion

I guess you can combine beam search with grammar sampling.
Interesting if we can make this work

cc @ejones and @mattpulver for some insights

[AutonomicPerfectionist]

One optimization to consider is if we can avoid having separate KV caches for the common prefix of the parallel runs.

@ggerganov You can use shared memory/anonymous pages and mmap to map the same physical page to multiple virtual pages, allowing you to reuse the common prompt context without copying it. Only works for CPU side of course, and you can only share full pages of context like this, any cached values that end up in the partial page would need to be either copied manually or recalculated.

So the workflow would go like this:

1. Initially create a shared memory segment and ftruncate to the calculated cache size
2. mmap it for prompt processing
3. When it's time to begin the batched run, calculate how many full pages the prompt context consumes so you know how big you want the shared map
4. Map an anonymous segment with the calculated cache size so the kernel finds a place in the virtual memory space with enough contiguous, free addresses
5. Memory map the prompt context shared segment over the start of the new anonymous segment with MAP_FIXED to force clobber the overlapping pages

Disclaimer: I have not looked too deep into the KV cache code, I just assumed it's a large contiguous segment

Reference:

https://nullprogram.com/blog/2016/04/10/

https://man7.org/linux/man-pages/man2/mmap.2.html

https://groups.google.com/g/comp.os.linux.development.system/c/Prx7ExCzsv4

[ejones]

@ggerganov grammar sampling should work with any stream of tokens. Based on first impressions from the beam search code I wonder if each beam could maintain a grammar (accepting tokens as they're added to a beam) and apply the grammar penalties in the top-k selection? cc @mattpulver

[mattpulver]

Currently llama_grammar has 3 member variables:
https://github.com/ggerganov/llama.cpp/blob/53885d7256909ec3e2176cdc2477f3986c15ec69/llama.cpp#L3434-L3440

If each beam is to hold its own grammar state, then presumably they can share a const reference to rules but must have their own copies of stacks and partial_utf8. These will be copied as beams branch. These can be added to
https://github.com/ggerganov/llama.cpp/blob/53885d7256909ec3e2176cdc2477f3986c15ec69/llama.cpp#L4334-L4337

Then we can abstract out the call to get the next_tokens for each beam, at or around

https://github.com/ggerganov/llama.cpp/blob/53885d7256909ec3e2176cdc2477f3986c15ec69/llama.cpp#L4457

and replace it with a call to a user-supplied callback function that will determine the next_tokens for each beam. The callback will have access to the above llama_grammar variables, and will only need to select the top n_beams tokens for each beam.

Does that sound satisfactory?

[ejones]

Yeah, that sounds about right! rules can definitely be shared. And then as the beams are created/tokens added to beams, the grammar state needs to be updated with (the equivalent of) llama_grammar_accept_token.

[ggerganov]

If each beam is to hold its own grammar state, then presumably they can share a const reference to rules but must have their own copies of stacks and partial_utf8.

How big in terms of memory is rules? If not too big, I would rather just have a copy of the entire grammar in each beam to make things simpler.

[mattpulver]

There was an earlier experiment w/ SQL grammar in which the rules got somewhat large. However to your point, sharing rules is an optimization we can keep separate from this PR either way. Thus the first iteration of this PR can continue to use the existing struct llama_grammar without defining a light-weight llama_grammar_view, etc.

If/when we do decide to implement that optimization, consider replacing

const std::vector<std::vector<llama_grammar_element>> rules;

with

using llama_grammar_rules = std::vector<std::vector<llama_grammar_element>>;
std::shared_ptr<const llama_grammar_rules> rules;

Then the rules can be lightly/freely copied after instantiation.