BatchedExecutor Discussion

[martindevans]

Continuing discussion started here: #560 (comment)

Huh, it indeed does, I should've checked it out.
As for feedback, my use case is a bit weird for this project I guess, but I shall share it nontheless :)
Basically the use case for me is that I am making a game in Unity with a lot of prompts.
I started this project of mine some time ago and decided to write my own high-level interface based on NativeAPI/handles directly.

Some notes on what would make BatchExecutor viable for my use-case:

I found that it's impossible for me to keep everything in one KV cache on GPU because of the limited context size, so Save/Load state in-memory is very important. Any ways to spill the context into RAM in seamless and optimal way, while still being able to run inference in batch mode would be very appreciated.

On a side note: I found out that vanilla Executors actually do something like this while writing this PR (I've cleaned KV cache, but if I did not remove Executor state, then LLM actually continued conversation that should've been removed from the KV cache). Really curious if that can be applied to batch setting.

If I process batches too big - frames per second suffer, so I need to be able to batch requests implicitly (e.g. split single prompt/message into multiple runs). If I am not doing it manually, llama.cpp just crashes whole application with assert that batch size is too small, so C# side needs to be smart in this regard.

I'd be happy to help, if that direction would be useful for you.

[martindevans]

Any ways to spill the context into RAM in seamless and optimal way, while still being able to run inference in batch mode would be very appreciated.

Saving/Loading state is actually the thing I'm thinking about at the moment for the BatchedExecutor. I think my ideal interface would be:

• BatchedExecutor.Save() - save the entire executor
• static BatchedExecutor.Load() - create a new executor form a saved state
• Conversation.Save() - save a single conversation
• static BatchedExecutor.LoadConversation() - create a new Conversation resuming from a saved state

However right now I don't think the two Conversation oriented methods are possible. There are methods in llama.cpp to save/load an entire context object (which could be used for the BatchedExecutor load/save), but there's nothing for saving/loading just a part of the cache. I've opened an issue about this here.

If I process batches too big - frames per second suffer, so I need to be able to batch requests implicitly

I think the current executors will take tokens and tries to process them in batch-size chunks, as quickly as possible (i.e. just looping and executing chunks). As far as I am aware this should work with any batch size (down to 1 token).

The BatchedExecutor allows you to prompt as many or as few tokens as you like (here) before running inference. After that you don't have to sample a token if you don't want to, instead you can just immediately prompt with more tokens, so that could be used to feed in a chunk of text in smaller parts.

[eublefar]

Hey

I've tried to rewrite my project to use BatchedExecutor. Love the API and it was mostly drop-in replacement, but there are few issues I struggled with:

• I am not sure how to create empty conversation w/o string prompt, or at least to pass a list of tokens directly as first prompt. This leads to problems with batch size, as there is no direct way to split first prompt into batches of tokens. I handle batching in conversation wrapper in a way you mentioned, but creating conversation causes issues. I'll try passing empty strings and see how it goes, probably should only process BOS token (if this is the way, then maybe empty string being the default parameter is a good quick imporvement for the API + special token handling parameter for tokenization as it might mess with stuff like chat markdown special tokens in the first prompt for other users).
• Infer method throws an exception when internal batch is larger than the batch size from parameters, it would be great if it would split the internal batch into batch queue and process them one by one on overflow. Conversation instances feel like they should be safely used w/o looking at the Executor, but in reality they could "clog" the whole inference thread. I think workaround for me now could be just to check how much tokens are already in the Executor and wait for enough space to put more stuff there.
• I think BatchExecutor is not thread-safe and would benefit a lot from some internal thread-safety, cause now it adds some extra-steps to handle it and I am not 100% sure I am doing this right :)
• It would be great to be able to pre-emptively handle NoKVSlot issues on the users of Conversation class side (e.g. to have some KVSlotsLeft property or just throw NoKVSlot exception on Conversation.Prompt call). In my case it would be to not just dispose of conversation, but also compress conversation history and start a new one with smaller token count that continues the previous one. I handle it now by just having a list of all active conversations and throwing them a callback on NoKVSlot decoderesult, but it's some extra steps that feel like they could probably be avoided smh.

Again, I'd be happy to help implementing some of those changes (as I kinda doing it already on the library user side).
Thanks for the great work on the BatchedExecutor!

