Repository: ollama/ollama
File: runner/ollamarunner/runner.go (Go)

Code before:
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
						var reprocess *ErrReprocessInputs
						if errors.As(err, &reprocess) {
							// Prepend these inputs to the sequence's inputs queue for reprocessing
							seq.inputs = append(reprocess.Inputs, seq.inputs...)
							// Continue processing as normal
						} else {
							return err
						}
					}
				} else {
					break
				}
			}

			if i >= s.batchSize {
				break
			}

			// TODO(jessegross): Image inputs need to be rethought - it's
			// it doesn't work well for different types of models or multiple sequences
			if input.image != nil {
				if len(seq.pendingInputs) != len(options.Images) {
					break
				}

				if imgSeq != seqIdx && imgSeq != -1 {
					s.nextSeq = seqIdx
					break
				}

				imgSeq = seqIdx
				options.Images = append(options.Images, input.image)
				seq.pendingInputs = append(seq.pendingInputs, input)
				continue
			}

			options.Inputs = append(options.Inputs, input.token)
			options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
			options.Sequences = append(options.Sequences, seq.cache.Id)

			seq.iBatch = len(options.Outputs)
			if i+1 == len(seq.inputs) {
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
			seq.pendingInputs = append(seq.pendingInputs, input)
		}


Review comment (performance):
This loop is currently iterating over `seq.inputs` so I don't think that this will do the right thing, we will likely get some of the existing inputs processed before the ones we want to prepend.

We could just do `input := seq.inputs[i]` after the`if len(seq.cache.Inputs)+len(seq.pendingInputs)+1 > s.cache.numCtx {` block. We only shift if we have no pending inputs, so that makes things simpler.

Code after:
			if err != nil {
				panic(err)
			}
		}
	}
}

func (s *Server) processBatch() error {
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

	var batchInputs []int32
	var batch input.Batch

	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

		// if past the num predict limit
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
			s.removeSequence(i, "limit")
			continue
		}

		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
			seq.cache.Inputs = []input.Input{}
		}

		batchSize := s.batchSize

		for j, inp := range seq.inputs {
			// If we are required to put following inputs into a single batch then extend the
			// batch size. Since we are only extending the size the minimum amount possible, this
			// will cause a break if we have pending inputs.
			minBatch := 1 + inp.SameBatch
			if minBatch > batchSize {
				batchSize = minBatch
			}

			if len(seq.pendingInputs)+minBatch > batchSize {
				break
			}

			// If the sum of our working set (already processed tokens, tokens we added to this
			// batch, required following tokens) exceeds the context size, then trigger a shift
			// now so we don't have to do one later when we can't break the batch.