runner.go 23.2 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package llamarunner
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
21
	"unicode/utf8"
22

23
24
	"golang.org/x/sync/semaphore"

25
26
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/llama"
27
	"github.com/ollama/ollama/llm"
Jesse Gross's avatar
Jesse Gross committed
28
	"github.com/ollama/ollama/runner/common"
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
)

// input is an element of the prompt to process, either
// a token or an image embedding (generated from a vision projector)
type input struct {
	token int

	// embed is an image embedding
	embed []float32
}

type Sequence struct {
	// batch index
	iBatch int

	// number of tokens predicted so far
	numPredicted int

	// prompt inputs left to evaluate
	inputs []input

50
51
52
	// inputs that have been added to a batch but not yet submitted to Decode
	pendingInputs []input

53
54
55
56
57
58
	// tokens that have been generated but not returned yet (e.g. for stop sequences)
	pendingResponses []string

	// input cache being used by this sequence
	cache *InputCacheSlot

59
60
61
62
	// does this sequence require cross-attention layers to be processed? - if we have seen
	// an image for certain multi-modal models
	crossAttention bool

63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
	// channel to send responses over
	responses chan string

	// channel to stop decoding (such as if the remote connection is closed)
	quit chan bool

	// number of tokens to predict
	numPredict int

	samplingCtx *llama.SamplingContext

	// channel to send back the embedding if embedding only
	embedding chan []float32

	// stop sequences
	stop []string

	// number of inputs to keep at the beginning when shifting context window
	numKeep int

	// true if an embedding are to be returned instead of text generation
	embeddingOnly bool

86
	doneReason llm.DoneReason
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

	// Metrics
	startProcessingTime time.Time
	startGenerationTime time.Time
	numDecoded          int
	numPromptInputs     int
}

type NewSequenceParams struct {
	numPredict     int
	stop           []string
	numKeep        int
	samplingParams *llama.SamplingParams
	embedding      bool
}

103
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
	s.ready.Wait()

	startTime := time.Now()

	inputs, err := s.inputs(prompt, images)
	if err != nil {
		return nil, fmt.Errorf("failed to process inputs: %w", err)
	} else if len(inputs) == 0 {
		return nil, errors.New("no input provided")
	}

	if params.numKeep < 0 {
		params.numKeep = len(inputs)
	}

119
120
	if s.model.AddBOSToken() {
		params.numKeep += 1
121
122
	}

123
124
125
	// Ensure that at least 1 input can be discarded during shift
	params.numKeep = min(params.numKeep, s.cache.numCtx-1)

126
	if len(inputs) > s.cache.numCtx {
127
		discard := len(inputs) - s.cache.numCtx
128
		newInputs := inputs[:params.numKeep]
129
130
131
		newInputs = append(newInputs, inputs[params.numKeep+discard:]...)

		slog.Warn("truncating input prompt", "limit", s.cache.numCtx, "prompt", len(inputs), "keep", params.numKeep, "new", len(newInputs))
132
		inputs = newInputs
133
134
135
136
	}

	var sc *llama.SamplingContext
	if params.samplingParams != nil {
Jesse Gross's avatar
Jesse Gross committed
137
138
139
140
		sc, err = llama.NewSamplingContext(s.model, *params.samplingParams)
		if err != nil {
			return nil, err
		}
141
142
		for _, input := range inputs {
			if input.embed == nil {
143
				sc.Accept(input.token, false)
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
			}
		}
	}

	return &Sequence{
		inputs:              inputs,
		numPromptInputs:     len(inputs),
		startProcessingTime: startTime,
		numPredict:          params.numPredict,
		pendingResponses:    make([]string, 0),
		responses:           make(chan string, 100),
		quit:                make(chan bool, 1),
		embedding:           make(chan []float32, 1),
		samplingCtx:         sc,
		embeddingOnly:       params.embedding,
		stop:                params.stop,
		numKeep:             params.numKeep,
	}, nil
}

// inputs processes the prompt and images into a list of inputs
// by splitting the prompt on [img-<n>] tags, tokenizing text and
// generating image embeddings for each image
167
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
168
	var inputs []input
169
170
171
172
173
174
175
176
177
178
	var parts []string
	var matches [][]string

	if s.image != nil {
		re := regexp.MustCompile(`\[img-(\d+)\]`)
		parts = re.Split(prompt, -1)
		matches = re.FindAllStringSubmatch(prompt, -1)
	} else {
		parts = []string{prompt}
	}