[martindevans]

Thanks for this feedback, it's really valuable!

I am not sure how to create empty conversation w/o string prompt

I don't think there's a way to do that at the moment, even at empty string will prompt it with a token. This is something I've been thinking about that, because with the current API we really need to duplicate every single prompt overload from Conversation onto BatchedExecutor just for the initial prompt! I've considered a Create method like this:

BatchedExecutor b;
var conversation = b.Create();
conversation.Prompt("whatever");

i.e. remove the BatchedExecutor.Prompt(string) method entirely, and replace it with Create. Do you think that would be better?

edit: see related PR #613

Infer method throws an exception when internal batch is larger than the batch size from parameters

That's a good point. I can see two ways to fix that:

1. The Conversation.Prompt methods could check the batch size and refuse prompting if the batch size is reached. I think I'd do this by returning a count of how many tokens were consumed. This would be quite simple to implement, but introduces a bit of extra complexity to the caller.
2. The Conversation.Prompt methods could add to a per-conversation queue of tokens. When BatchedExecutor.Infer is called it pulls as many tokens as possible from conversations and runs that batch. This is a lot more complicated to implement, and introduces quite a lot of complexity into the systems.

Option 2 is quite a bit more complex, but I think it may actually solve another issue. With llava in a single inference call you can either prompt with an image or with text, not both. I was considering some kind of "mode" on the executor, so that as soon as you prompt a single conversation with an image, all attempts to prompt with text would throw until the next inference call (and vice versa). I haven't implemented that yet, because it's a horrible API! If instead the text/image prompts were simply buffered up until the next inference call that would solve the issue.

I think BatchExecutor is not thread-safe and would benefit a lot from some internal thread-safety

It wasn't designed with thread safety in mind, so that's probably true! Are you wanting to use it with one thread calling inference/sampling periodically, while other threads call Prompt?

It would be great to be able to pre-emptively handle NoKVSlot issues on the users of Conversation class side

Interesting idea. I'm not sure llama.cpp offers use the information needed to implement a KVSlotsLeft property, but I'll look into it.

compress conversation history and start a new one with smaller token count that continues the previous one

Have you looked into the ConversationExtensions at all? it's actually possible to re-arrange the KV cache in place, compressing it without having to re-read all of the token history (e.g. ShiftLeft). This system is intended to be extensible, so other KV cache manipulations can be implemented as custom extension methods (and maybe even upstreamed if they're useful).

[eublefar]

I don't think there's a way to do that at the moment, even at empty string will prompt it with a token.

I've tested empty string and BOS token initialization works for me, but create method also sounds good.

That's a good point. I can see two ways to fix that: ...

Maybe there is a 3rd option where BatchedExecutor has a queue of batches inside and instead of adding all to single batch it creates a new one on the overflow and adds it to the queue. Infer then would process single batch from the Queue (or could have some parameter e.g. processAll which would decide if Infer does one batch or all the Queue)

So like this would be some BatchedExecutor.AddToken method which would just check if batch size < params.batch_size and create a new batch on the queue otherwise.

Seems simpler than to handle it on Conversation side.

Does this make sense?

It wasn't designed with thread safety in mind, so that's probably true! Are you wanting to use it with one thread calling inference/sampling periodically, while other threads call Prompt?

I do exactly that: I have a separate thread (well rather async task on a threadpool) that runs Infer in a loop every few frames. Sampling and adding prompt happens in the Conversation wrapper on a thread pool as well. What I do now is that I just have a mutex that locks for BatchedExecutor.Infer, Conversation.Sample and Conversation.Prompt methods and it seems to work just fine.
This kind of approach feels natural to do with your API I think, where all the heavy lifting is offset to some worker thread and Conversations are just front-end objects.

it's actually possible to re-arrange the KV cache in place

What I had implented before was other two prompts that would summarize information from dialogue and then merge this summary with previous summaries, so I was thinking about how to convert my current solution into BatchedExecutor usage (specifically detecting context tokens left number is low). On a second thought now, smth like memory retrieval architecture + those operations on KV cache sound like a more optimized alternative to that so I guess I'll try it a bit later.

[eublefar]

Re: batch overflow. Had some free time so made a quick solution #619