sample: do all sorting in topK

sample/samplers.go

@@ -1,10 +1,11 @@
 package sample
 
 import (
+	"errors"
 	"math"
-	"math/rand"
+	"math/rand/v2"
+	"slices"
 	"sync"
-	"time"
 
 	"github.com/ollama/ollama/llama"
 )
@@ -83,59 +84,56 @@ func (s *Sampler) sample(tokens []token) (token, error) {
 		return greedy(tokens), nil
 	}
 
-	if s.topK > 0 {
-		tokens = topK(tokens, s.topK)
-	} else {
-		sortLogits(tokens)
-	}
+	// topK also sorts the tokens in descending order of logits
+	tokens = topK(tokens, s.topK)
+
+	// token logit values are updated to probabilities
+	tokens = temperature(tokens, s.temperature)
 
 	tokens = topP(tokens, s.topP)
 	tokens = minP(tokens, s.minP)
 
-	// token logit values are updated to probabilities
-	temperature(tokens, s.temperature)
-	softmax(tokens)
-	return tokens[dist(tokens, s.rng.Int63())], nil
+	// TODO: this should fall back to greedy sampling
+	// or topP, topK values etc should be such that
+	// there are always tokens to sample from
+	if len(tokens) == 0 {
+		return token{}, errors.New("no tokens to sample from")
+	}
 
-	// // TODO: this should fall back to greedy sampling
-	// // or topP, topK values etc should be such that
-	// // there are always tokens to sample from
-	// if len(tokens) == 0 {
-	// 	return token{}, errors.New("no tokens to sample from")
-	// }
+	var r float32
+	if s.rng != nil {
+		r = s.rng.Float32()
+	} else {
+		r = rand.Float32()
+	}
 
-	// var r float32
-	// if s.rng != nil {
-	// 	r = s.rng.Float32()
-	// } else {
-	// 	r = rand.Float32()
-	// }
+	// Calculate cumulative sum of probabilities
+	var sum float32
+	for i := range tokens {
+		sum += tokens[i].value
+		tokens[i].value = sum
+	}
+	r *= tokens[len(tokens)-1].value
 
-	// // Calculate cumulative sum of probabilities
-	// var sum float32
-	// for i := range tokens {
-	// 	sum += tokens[i].value
-	// 	tokens[i].value = sum
-	// }
-	// r *= tokens[len(tokens)-1].value
+	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
+		if token.value < target {
+			return -1
+		}
+		return 1
+	})
 