179
180
181

	for i, part := range parts {
		// text - tokenize
182
183
184
185
		tokens, err := s.lc.Model().Tokenize(part, i == 0, true)
		if err != nil {
			return nil, err
		}
186

187
188
		for _, t := range tokens {
			inputs = append(inputs, input{token: t})
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
		}

		// image - generate image embedding
		if i < len(matches) {
			n, _ := strconv.Atoi(matches[i][1])

			imageIndex := -1
			for j := range images {
				if images[j].ID == n {
					imageIndex = j
					break
				}
			}

			if imageIndex < 0 {
				return nil, fmt.Errorf("invalid image index: %d", n)
			}

Jesse Gross's avatar
Jesse Gross committed
207
208
209
210
211
			embed, err := s.image.NewEmbed(s.lc, images[imageIndex].Data, images[imageIndex].AspectRatioID)
			if err != nil {
				return nil, err
			}

212
213
214
215
216
217
218
219
220
221
			for _, e := range embed {
				inputs = append(inputs, input{embed: e})
			}
		}
	}

	return inputs, nil
}

type Server struct {
222
223
224
225
226
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
227
228
	model *llama.Model

229
	// image model context for multi-modal models
230
	image *ImageContext
231

232
	// status for external health reporting - loading, ready to serve, etc.
233
	status llm.ServerStatus
234
235
236
237
238
239
240
241

	// current progress on loading the model
	progress float32

	// number of simultaneous requests to handle
	parallel int

	// maximum number of elements in a batch (per sequence)
242
	// TODO (jmorganca): make this n_batch
243
244
	batchSize int

245
246
247
248
249
250
251
252
253
	// protects access to everything below this line
	// this is context state needed for decoding
	mu sync.Mutex

	// indicates that data is ready for processing
	cond *sync.Cond

	// decoding state
	lc *llama.Context
254

255
	// the list of simultaneous sequences being evaluated
256
257
	seqs []*Sequence

258
259
260
261
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
	// KV cache
	cache *InputCache

	// next sequence for prompt processing to avoid starvation
	nextSeq int
}

func (s *Server) allNil() bool {
	for _, item := range s.seqs {
		if item != nil {
			return false
		}
	}
	return true
}

func flushPending(seq *Sequence) bool {
279
280
281
282
283
284
285
286
287
288
289
	joined := strings.Join(seq.pendingResponses, "")
	seq.pendingResponses = []string{}

	// Check if there are any partial UTF-8 characters remaining.
	// We already check and queue as we are generating but some may
	// still make it here:
	// - Sequence is ending, e.g. generation limit has been hit
	// - Invalid characters in the middle of a string
	// This is a stricter check to ensure we never output invalid Unicode.
	for !utf8.ValidString(joined) {
		joined = joined[:len(joined)-1]
290
291
	}

292
293
294
295
296
297
298
299
300
301
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
302
303
}

304
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
305
306
307
308
309
310
311
312
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
313
	s.seqsSem.Release(1)
314
315
316
317
318
}

func (s *Server) run(ctx context.Context) {
	s.ready.Wait()

319
	// Logically these batches are used only within the context of processBatch
320
	// but it is better for performance to allocate them once here
Jesse Gross's avatar
Jesse Gross committed
321
322
323
324
	tokenBatch, err := llama.NewBatch(s.batchSize, len(s.seqs), 0)
	if err != nil {
		panic(err)
	}
325
326
	defer tokenBatch.Free()

327
328
329
	var embedBatch *llama.Batch
	embedBatchSize := s.image.BatchSize(s.batchSize)
	if embedBatchSize != 0 {
Jesse Gross's avatar
Jesse Gross committed
330
331
332
333
		embedBatch, err = llama.NewBatch(embedBatchSize, len(s.seqs), s.image.EmbedSize(s.lc))
		if err != nil {
			panic(err)
		}
334
335
336
337
		defer embedBatch.Free()
	} else {
		embedBatch = &llama.Batch{}
	}
338
339
340
341
342
343

	for {
		select {
		case <-ctx.Done():
			return
		default:
344
345
346
347
348
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

349
350
351
352
353
354
355
356
357
358
359
360
361
			tokenBatch.Clear()
			embedBatch.Clear()
		}
	}
}

