fix: fixes a memory leak in bfloat16 package

This change vendors in the bfloat16 package from
github.com/d4l3k/go-bfloat16/ and fixes a memory leak which
was being caused by using unsafe pointers instead of the
math package.

convert/convert.go

@@ -201,7 +201,7 @@ func ConvertModel(fsys fs.FS, ws io.WriteSeeker) error {
 	case "CohereForCausalLM":
 		conv = &commandrModel{}
 	default:
-		return fmt.Errorf("unsupported architecture %q", p.Architectures[0])
+		return errors.New("unsupported architecture")
 	}
 
 	if err := json.Unmarshal(bts, conv); err != nil {

convert/reader_safetensors.go

@@ -11,9 +11,10 @@ import (
 	"slices"
 	"strings"
 
-	"github.com/d4l3k/go-bfloat16"
 	"github.com/x448/float16"
 	"golang.org/x/exp/maps"
+
+	"github.com/ollama/ollama/types/bfloat16"
 )
 
 type safetensorMetadata struct {

go.mod

@@ -16,7 +16,6 @@ require (
 
 require (
 	github.com/agnivade/levenshtein v1.1.1
-	github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1
 	github.com/dlclark/regexp2 v1.11.4
 	github.com/emirpasic/gods/v2 v2.0.0-alpha
 	github.com/google/go-cmp v0.6.0

go.sum

@@ -35,8 +35,6 @@ github.com/containerd/console v1.0.3 h1:lIr7SlA5PxZyMV30bDW0MGbiOPXwc63yRuCP0ARu
 github.com/containerd/console v1.0.3/go.mod h1:7LqA/THxQ86k76b8c/EMSiaJ3h1eZkMkXar0TQ1gf3U=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
-github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1 h1:cBzrdJPAFBsgCrDPnZxlp1dF2+k4r1kVpD7+1S1PVjY=
-github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1/go.mod h1:uw2gLcxEuYUlAd/EXyjc/v55nd3+47YAgWbSXVxPrNI=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=

ml/backend/ggml/ggml.go

@@ -330,7 +330,7 @@ func New(r *os.File, params ml.BackendParams) (ml.Backend, error) {
 		}
 	}
 
-	if err := g.Wait(); err != nil {
+	if g.Wait() != nil {
 		return nil, err
 	}
 

model/models/gemma2/model.go

@@ -179,7 +179,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

model/models/gemma3/model.go

@@ -150,7 +150,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

model/models/llama/model.go

@@ -150,7 +150,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

model/models/mllama/model.go

@@ -154,7 +154,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

runner/common/stop.go

@@ -2,8 +2,6 @@ package common
 
 import (
 	"strings"
-
-	"github.com/ollama/ollama/llm"
 )
 
 func FindStop(sequence string, stops []string) (bool, string) {
@@ -31,43 +29,40 @@ func ContainsStopSuffix(sequence string, stops []string) bool {
 // truncateStop removes the provided stop string from pieces,
 // returning the partial pieces with stop removed, including truncating
 // the last piece if required (and signalling if this was the case)
-func TruncateStop(resps []llm.CompletionResponse, stop string) ([]llm.CompletionResponse, bool) {
-	var sequence string
-	for _, resp := range resps {
-		sequence += resp.Content
+func TruncateStop(pieces []string, stop string) ([]string, bool) {
+	joined := strings.Join(pieces, "")
+
+	index := strings.Index(joined, stop)
+	if index == -1 {
+		return pieces, false
 	}
 
-	idx := strings.Index(sequence, stop)
-	if idx < 0 {
-		return resps, false
+	joined = joined[:index]
+
+	// Split truncated string back into pieces of original lengths
+	lengths := make([]int, len(pieces))
+	for i, piece := range pieces {
+		lengths[i] = len(piece)
 	}
 
-	truncated := sequence[:idx]
-	if len(truncated) == 0 {
-		return nil, true
-	}
-
-	result := make([]llm.CompletionResponse, 0, len(resps))
-
-	// Track position in truncated sequence
-	pos := 0
-	truncationHappened := false
-	for _, resp := range resps {
-		if pos >= len(truncated) {
+	var result []string
+	tokenTruncated := false
+	start := 0
+	for _, length := range lengths {
+		if start >= len(joined) {
 			break
 		}
 
-		chunk := truncated[pos:min(pos+len(resp.Content), len(truncated))]
-		if len(chunk) < len(resp.Content) {
-			truncationHappened = true
+		end := start + length
+		if end > len(joined) {
+			end = len(joined)
+			tokenTruncated = true
 		}
-		if len(chunk) > 0 {
-			result = append(result, llm.CompletionResponse{Content: chunk})
-		}
-		pos += len(resp.Content)
+		result = append(result, joined[start:end])
+		start = end
 	}
 
-	return result, truncationHappened
+	return result, tokenTruncated
 }
 
 func IncompleteUnicode(token string) bool {

runner/common/stop_test.go

@@ -1,84 +1,51 @@
 package common
 
 import (
-	"fmt"
 	"reflect"
 	"testing"
-
-	"github.com/ollama/ollama/llm"
 )
 
 func TestTruncateStop(t *testing.T) {
 	tests := []struct {
 		name          string
-		pieces        []llm.CompletionResponse
+		pieces        []string
 		stop          string
-		expected      []llm.CompletionResponse
+		expected      []string
 		expectedTrunc bool
 	}{
 		{
-			name: "Single word",
-			pieces: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: "world"},
-			},
-			stop: "world",
-			expected: []llm.CompletionResponse{
-				{Content: "Hello"},
-			},
+			name:          "Single word",
+			pieces:        []string{"hello", "world"},
+			stop:          "world",
+			expected:      []string{"hello"},
 			expectedTrunc: false,
 		},
 		{
-			name: "Partial",
-			pieces: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " wor"},
-			},
-			stop: "or",
-			expected: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " w"},
-			},
+			name:          "Partial",
+			pieces:        []string{"hello", "wor"},
+			stop:          "or",
+			expected:      []string{"hello", "w"},
 			expectedTrunc: true,
 		},
 		{
-			name: "Suffix",
-			pieces: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " there"},
-				{Content: "!"},
-			},
-			stop: "!",
-			expected: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " there"},
-			},
+			name:          "Suffix",
+			pieces:        []string{"Hello", " there", "!"},
+			stop:          "!",
+			expected:      []string{"Hello", " there"},
 			expectedTrunc: false,
 		},
 		{
-			name: "Suffix partial",
-			pieces: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " the"},
-				{Content: "re!"},
-			},
-			stop: "there!",
-			expected: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " "},
-			},
+			name:          "Suffix partial",
+			pieces:        []string{"Hello", " the", "re!"},
+			stop:          "there!",
+			expected:      []string{"Hello", " "},
 			expectedTrunc: true,
 		},
 		{
-			name: "Middle",
-			pieces: []llm.CompletionResponse{
-				{Content: "Hello"},
-				{Content: " wo"},
-			},
-			stop: "llo w",
-			expected: []llm.CompletionResponse{
-				{Content: "He"},
-			},
+			name:          "Middle",
+			pieces:        []string{"hello", " wor"},
+			stop:          "llo w",
+			expected:      []string{"he"},
 			expectedTrunc: true,
 		},
 	}
