extras

sample/samplers.go

@@ -1,11 +1,10 @@
 package sample
 
 import (
-	"errors"
 	"math"
-	"math/rand/v2"
+	"math/rand"
-	"slices"
 	"sync"
+	"time"
 
 	"github.com/ollama/ollama/llama"
 )
@@ -90,53 +89,53 @@ func (s *Sampler) sample(tokens []token) (token, error) {
 		sortLogits(tokens)
 	}
 
-	// token logit values are updated to probabilities
-	tokens = temperature(tokens, s.temperature)
-
 	tokens = topP(tokens, s.topP)
 	tokens = minP(tokens, s.minP)
 
-	// TODO: this should fall back to greedy sampling
+	// token logit values are updated to probabilities
-	// or topP, topK values etc should be such that
+	temperature(tokens, s.temperature)
-	// there are always tokens to sample from
+	softmax(tokens)
-	if len(tokens) == 0 {
+	return tokens[dist(tokens, s.rng.Int63())], nil
-		return token{}, errors.New("no tokens to sample from")
-	}
 
-	var r float32
+	// // TODO: this should fall back to greedy sampling
-	if s.rng != nil {
+	// // or topP, topK values etc should be such that
-		r = s.rng.Float32()
+	// // there are always tokens to sample from
-	} else {
+	// if len(tokens) == 0 {
-		r = rand.Float32()
+	// 	return token{}, errors.New("no tokens to sample from")
-	}
+	// }
 
-	// Calculate cumulative sum of probabilities
+	// var r float32
-	var sum float32
+	// if s.rng != nil {
-	for i := range tokens {
+	// 	r = s.rng.Float32()
-		sum += tokens[i].value
+	// } else {
-		tokens[i].value = sum
+	// 	r = rand.Float32()
-	}
+	// }
-	r *= tokens[len(tokens)-1].value
 
-	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
+	// // Calculate cumulative sum of probabilities
-		if token.value < target {
+	// var sum float32
-			return -1
+	// for i := range tokens {
-		}
+	// 	sum += tokens[i].value
-		return 1
+	// 	tokens[i].value = sum
-	})
+	// }
+	// r *= tokens[len(tokens)-1].value
 
-	return tokens[idx], nil
+	// idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
+	// 	if token.value < target {
+	// 		return -1
+	// 	}
+	// 	return 1
+	// })
+
+	// 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 {
-		// PCG requires two parameters: sequence and stream
+		rng = rand.New(rand.NewSource(int64(seed)))
-		// Use original seed for sequence
+	} else {
-		sequence := uint64(seed)
+		rng = rand.New(rand.NewSource(time.Now().UnixNano()))
-		// 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

@@ -1,92 +1,58 @@
 package sample
 
 import (
+	"container/heap"
 	"math"
+	"math/rand"
 	"slices"
 )
 
-// temperature applies scaling and softmax to the logits
+// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
-func temperature(ts []token, temp float32) []token {
+type tokenHeap []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)
+func (h tokenHeap) Len() int           { return len(h) }
-	temp = max(temp, 1e-7)
+func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value } // Use < for min-heap to track largest elements
-	var sum float32
+func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
-	for i, v := range ts {
-		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
-		sum += ts[i].value
-	}
 
-	// Normalize
+func (h *tokenHeap) Push(x any) {
-	for i := range ts {
+	*h = append(*h, x.(token))
-		ts[i].value /= sum
-	}
-
-	return ts
 }
 
-// siftDown maintains a min-heap property by recursively moving larger elements down the heap.
+func (h *tokenHeap) Pop() any {
-//
+	old := *h
-// The heap is represented as an array where for any node at index i:
+	n := len(old)
-// - Left child is at index 2i + 1
+	x := old[n-1]
-// - Right child is at index 2i + 2
+	*h = old[0 : n-1]
-// - Parent is at index (i-1)/2
+	return x
-//
-// The function compares a node with its children and:
-// 1. Finds the smallest value between the node and its children
-// 2. If the node is not the smallest, swaps it with its smallest child
-// 3. Continues this process down the affected path until the min-heap property is restored
-func siftDown(data []token, start, end int) {
-	root := start
-	for {
-		child := 2*root + 1
-		if child >= end {
-			break
-		}
-		// Find smaller child (we want min heap)
-		if child+1 < end && data[child+1].value < data[child].value {
-			child++
-		}
-		// Exit if root is already smaller than children
-		if data[root].value <= data[child].value {
-			break
-		}
-		// Swap with smaller child and continue
-		data[root], data[child] = data[child], data[root]
-		root = child
-	}
 }
 
 // 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)
 		return ts
 	}