// TODO (jmorganca): processBatch should be simplified, removing:
// * sampling
// * stop token checking
// * metrics
// these should instead be handled by the handlers
// it should only be responsible for accepting tokens or embeddings and
// processing batches as fast as possible
362
func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch) error {
363
364
365
366
367
368
369
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

	var batch *llama.Batch
370
	crossAttention := false
371
372
373
374
375
376
377
378
379
380
381

	seqIdx := s.nextSeq - 1
	for range s.seqs {
		seqIdx = (seqIdx + 1) % len(s.seqs)
		seq := s.seqs[seqIdx]

		if seq == nil {
			continue
		}

		// if past the num predict limit
382
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
383
			s.removeSequence(seqIdx, llm.DoneReasonLength)
384
385
386
387
			continue
		}

		for i, input := range seq.inputs {
388
389
390
391
			if 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 {
392
393
394
395
396
397
398
399
400
						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
							continue
						} else {
							return err
						}
401
					}
402
403
404
405
406
				} else {
					break
				}
			}

407
408
409
410
411
412
413
414
415
416
417
			embedding := input.embed != nil

			// If we don't currently have a batch, use one of the correct type and
			// fill it up as much as possible across all sequences. If we encounter an
			// input of the opppsite type, stop for that sequence but then pick up from
			// there for the next batch, ensuring that we alternate types
			if batch == nil {
				if !embedding {
					batch = tokenBatch
				} else {
					batch = embedBatch
418
					seq.crossAttention = s.image.NeedCrossAttention(input)
419
				}
420
			} else if embedding != batch.IsEmbedding() || crossAttention != seq.crossAttention {
421
422
423
424
				s.nextSeq = seqIdx
				break
			}

425
			if i >= batch.Size() {
426
427
428
				break
			}

429
			crossAttention = seq.crossAttention
430
431
			batch.Add(input.token, input.embed, len(seq.cache.Inputs)+len(seq.pendingInputs), i+1 == len(seq.inputs), seq.cache.Id)
			seq.pendingInputs = append(seq.pendingInputs, input)
432
433
			seq.iBatch = batch.NumTokens() - 1
		}
434
435

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
436
437
438
	}

	if batch == nil || batch.NumTokens() == 0 {
439
		return nil
440
441
	}

442
443
	s.lc.SetCrossAttention(crossAttention)

444
445
	err := s.lc.Decode(batch)
	if err != nil {
446
		return fmt.Errorf("failed to decode batch: %w", err)
447
448
	}

449
450
451
452
453
454
455
	if crossAttention {
		// synchronize state to ensure the cross attention batch is complete.
		// needed specifically for multi-GPU systems otherwise an inflight
		// task may be incorrectly invalidated causing a crash
		s.lc.Synchronize()
	}

456
457
458
459
460
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

461
462
463
464
465
466
		// After calling Decode, pending inputs are now in the cache
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
			seq.pendingInputs = []input{}
		}

467
468
469
470
471
472
473
474
475
476
477
478
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
			continue
		}

		seq.numDecoded += 1
		if seq.numDecoded == 1 {
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
479
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
480
481
482
483
484
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
485
			s.removeSequence(i, llm.DoneReasonStop)
486
487
488
489
			continue
		}

		// sample a token
490
491
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
492
493
494
495
496
497
498
499
500
501
		piece := s.model.TokenToPiece(token)

		seq.numPredicted++

		// if it's an end of sequence token, break
		if s.model.TokenIsEog(token) {
			// TODO (jmorganca): we should send this back
			// as it's important for the /api/generate context
			// seq.responses <- piece

502
			s.removeSequence(i, llm.DoneReasonStop)
503
504
505
506
507
508
509
510
			continue
		}

		seq.inputs = []input{{token: token}}

		seq.pendingResponses = append(seq.pendingResponses, piece)
		sequence := strings.Join(seq.pendingResponses, "")

Jesse Gross's avatar
Jesse Gross committed
511
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
512
513
514
515
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
516
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
			newLen := len(seq.pendingResponses)

			// Update the cache based on the tokens that will be returned:
			// - We have 1 token more than is currently in the cache because
			// the last one generated wasn't submitted to Decode
			// - Remove any stop sequences that we stripped out
			// - If truncateStop removed a portion of a token, drop that
			// - As defense-in-depth, if truncatedToken didn't find a stop token
			// remove the extra one that we added to the cache len
			tokenLen := len(seq.cache.Inputs) + 1
			tokenLen -= origLen - newLen
			if tokenTruncated || origLen == newLen {
				tokenLen--
			}
			seq.cache.Inputs = seq.cache.Inputs[:tokenLen]