@@ -87,27 +54,12 @@ func TestTruncateStop(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			result, resultTrunc := TruncateStop(tt.pieces, tt.stop)
 			if !reflect.DeepEqual(result, tt.expected) || resultTrunc != tt.expectedTrunc {
-				t.Errorf("truncateStop(%v, %v):\n%shave truncated %v\nwant truncated %v",
-					tt.pieces, tt.stop, formatContentDiff(result, tt.expected), resultTrunc, tt.expectedTrunc)
+				t.Errorf("truncateStop(%v, %s): have %v (%v); want %v (%v)", tt.pieces, tt.stop, result, resultTrunc, tt.expected, tt.expectedTrunc)
 			}
 		})
 	}
 }
 
-func formatContentDiff(result, expected []llm.CompletionResponse) string {
-	var s string
-	for i := 0; i < len(result) || i < len(expected); i++ {
-		if i < len(result) && i < len(expected) && result[i].Content != expected[i].Content {
-			s += fmt.Sprintf("[%d] %q vs %q\n", i, result[i].Content, expected[i].Content)
-		} else if i < len(result) && i >= len(expected) {
-			s += fmt.Sprintf("[%d] extra %q\n", i, result[i].Content)
-		} else if i >= len(result) && i < len(expected) {
-			s += fmt.Sprintf("[%d] missing %q\n", i, expected[i].Content)
-		}
-	}
-	return s
-}
-
 func TestIncompleteUnicode(t *testing.T) {
 	tests := []struct {
 		name     string

runner/llamarunner/runner.go

@@ -51,7 +51,7 @@ type Sequence struct {
 	pendingInputs []input
 
 	// tokens that have been generated but not returned yet (e.g. for stop sequences)
-	pendingResponses []llm.CompletionResponse
+	pendingResponses []string
 
 	// input cache being used by this sequence
 	cache *InputCacheSlot
@@ -61,7 +61,7 @@ type Sequence struct {
 	crossAttention bool
 
 	// channel to send responses over
-	responses chan llm.CompletionResponse
+	responses chan string
 
 	// channel to stop decoding (such as if the remote connection is closed)
 	quit chan bool
@@ -150,8 +150,8 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 		numPromptInputs:     len(inputs),
 		startProcessingTime: startTime,
 		numPredict:          params.numPredict,
-		pendingResponses:    make([]llm.CompletionResponse, 0),
-		responses:           make(chan llm.CompletionResponse, 100),
+		pendingResponses:    make([]string, 0),
+		responses:           make(chan string, 100),
 		quit:                make(chan bool, 1),
 		embedding:           make(chan []float32, 1),
 		samplingCtx:         sc,
@@ -276,28 +276,29 @@ func (s *Server) allNil() bool {
 }
 
 func flushPending(seq *Sequence) bool {
-	pending := seq.pendingResponses
-	seq.pendingResponses = []llm.CompletionResponse{}
+	joined := strings.Join(seq.pendingResponses, "")
+	seq.pendingResponses = []string{}
 
-	for i, r := range pending {
-		if i == len(pending)-1 {
-			// Check and trim any trailing partial UTF-8 characters
-			content := r.Content
-			for !utf8.ValidString(content) {
-				content = content[:len(content)-1]
-			}
-			r.Content = content
-		}
+	// 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]
+	}
 
-		select {
-		case seq.responses <- r:
-			return true
-		case <-seq.quit:
-			return false
-		}
+	if len(joined) == 0 {
+		return true
+	}
+
+	select {
+	case seq.responses <- joined:
+		return true
+	case <-seq.quit:
+		return false
 	}
-	// no pending responses to send
-	return true
 }
 
 func (s *Server) removeSequence(seqIndex int, reason string) {
@@ -496,11 +497,8 @@ func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch)
 
 		seq.inputs = []input{{token: token}}
 
-		seq.pendingResponses = append(seq.pendingResponses, llm.CompletionResponse{Content: piece})
-		sequence := ""
-		for _, r := range seq.pendingResponses {
-			sequence += r.Content
-		}
+		seq.pendingResponses = append(seq.pendingResponses, piece)
+		sequence := strings.Join(seq.pendingResponses, "")
 
 		if ok, stop := common.FindStop(sequence, seq.stop); ok {
 			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)
@@ -639,7 +637,9 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 			return
 		case content, ok := <-seq.responses:
 			if ok {
-				if err := json.NewEncoder(w).Encode(&content); err != nil {
+				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
+					Content: content,
+				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
 					close(seq.quit)
 					return

runner/ollamarunner/runner.go

@@ -53,13 +53,13 @@ type Sequence struct {
 	pendingInputs []input.Input
 
 	// tokens that have been generated but not returned yet (e.g. for stop sequences)
-	pendingResponses []llm.CompletionResponse
+	pendingResponses []string
 
 	// input cache being used by this sequence
 	cache *InputCacheSlot
 
 	// channel to send responses over
-	responses chan llm.CompletionResponse
+	responses chan string
 
 	// channel to stop decoding (such as if the remote connection is closed)
 	quit chan bool
@@ -138,8 +138,8 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 		numPromptInputs:     len(inputs),
 		startProcessingTime: startTime,
 		numPredict:          params.numPredict,
-		pendingResponses:    make([]llm.CompletionResponse, 0),
-		responses:           make(chan llm.CompletionResponse, 100),
+		pendingResponses:    make([]string, 0),
+		responses:           make(chan string, 100),
 		quit:                make(chan bool, 1),
 		embedding:           make(chan []float32, 1),
 		sampler:             params.sampler,
@@ -288,28 +288,29 @@ func (s *Server) allNil() bool {
 }
 
 func flushPending(seq *Sequence) bool {
-	pending := seq.pendingResponses
-	seq.pendingResponses = []llm.CompletionResponse{}
+	joined := strings.Join(seq.pendingResponses, "")
+	seq.pendingResponses = []string{}
 
-	for i, r := range pending {
-		if i == len(pending)-1 {
-			// Check and trim any trailing partial UTF-8 characters
-			content := r.Content
-			for !utf8.ValidString(content) {
-				content = content[:len(content)-1]
-			}
-			r.Content = content
-		}
+	// 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]
+	}
 
-		select {
-		case seq.responses <- r:
-			return true
-		case <-seq.quit:
-			return false
-		}
+	if len(joined) == 0 {
+		return true
+	}
+
+	select {
+	case seq.responses <- joined:
+		return true
+	case <-seq.quit:
+		return false
 	}
-	// no pending responses to send
-	return true
 }
 
 func (s *Server) removeSequence(seqIndex int, reason string) {
@@ -483,11 +484,8 @@ func (s *Server) processBatch() error {
 
 		seq.inputs = []input.Input{{Token: token}}
 
-		seq.pendingResponses = append(seq.pendingResponses, llm.CompletionResponse{Content: piece})
-		sequence := ""
-		for _, r := range seq.pendingResponses {
-			sequence += r.Content
-		}
+		seq.pendingResponses = append(seq.pendingResponses, piece)
+		sequence := strings.Join(seq.pendingResponses, "")
 
 		if ok, stop := common.FindStop(sequence, seq.stop); ok {
 			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)
@@ -625,7 +623,9 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 			return
 		case content, ok := <-seq.responses:
 			if ok {
-				if err := json.NewEncoder(w).Encode(&content); err != nil {
+				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
+					Content: content,
+				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
 					close(seq.quit)
 					return

types/bfloat16/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Tristan Rice
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

types/bfloat16/bfloat16.go

@@ -0,0 +1,57 @@
+// Vendored code from https://github.com/d4l3k/go-bfloat16
+// unsafe pointer replaced by "math"
+package bfloat16
+
+import "math"
+
+type BF16 uint16
+
+func FromBytes(buf []byte) BF16 {
+	return BF16(uint16(buf[0]) + uint16(buf[1])<<8)
+}
+
+func ToBytes(b BF16) []byte {
+	return []byte{byte(b & 0xFF), byte(b >> 8)}
+}
+
+func Decode(buf []byte) []BF16 {
+	var out []BF16
+	for i := 0; i < len(buf); i += 2 {
+		out = append(out, FromBytes(buf[i:]))
+	}
+	return out
+}
+
+func Encode(f []BF16) []byte {
+	var out []byte
+	for _, a := range f {
+		out = append(out, ToBytes(a)...)
+	}
+	return out
+}
+
+func DecodeFloat32(buf []byte) []float32 {
+	var out []float32
+	for i := 0; i < len(buf); i += 2 {
+		out = append(out, ToFloat32(FromBytes(buf[i:])))
+	}
+	return out
+}
+
+func EncodeFloat32(f []float32) []byte {
+	var out []byte
+	for _, a := range f {
+		out = append(out, ToBytes(FromFloat32(a))...)
+	}
+	return out
+}
+
+func ToFloat32(b BF16) float32 {
+	u32 := uint32(b) << 16
+	return math.Float32frombits(u32)
+}
+
+func FromFloat32(f float32) BF16 {
+	u32 := math.Float32bits(f)
+	return BF16(u32 >> 16)
+}

types/bfloat16/bfloat16_test.go

@@ -0,0 +1,53 @@
+package bfloat16
+
+import (
+	"crypto/rand"
+	"reflect"
+	"testing"
+)
+
+func randomBytes(n int) []byte {
+	out := make([]byte, n)
+	if _, err := rand.Read(out); err != nil {
+		panic(err)
+	}
+	return out
+}
+
+func TestEncodeDecode(t *testing.T) {
+	b := randomBytes(1024)
+	bf16 := Decode(b)
+	out := Encode(bf16)
+	if !reflect.DeepEqual(b, out) {
+		t.Fatalf("%+v != %+v", b, out)
+	}
+}
+
+func TestEncodeDecodeFloat32(t *testing.T) {
+	b := randomBytes(1024)
+	bf16 := DecodeFloat32(b)
+	out := EncodeFloat32(bf16)
+	if !reflect.DeepEqual(b, out) {
+		t.Fatalf("%+v != %+v", b, out)
+	}
+}
+
+func TestBasicFloat32(t *testing.T) {
+	var in float32 = 1.0
+	out := ToFloat32(FromFloat32(in))
+	if !reflect.DeepEqual(in, out) {
+		t.Fatalf("%+v != %+v", in, out)
+	}
+}
+
+func TestComplexFloat32(t *testing.T) {
+	var in float32 = 123456789123456789.123456789
+	var want float32 = 123286039799267328.0
+	out := ToFloat32(FromFloat32(in))
+	if in == out {
+		t.Fatalf("no loss of precision")
+	}
+	if out != want {
+		t.Fatalf("%.16f != %.16f", want, out)
+	}
+}