Repository navigation
Feature request: batch creation of embeddings #1577
Description
The program does not currently allow for true batch processing when creating embeddings, which would greatly speed up creating embeddings. Upstream allows for batch processing but Kobold doesn't implement it. Here's an outline that's hopefully helpful:
- llama.h allows the batching of multiple sequences of tokens.
- common.cpp adds each token of a prompt to the "llama_batch" and associates each token with a sequence id
- llama_encode is used, see more here context : allow cache-less context for embeddings ggml-org/llama.cpp#13108
- embeddings are retrieved using "llama_get_embeddings_seq" or "llama_get_embeddings_ith"
- "n_batch" within llama.cpp refers to the maximum number of tokens that can be sent to "llama_batch". It is not traditional. For example, with pytorch and sentence transformers each "batch" consists of multiple "sequences" (i.e. chunks of text to be embedded) and each one must not exceed the context limit of a particular embedding model. Further each "batch" is padded to properly create a "tensor." However, with llama.cpp, each "batch" (i.e. a sequence/chunk of text to be embedded) becomes part of a flat token stream and different sequences are identified with a "seq_id" and, therefore, padding is unnecessary. However, this means that all sequences must be within "n_batch" and, more importantly, "n_batch" must never exceed an embedding model's maximum context limit. Again, this is different than pytorch and sentence transformers.
- Historically, embedding models had a maximum context limit of 512 (and many still do), which made true batch processing with llama.cpp unattractive. However, many embedding models now have limits of 8192 or more...
Batch processing would greatly improve embedding large documents.
Within Kobold, embeddingstype_generate() processes a single prompt although it's structured allow future expansion.
Further, koboldcpp.py loops over user inputs and only passes one prompt to the native embeddings_generate() C++ function.
for prompt in prompts:
inputs = embeddings_generation_inputs()
inputs.prompt = prompt.encode()
ret = handle.embeddings_generate(inputs)It then aggregates results into an openai-style response...
Basically, multiple prompts would be processed one at a time at the python layer, and internally "embeddings_generate" only receives and batches one prompt at a time.
Activity
[-]Feature request:[/-][+]Feature request: batch creation of embeddings[/+]LostRuins commentedon Jun 1, 2025
I don't think you'll get much speedup with this, because most of the time you are already occupying a large portion of the available context (which is loaded equal to the
n_ctx_trainand maximum embedding size. It will only be helpful if sending a large number of small texts to embed, which would be fast enough in either case anyway. Moreover, doing more than 1 batch at once would require correspondingly more memory too.Are your embeddings taking a lot of time? I generated a few hundred within a minute, so it didn't really seem like an issue to me, so long as you are using a GPU.
BBC-Esq commentedon Jun 1, 2025
I understand your point regarding a large number of small texts. Typically, assuming an average sequence length of 250 tokens, that's only two sequences per "batch" as llama.cpp understands that term. However, the major speed up would occur with most new embedding models that support 8192 token limits, 32k, and even more now. Currently, the highest ranking models are all like this now.
But yes, there would be more memory and I noticed a portion of the source code outlining how llama.cpp allocates it...However, I didn't get a chance to do an estimated calculation. So, yes, there should be major speedups with models that have large context windows...but the memory issue remains to be seen.
Which embedding model, approximate chunk size, and hardware did you use so I can get an approximation and do a comparison on my end if you don't mind?
LostRuins commentedon Jun 1, 2025
I tested with snowflake-arctic-embed, various sizes. This is on an rtx 4090
BBC-Esq commentedon Jun 1, 2025
Thanks, was it the laptop version I think you bought a laptop or something?
LostRuins commentedon Jun 1, 2025
Yes, it's the laptop version, 16GB VRAM
BBC-Esq commentedon Jun 3, 2025
chunks.zip
Can you put a few benchmarks down so we can get an idea of the speed? I understand that our hardware is different, but it's close enough to get an idea of whether I need to spend X hour testing this to the Nth degree.
Setup
n_batchto 512.Torch+sentence-transformers:
Batch size 8 - 13.52 seconds (slight regression, which is why I set the batch to 7 in my program)
Batch size 4 = 13.14 seconds ~299 chunks/second
Batch size 2 = 17.54 seconds ~224 chunks/second
Batch size 1 = 29.58 seconds ~133 chunks/second
Seems faster than the "few hundred per minute" that you mentioned...but again, this is ballparking since we have slightly different hardware. Anyways, thought this was also interesting
LostRuins commentedon Jun 3, 2025
But either way, i believe all the above speeds are more than fast enough for most use cases?
BBC-Esq commentedon Jun 3, 2025
Are you willing to benchmark the chunks I just gave? I'd still like to delve into the speed issue if not the memory issue.
LostRuins commentedon Jun 3, 2025
Sure. Do you have a benchmarking script I can run?
BBC-Esq commentedon Jun 3, 2025
Unfortunately I don't. I thought you did since you said you'd been testing them? All you'd have to do is import the chunks from the .json lines file and send to the endpoint. If you send me the script you were using I can take a look to make sure it's done correctly and the time keeping is right?
LostRuins commentedon Jun 3, 2025
I've been testing it directly in Lite. Using the TextDB embedding feature.
BBC-Esq commentedon Jun 3, 2025
Ah, well I'll meet you halfway. I've done the heavy lifting of extracting/splitting and testing with hf+ST if you'll get the chunks and test them in Kobold somehow.
ceruleandeep commentedon Jun 26, 2025
I tried running those chunks. I get "hundreds of generations a minute" from KoboldCpp and "hundreds of generations per second" from vllm, which are kinda the numbers that both of you got? These are big differences. I'm probably 'holding it wrong' somehow, but I couldn't work out how.
4090 desktop
I did try configuring llama-server for batched inferencing, I couldn't find a configuration that went any faster.
LostRuins commentedon Jun 26, 2025
Yeah this is probably as fast as it will go in kcpp for the large model. So in this case vllm will be significantly faster. This is in reqs per second tho, each request is probably 100s of tokens, so it is doing thousands to t/s
so this is correct and similar to what i got, If you use the small model together with
--embeddingsgputhat should be a bunch higher.