532

533
			s.removeSequence(i, llm.DoneReasonStop)
534
535
536
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
537
		if common.ContainsStopSuffix(sequence, seq.stop) {
538
539
540
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
541
		if common.IncompleteUnicode(sequence) {
542
543
544
545
			continue
		}

		if !flushPending(seq) {
546
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
547
548
		}
	}
549
550

	return nil
551
552
553
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
554
	var req llm.CompletionRequest
555
556
557
558
559
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

560
561
562
563
564
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

565
566
567
568
569
570
571
572
573
574
	// Set the headers to indicate streaming
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Transfer-Encoding", "chunked")

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", http.StatusInternalServerError)
		return
	}

575
576
577
578
579
580
581
582
583
584
585
586
587
588
	// Extract options from the CompletionRequest
	samplingParams := llama.SamplingParams{
		TopK:           req.Options.TopK,
		TopP:           req.Options.TopP,
		MinP:           req.Options.MinP,
		TypicalP:       req.Options.TypicalP,
		Temp:           req.Options.Temperature,
		RepeatLastN:    req.Options.RepeatLastN,
		PenaltyRepeat:  req.Options.RepeatPenalty,
		PenaltyFreq:    req.Options.FrequencyPenalty,
		PenaltyPresent: req.Options.PresencePenalty,
		Seed:           uint32(req.Options.Seed),
		Grammar:        req.Grammar,
	}
589
590

	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
591
592
593
		numPredict:     req.Options.NumPredict,
		stop:           req.Options.Stop,
		numKeep:        req.Options.NumKeep,
594
595
596
597
598
599
600
601
		samplingParams: &samplingParams,
		embedding:      false,
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

602
	// Ensure there is a place to put the sequence, released when removed from s.seqs
603
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
604
605
606
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
607
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
608
		}
609
610
611
		return
	}

612
	s.mu.Lock()
613
	found := false
614
615
	for i, sq := range s.seqs {
		if sq == nil {
616
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
617
618
			if err != nil {
				s.mu.Unlock()
619
				s.seqsSem.Release(1)
620
621
622
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
623

624
625
			seq.crossAttention = s.image.NeedCrossAttention(seq.cache.Inputs...)

626
627
			s.seqs[i] = seq
			s.cond.Signal()
628
			found = true
629
630
631
632
633
			break
		}
	}
	s.mu.Unlock()

634
	if !found {
635
		s.seqsSem.Release(1)
636
637
638
639
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

640
641
642
643
644
645
646
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
647
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
648
649
650
651
652
653
654
655
656
					Content: content,
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
657
658
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
659
					DoneReason:         seq.doneReason,
660
661
662
663
					PromptEvalCount:    seq.numPromptInputs,
					PromptEvalDuration: seq.startGenerationTime.Sub(seq.startProcessingTime),
					EvalCount:          seq.numDecoded,
					EvalDuration:       time.Since(seq.startGenerationTime),
664
665
666
667
668
669
670
671
672
673
674
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
				}

				return
			}
		}
	}
}

func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
675
	var req llm.EmbeddingRequest
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, fmt.Sprintf("bad request: %s", err), http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/json")

	slog.Debug("embedding request", "content", req.Content)

	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{embedding: true})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

691
	// Ensure there is a place to put the sequence, released when removed from s.seqs
692
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
693
694
695
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
696
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
697
		}
698
699
700
		return
	}

701
	s.mu.Lock()
702
	found := false
