소스 검색

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.
Jesse Gross 8 달 전
부모
커밋
8e1554c91d
1개의 변경된 파일1개의 추가작업 그리고 2개의 파일을 삭제
  1. 1 2
      llama/runner/runner.go

+ 1 - 2
llama/runner/runner.go

@@ -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