From 5d34320b7c43bfde9859096fd0e3c10a83149e49 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Tue, 13 Aug 2024 10:51:50 -0700 Subject: [PATCH] runner.go: Fix off by one in batch size check When adding tokens to a batch, the index is zero based but is checked against being greater than the max batch size. This results in an out-of-bounds access when the final token is added. --- llama/runner/runner.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llama/runner/runner.go b/llama/runner/runner.go index 9aca6cfe..c78f2ee8 100644 --- a/llama/runner/runner.go +++ b/llama/runner/runner.go @@ -178,7 +178,7 @@ func (s *Server) run(ctx context.Context) { for j, t := range seq.tokens { // todo: make this n_batch - if j > s.batchSize { + if j >= s.batchSize { break } batch.Add(t, seq.nPast, []int{i}, !seq.prompt())