runner.go: Scale batches to be processed by numParallel

We should process a batch of tokens for each parallel request, rather
than having a shared pool. Otherwise, a single request can fill the
batch and then subsequent ones will fail or get starved.

Server.cpp used the KV cache size allocated for each parallel request
as the allocated size for the batch. This is the upper bound for the
batch but since we know how many tokens we will actually put in a batch
there is no need to over allocate.
This commit is contained in:
Jesse Gross 2024-08-23 13:44:30 -07:00 committed by jmorganca
parent f52d4b9879
commit 8e1554c91d

View File

@ -198,8 +198,7 @@ func incompleteUnicode(token string) bool {
}
func (s *Server) run(ctx context.Context) {
// TODO - should this be n_ctx / parallel like the old server.cpp setup?
batch := llama.NewBatch(s.batchSize, 0, s.parallel)
batch := llama.NewBatch(s.batchSize*len(s.seqs), 0, len(s.seqs))
defer batch.Free()
// build up stop sequences as we recognize them