-	// Heapify + siftDown - O(nlog(k))
-	// Build min-heap of first k elements
-	heap := ts[:k]
-	for i := k/2 - 1; i >= 0; i-- {
-		siftDown(heap, i, k)
-	}
 
-	// Process remaining elements - if larger than heap root, replace root
+	// Initialize min-heap with first k elements
+	h := make(tokenHeap, k)
+	copy(h, ts[:k])
+	heap.Init(&h)
+
+	// Process remaining elements
 	for i := k; i < len(ts); i++ {
-		if ts[i].value > heap[0].value {
+		if ts[i].value > h[0].value {
-			heap[0] = ts[i]
+			heap.Pop(&h)
-			siftDown(heap, 0, k)
+			heap.Push(&h, ts[i])
 		}
 	}
 
-	slices.Reverse(heap)
+	// Convert heap to sorted slice in descending order
+	result := make([]token, k)
+	for i := k - 1; i >= 0; i-- {
+		result[i] = heap.Pop(&h).(token)
+	}
 
-	ts = heap
+	return result
-	return ts
 }
 
 // topP limits tokens to those with cumulative probability p
@@ -135,61 +101,133 @@ func minP(ts []token, p float32) []token {
 	return ts
 }
 
-// TODO(parthsareen): possibly replace with simpler implementation https://github.com/ollama/ollama/issues/9584
+// partialSortLogits uses quickselect to efficiently find and sort the top n tokens
-// sortLogits sorts implementation to sort tokens by logits using counting sort
+func partialSortLogits(ts []token, n int) []token {
-// counting sort is faster than built-in sort for this use case
+	if n >= len(ts) {
-func sortLogits(tokens []token) {
+		n = len(ts)
-	if len(tokens) <= 1 {
+	}
+
+	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/min in a single pass
+	// Find max logit for numerical stability
-	minLogit, maxLogit := tokens[0].value, tokens[0].value
+	maxLogit := ts[0].value
-	for _, t := range tokens[1:] {
+	for _, t := range ts {
-		if t.value < minLogit {
+		if t.value > maxLogit {
-			minLogit = t.value
-		} else if t.value > maxLogit {
 			maxLogit = t.value
 		}
 	}
 
-	// Calculate scaling to map to uint32 range
+	// Compute exp(logit - maxLogit) and sum them
-	logitRange := maxLogit - minLogit
+	var sumExp float32
-	if logitRange < 1e-6 {
+	for i, t := range ts {
-		return // All values effectively equal
+		expVal := float32(math.Exp(float64(t.value - maxLogit)))
+		ts[i].value = expVal
+		sumExp += expVal
 	}
 
-	// Count frequencies directly from tokens
+	// Normalize probabilities
-	const maxInt = (1 << 24) - 1 // Use 24 bits for good granularity
+	for i := range ts {
-	var counts [256]int          // For first byte
+		ts[i].value /= sumExp
-
-	// First pass: count frequencies
-	for _, t := range tokens {
-		// Map to [0, maxInt] range
-		score := min(uint32((t.value-minLogit)*float32(maxInt)/logitRange), maxInt)
-		counts[score>>16]++
 	}
-
+}
-	// Calculate offsets
+
-	var offset int
+// applyDist selects a token based on probabilities and seed
-	for i := range counts {
+func dist(ts []token, seed int64) int {
-		count := counts[i]
+	rng := rand.New(rand.NewSource(seed))
-		counts[i] = offset
+
-		offset += count
+	cdf := make([]float32, len(ts))
-	}
+	var cumSum float32
-
+	for i, t := range ts {
-	// Second pass: place elements in correct position
+		cumSum += t.value
-	output := make([]token, len(tokens))
+		cdf[i] = cumSum
-	// Track current positions
+	}
-	countsCopy := counts
+
-
+	r := rng.Float32() * cumSum
-	for i, t := range tokens {
+
-		score := min(uint32((t.value-minLogit)*float32(maxInt)/logitRange), maxInt)
+	// Select token based on CDF
-
+	for i, probSum := range cdf {
-		pos := countsCopy[score>>16]
+		if r < probSum {
-		countsCopy[score>>16]++
+			return i
-		output[len(tokens)-1-pos] = tokens[i]
+		}
-	}
+	}
-
+
-	copy(tokens, output)
+	return len(ts) - 1
 }

sample/transforms_test.go

@@ -1,39 +1,44 @@
 package sample
 
 import (
+	"encoding/binary"
+	"errors"
 	"math"
 	"math/rand/v2"
+	"os"
+	"path/filepath"
+	"runtime"
 	"testing"
 )
 
-// Helper to convert float64 slice to logit slice
+// Helper to convert float32 slice to logit slice
-func toTokens(values []float64) []token {
+func toTokens(values []float32) []token {
 	tokens := make([]token, len(values))
 	for i, v := range values {
 		tokens[i] = token{
 			id:    int32(i),
-			value: float32(v),
+			value: v,
 		}
 	}
 	return tokens
 }
 
 // Helper to compare logit slices
-func compareLogits(t *testing.T, name string, want []float64, got []token) {
+func compareLogits(t *testing.T, name string, want []float32, got []token) {
 	t.Helper()
 	if len(want) != len(got) {
 		t.Errorf("%s: length mismatch: want %d, got %d", name, len(want), len(got))
 		return
 	}
 	for i := range want {
-		if math.Abs(float64(got[i].value)-want[i]) > 1e-6 {
+		if math.Abs(float64(got[i].value-want[i])) > 1e-6 {
 			t.Errorf("%s: index %d: want %f, got %f", name, i, want[i], got[i].value)
 		}
 	}
 }
 
 func TestTemperatureAndSoftmax(t *testing.T) {
-	input := []float64{1, 4, -2, 0}
+	input := []float32{1, 4, -2, 0}
 	got := temperature(toTokens(input), 0.5)
 
 	// Check probabilities sum to 1
@@ -41,7 +46,7 @@ func TestTemperatureAndSoftmax(t *testing.T) {
 	for _, token := range got {
 		sum += token.value
 	}
-	if math.Abs(float64(sum)-1.0) > 1e-6 {
+	if math.Abs(float64(sum-1.0)) > 1e-6 {
 		t.Errorf("probabilities don't sum to 1: got %f", sum)
 	}
 
@@ -51,30 +56,31 @@ func TestTemperatureAndSoftmax(t *testing.T) {
 	for _, token := range got {
 		sum += token.value
 	}
-	if math.Abs(float64(sum)-1.0) > 1e-6 {
+	if math.Abs(float64(sum-1.0)) > 1e-6 {
 		t.Errorf("probabilities don't sum to 1: got %f", sum)
 	}
 }
 
 func TestTopK(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4}
+	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
-	got := topK(toTokens(input), 3)
+	got := topK(toTokens(input), 5)
-	if len(got) != 3 {
+	if len(got) != 5 {
-		t.Errorf("topK(3): wrong length: want 3, got %d", len(got))
+		t.Errorf("topK(5): wrong length: want 5, got %d", len(got))
 	}
-	// Should keep highest 3 values: 4, 2, 1
+	// Should keep highest 3 values in descending order
-	want := []float64{4, 2, 1}
+	want := []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154}
 	compareLogits(t, "topK(3)", want, got)
 
-	// Test k > len
+	got = topK(toTokens(input), 20)
-	got = topK(toTokens(input), 10)
+	if len(got) != len(input) {
-	compareLogits(t, "topK(10)", input, got)
+		t.Errorf("topK(20): wrong length: want %d, got %d", len(input), len(got))
+	}
 }
 
 func TestTopP(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4}
+	input := []float32{-3, -2, -1, 0, 1, 2, 4}
 	tokens := toTokens(input)
 
 	// First apply temperature and softmax to get probabilities
@@ -92,7 +98,7 @@ func TestTopP(t *testing.T) {
 }
 
 func TestMinP(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4, 3}
+	input := []float32{-3, -2, -1, 0, 1, 2, 4, 3}
 	tokens := toTokens(input)
 
 	// First apply temperature and softmax
@@ -108,7 +114,7 @@ func TestMinP(t *testing.T) {
 }
 
 func TestSortLogits(t *testing.T) {
-	input := []float64{3, 1, 4, 2, -1, 0, -2}
+	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)
@@ -120,10 +126,102 @@ func TestSortLogits(t *testing.T) {
 		}
 	}
 
-	want := []float64{4, 3, 2, 1, 0, -1, -2}
+	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}
 	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)