703
704
	for i, sq := range s.seqs {
		if sq == nil {
705
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
706
707
			if err != nil {
				s.mu.Unlock()
708
				s.seqsSem.Release(1)
709
710
711
712
713
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
			s.seqs[i] = seq
			s.cond.Signal()
714
			found = true
715
716
717
718
719
			break
		}
	}
	s.mu.Unlock()

720
	if !found {
721
		s.seqsSem.Release(1)
722
723
724
725
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

726
727
	embedding := <-seq.embedding

728
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
729
730
731
732
733
734
735
736
		Embedding: embedding,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
737
738
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
739
740
741
742
743
744
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

745
746
747
748
749
750
751
752
753
754
755
type multiLPath []string

func (m *multiLPath) Set(value string) error {
	*m = append(*m, value)
	return nil
}

func (m *multiLPath) String() string {
	return strings.Join(*m, ", ")
}

756
757
758
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
759
	lpath multiLPath,
760
761
	ppath string,
	kvSize int,
762
	kvCacheType string,
763
764
765
766
	flashAttention bool,
	threads int,
	multiUserCache bool,
) {
767
768
769
770
771
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
772

773
	ctxParams := llama.NewContextParams(kvSize, s.batchSize*s.parallel, s.parallel, threads, flashAttention, kvCacheType)
774
775
776
777
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
778

779
780
781
782
783
784
	if lpath.String() != "" {
		for _, path := range lpath {
			err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
			if err != nil {
				panic(err)
			}
785
786
787
788
		}
	}

	if ppath != "" {
789
		var err error
790
		s.image, err = NewImageContext(s.lc, ppath)
791
792
793
		if err != nil {
			panic(err)
		}
794
795
	}

796
797
798
799
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
800

801
	s.status = llm.ServerStatusReady
802
803
804
	s.ready.Done()
}

805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
func Execute(args []string) error {
	fs := flag.NewFlagSet("runner", flag.ExitOnError)
	mpath := fs.String("model", "", "Path to model binary file")
	ppath := fs.String("mmproj", "", "Path to projector binary file")
	parallel := fs.Int("parallel", 1, "Number of sequences to handle simultaneously")
	batchSize := fs.Int("batch-size", 512, "Batch size")
	nGpuLayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGpu := fs.Int("main-gpu", 0, "Main GPU")
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
	kvSize := fs.Int("ctx-size", 2048, "Context (or KV cache) size")
	kvCacheType := fs.String("kv-cache-type", "", "quantization type for KV cache (default: f16)")
	port := fs.Int("port", 8080, "Port to expose the server on")
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
	noMmap := fs.Bool("no-mmap", false, "do not memory-map model (slower load but may reduce pageouts if not using mlock)")
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
822

823
	var lpaths multiLPath
824
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
825

826
827
828
829
830
831
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
	}
	level := slog.LevelInfo
	if *verbose {
		level = slog.LevelDebug
	}
	handler := slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level:     level,
		AddSource: true,
		ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
			if attr.Key == slog.SourceKey {
				source := attr.Value.Any().(*slog.Source)
				source.File = filepath.Base(source.File)
			}
			return attr
		},
	})
	slog.SetDefault(slog.New(handler))
	slog.Info("starting go runner")
850
851

	llama.BackendInit()
852
853
854
855
856

	server := &Server{
		batchSize: *batchSize,
		parallel:  *parallel,
		seqs:      make([]*Sequence, *parallel),
857
		seqsSem:   semaphore.NewWeighted(int64(*parallel)),
858
		status:    llm.ServerStatusLoadingModel,
859
860
861
862
	}

	var tensorSplitFloats []float32
	if *tensorSplit != "" {
863
864
865
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
866
			f, _ := strconv.ParseFloat(s, 32)
867
			tensorSplitFloats[i] = float32(f)
868
869
870
871
872
873
		}
	}

	params := llama.ModelParams{
		NumGpuLayers: *nGpuLayers,
		MainGpu:      *mainGpu,
874
		UseMmap:      !*noMmap && lpaths.String() == "",
875
876
877
878
879
880
881
		TensorSplit:  tensorSplitFloats,
		Progress: func(progress float32) {
			server.progress = progress
		},
	}

	server.ready.Add(1)
882
	go server.loadModel(params, *mpath, lpaths, *ppath, *kvSize, *kvCacheType, *flashAttention, *threads, *multiUserCache)
883
884
885
886

	server.cond = sync.NewCond(&server.mu)

	ctx, cancel := context.WithCancel(context.Background())
Michael Yang's avatar
Michael Yang committed
887
888
	defer cancel()

889
890
891
892
893
894
	go server.run(ctx)

	addr := "127.0.0.1:" + strconv.Itoa(*port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Println("Listen error:", err)
895
		return err
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
	}
	defer listener.Close()

	mux := http.NewServeMux()
	mux.HandleFunc("/embedding", server.embeddings)
	mux.HandleFunc("/completion", server.completion)
	mux.HandleFunc("/health", server.health)

	httpServer := http.Server{
		Handler: mux,
	}

	log.Println("Server listening on", addr)
	if err := httpServer.Serve(listener); err != nil {
		log.Fatal("server error:", err)
911
		return err
912
913
	}

914
	return nil
915
}