-	// idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
-	// 	if token.value < target {
-	// 		return -1
-	// 	}
-	// 	return 1
-	// })
-
-	// return tokens[idx], nil
+	return tokens[idx], nil
 }
 
 // TODO(parthsareen): update sampler interface to use json unmarshal https://github.com/ollama/ollama/issues/9278
 func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int, grammar *Grammar) Sampler {
 	var rng *rand.Rand
 	if seed != -1 {
-		rng = rand.New(rand.NewSource(int64(seed)))
-	} else {
-		rng = rand.New(rand.NewSource(time.Now().UnixNano()))
+		// PCG requires two parameters: sequence and stream
+		// Use original seed for sequence
+		sequence := uint64(seed)
+		// Use golden ratio hash to generate statistically independent seeds
+		rng = rand.New(rand.NewPCG(sequence, sequence^0x9E3779B9))
 	}
 	if temperature < 0.0 {
 		temperature = 0.0

sample/transforms.go

@@ -3,7 +3,6 @@ package sample
 import (
 	"container/heap"
 	"math"
-	"math/rand"
 	"slices"
 )
 
@@ -11,7 +10,7 @@ import (
 type tokenHeap []token
 
 func (h tokenHeap) Len() int           { return len(h) }
-func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value } // Use < for min-heap to track largest elements
+func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
 func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
 
 func (h *tokenHeap) Push(x any) {
@@ -26,10 +25,45 @@ func (h *tokenHeap) Pop() any {
 	return x
 }
 
+// temperature applies scaling and softmax to the logits
+func temperature(ts []token, temp float32) []token {
+	// Find max logit for numerical stability
+	maxLogit := float32(math.Inf(-1))
+	for _, t := range ts {
+		if t.value > maxLogit {
+			maxLogit = t.value
+		}
+	}
+
+	// Apply temperature and compute exp(x - max)
+	temp = max(temp, 1e-7)
+	var sum float32
+	for i, v := range ts {
+		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
+		sum += ts[i].value
+	}
+
+	// Normalize
+	for i := range ts {
+		ts[i].value /= sum
+	}
+
+	return ts
+}
+
 // topK limits the number of tokens considered to the k highest logits
 func topK(ts []token, k int) []token {
-	if k >= len(ts) {
-		sortLogits(ts)
+	if k >= len(ts) || k <= 0 {
+		slices.SortFunc(ts, func(a, b token) int {
+			switch {
+			case a.value < b.value:
+				return 1
+			case a.value > b.value:
+				return -1
+			default:
+				return 0
+			}
+		})
 		return ts
 	}
 
@@ -47,7 +81,7 @@ func topK(ts []token, k int) []token {
 	}
 
 	// Convert heap to sorted slice in descending order
-	result := make([]token, k)
+	result := make([]token, len(h))
 	for i := k - 1; i >= 0; i-- {
 		result[i] = heap.Pop(&h).(token)
 	}
@@ -100,134 +134,3 @@ func minP(ts []token, p float32) []token {
 	ts = validTokens
 	return ts
 }
-
-// partialSortLogits uses quickselect to efficiently find and sort the top n tokens
-func partialSortLogits(ts []token, n int) []token {
-	if n >= len(ts) {
-		n = len(ts)
-	}
-
-	left, right := 0, len(ts)-1
-	target := n - 1
-
-	// Quickselect algorithm to partition array around pivot
-	for left < right {
-		// Choose middle element as pivot and move it to the end
-		pivot := left + (right-left)/2
-		ts[pivot], ts[right] = ts[right], ts[pivot]
-
-		// storeIndex tracks where to put next element greater than pivot
-		storeIndex := left
-		pivotValue := ts[right].value
-
-		// Partition array into elements >= pivot and < pivot
-		// Elements >= pivot go to the left side
-		for i := left; i < right; i++ {
-			if ts[i].value >= pivotValue {
-				ts[storeIndex], ts[i] = ts[i], ts[storeIndex]
-				storeIndex++
-			}
-		}
-
-		// Move pivot to its final position
-		ts[right], ts[storeIndex] = ts[storeIndex], ts[right]
-
-		// If pivot is at target position, we're done
-		// Otherwise recursively partition the half containing target
-		if storeIndex == target {
-			break
-		} else if storeIndex < target {
-			left = storeIndex + 1 // Target is in right half
-		} else {
-			right = storeIndex - 1 // Target is in left half
-		}
-	}
-
-	// Sort just the top n elements in descending order
-	slices.SortFunc(ts[:n], func(a, b token) int {
-		if a.value > b.value {
-			return -1
-		}
-		if a.value < b.value {
-			return 1
-		}
-		return 0
-	})
-
-	return ts[:n]
-}
-
-// sortLogits uses partialSortLogits to efficiently sort tokens
-// It sorts approximately sqrt(len(tokens)) elements which balances
-// between having enough tokens for sampling while avoiding full sort
-func sortLogits(ts []token) {
-	// Use sqrt of token length as a heuristic for partial sort size
-	// This provides a good balance between performance and having enough tokens
-	n := int(math.Sqrt(float64(len(ts)))) + 1
-
-	// Ensure we have at least 100 tokens and at most 1000
-	switch {
-	case n < 100:
-		n = 100
-	case n > 1000:
-		n = 1000
-	}
-
-	partialSortLogits(ts, n)
-}
-
-func temperature(ts []token, temp float32) {
-	for i := range ts {
-		ts[i].value /= temp
-	}
-}
-
-func softmax(ts []token) {
-	if len(ts) == 0 {
-		return
-	}
-
-	// Find max logit for numerical stability
-	maxLogit := ts[0].value
-	for _, t := range ts {
-		if t.value > maxLogit {
-			maxLogit = t.value
-		}
-	}
-
-	// Compute exp(logit - maxLogit) and sum them
-	var sumExp float32
-	for i, t := range ts {
-		expVal := float32(math.Exp(float64(t.value - maxLogit)))
-		ts[i].value = expVal
-		sumExp += expVal
-	}
-
-	// Normalize probabilities
-	for i := range ts {
-		ts[i].value /= sumExp
-	}
-}
-
-// applyDist selects a token based on probabilities and seed
-func dist(ts []token, seed int64) int {
-	rng := rand.New(rand.NewSource(seed))
-
-	cdf := make([]float32, len(ts))
-	var cumSum float32
-	for i, t := range ts {
-		cumSum += t.value
-		cdf[i] = cumSum
-	}
-
-	r := rng.Float32() * cumSum
-
-	// Select token based on CDF
-	for i, probSum := range cdf {
-		if r < probSum {
-			return i
-		}
-	}
-
-	return len(ts) - 1
-}

sample/transforms_test.go

@@ -1,13 +1,8 @@
 package sample
 
 import (
-	"encoding/binary"
-	"errors"
 	"math"
 	"math/rand/v2"
-	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
 )
 
@@ -64,7 +59,7 @@ func TestTemperatureAndSoftmax(t *testing.T) {
 func TestTopK(t *testing.T) {
 	input := []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
 
-	// Test k=3
+	// Test k=5
 	got := topK(toTokens(input), 5)
 	if len(got) != 5 {
 		t.Errorf("topK(5): wrong length: want 5, got %d", len(got))
@@ -77,6 +72,24 @@ func TestTopK(t *testing.T) {
 	if len(got) != len(input) {
 		t.Errorf("topK(20): wrong length: want %d, got %d", len(input), len(got))
 	}
+
+	// Test k=-1
+	input = []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
+	want = []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154, 0.043722924, 0.036774673, 0.026986899, 0.01681367, 0.0046718004, 0.00412893, 0.0030491839}
+	got = topK(toTokens(input), -1)
+	if len(got) != len(input) {
+		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(got))
+	}
+	compareLogits(t, "topK(-1)", want, got)
+
+	// Test k=0
+	input = []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
+	want = []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154, 0.043722924, 0.036774673, 0.026986899, 0.01681367, 0.0046718004, 0.00412893, 0.0030491839}
+	got = topK(toTokens(input), 0)
+	if len(got) != len(input) {
+		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(got))
+	}
+	compareLogits(t, "topK(-1)", want, got)
 }
 
 func TestTopP(t *testing.T) {
@@ -85,7 +98,7 @@ func TestTopP(t *testing.T) {
 
 	// First apply temperature and softmax to get probabilities
 	tokens = temperature(tokens, 1)
-	sortLogits(tokens)
+	tokens = topK(tokens, 20)
 
 	// Then apply topP
 	got := topP(tokens, 0.95)
@@ -117,7 +130,7 @@ func TestSortLogits(t *testing.T) {
 	input := []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
 	tokens := toTokens(input)
 
-	sortLogits(tokens)
+	tokens = topK(tokens, 20)
 
 	for i := 1; i < len(tokens); i++ {
 		if tokens[i].value > tokens[i-1].value {
@@ -130,98 +143,6 @@ func TestSortLogits(t *testing.T) {
 	compareLogits(t, "sortLogits", want, tokens)
 }
 
-// TestSortLogitsWithRealData tests sorting behavior using real model logit distributions
-func TestSortLogitsWithRealData(t *testing.T) {
-	// This will be populated from testdata/logits.bin
-	// Format: 32-bit float array in binary format
-	logits, err := loadTestLogits(t)
-	if err != nil {
-		t.Skipf("Skipping real logit test: %v", err)
-		return
-	}
-
-	tokens := toTokens(logits)
-	sortLogits(tokens)
-
-	// Calculate n for verification
-	n := int(math.Sqrt(float64(len(tokens)))) + 1
-	if n > 1000 {
-		n = 1000
-	} else if n < 100 {
-		n = 100
-	}
-
-	t.Logf("Testing with %d tokens, partial sorting top %d", len(tokens), n)
-
-	// Only verify the top n elements are sorted (which is what we guarantee)
-	// This is much faster than checking the entire array
-	topN := tokens[:n]
-	for i := 1; i < len(topN); i++ {
-		if topN[i].value > topN[i-1].value {
-			t.Fatalf("top %d tokens not properly sorted at index %d: %.15f > %.15f",
-				n, i, topN[i].value, topN[i-1].value)
-		}
-	}
-
-	// Verify we didn't lose any high value tokens by checking that
-	// all tokens after position n are <= the nth token
-	// Do this in chunks to avoid timeouts on large arrays
-	nthValue := tokens[n-1].value
-	const chunkSize = 1000
-
-	for start := n; start < len(tokens); start += chunkSize {
-		end := min(start+chunkSize, len(tokens))
-		for i := start; i < end; i++ {
-			if tokens[i].value > nthValue {
-				t.Fatalf("found higher value token after position %d: tokens[%d].value = %.15f > %.15f",
-					n, i, tokens[i].value, nthValue)
-			}
-		}
-	}
-}
-
-// loadTestLogits loads logit test data from testdata/logits.bin
-func loadTestLogits(t *testing.T) ([]float32, error) {
-	t.Helper()
-
-	_, currFile, _, ok := runtime.Caller(0)
-	if !ok {
-		return nil, errors.New("could not determine test file path")
-	}
-	testDataPath := filepath.Join(filepath.Dir(currFile), "testdata", "logits.bin")
-
-	file, err := os.Open(testDataPath)
-	if err != nil {
-		return nil, err
-	}
-	defer file.Close()
-
-	stat, err := file.Stat()
-	if err != nil {
-		return nil, err
-	}
-
-	numFloats := stat.Size() / 4 // each float32 is 4 bytes
-	if numFloats*4 != stat.Size() {
-		return nil, errors.New("logits.bin has invalid size: not a multiple of 4 bytes")
-	}
-
-	logits := make([]float32, numFloats)
-	for i := range logits {
-		var val uint32
-		if err := binary.Read(file, binary.LittleEndian, &val); err != nil {
-			return nil, err
-		}
-		logits[i] = math.Float32frombits(val)
-	}
-
-	if len(logits) == 0 {
-		return nil, errors.New("logits.bin is empty")
-	}
-
-	return logits, nil
-}
-
 func BenchmarkTransforms(b *testing.B) {
 	// Generate random logits
 	tokens := make([]token, 1<<16)
@@ -270,7 +191,7 @@ func BenchmarkTransforms(b *testing.B) {
 		b.ResetTimer()
 		for b.Loop() {
 			copy(tokensCopy, tokens)
-			sortLogits(tokensCopy)
+			topK(tokensCopy, 200000)
 		}
 	})
 }