ollama/mlxrunner/sample/sample_test.go

package sample

import (
	"math"
	"slices"
	"testing"

	"github.com/ollama/ollama/mlx"
	"github.com/ollama/ollama/mlx/mlxtest"
)

// slotLogits builds a [1, V] logits tensor for a single-slot Sample call.
func slotLogits(values []float32) *mlx.Array {
	return mlx.FromValues(values, 1, len(values))
}

// batchLogits stacks per-row float32 slices of equal length into a [B, V]
// logits tensor.
func batchLogits(rows ...[]float32) *mlx.Array {
	v := len(rows[0])
	flat := make([]float32, 0, len(rows)*v)
	for _, r := range rows {
		if len(r) != v {
			panic("batchLogits: rows must share vocab size")
		}
		flat = append(flat, r...)
	}
	return mlx.FromValues(flat, len(rows), v)
}

// sampleOne runs Sample on a freshly-added single slot and returns the
// sampled token id. Used both for the single-slot options table and as the
// reference oracle for the batched-equivalence test.
func sampleOne(t *mlxtest.T, opts Options, priorTokens []int32, values []float32) int32 {
	t.Helper()
	s := New(128)
	t.Cleanup(func() {
		s.Free()
	})
	s.Add(0, opts, priorTokens)

	got := s.Sample([]int{0}, slotLogits(values)).Token
	mlx.Eval(got)
	return got.Int()
}

// logOf returns log(p) as a float32 so tests can build logits that softmax to
// a chosen probability distribution.
func logOf(p float64) float32 { return float32(math.Log(p)) }

// TestSampleSingleSlotOptions pins the per-slot behavior of each Options
// knob against a concrete expected token. Expected values are worked out by
// hand from the math of each transform, not from a second call into the
// sampler — so a regression in any single transform shows up here.
func TestSampleSingleSlotOptions(t *testing.T) {
	mlxtest.SkipIfUnavailable(t)

	cases := []struct {
		name   string
		opts   Options
		priors []int32
		logits []float32
		want   int32
	}{
		{
			name:   "presence penalty",
			opts:   Options{RepeatLastN: 1, PresencePenalty: 6},
			priors: []int32{1},
			logits: []float32{0, 5, 4},
			want:   2, // token 1: 5 - 6 = -1, argmax shifts to 2
		},
		{
			name:   "repeat penalty on positive logits",
			opts:   Options{RepeatLastN: 1, RepeatPenalty: 2},
			priors: []int32{1},
			logits: []float32{0, 5, 4},
			want:   2, // token 1 positive → divided: 5/2 = 2.5, argmax shifts to 2
		},
		{
			name:   "repeat penalty on negative logits",
			opts:   Options{RepeatLastN: 1, RepeatPenalty: 4},
			priors: []int32{1},
			logits: []float32{-5, -1, -3},
			want:   2, // token 1 negative → multiplied: -1*4 = -4, argmax shifts to 2
		},
		{
			name:   "frequency penalty",
			opts:   Options{RepeatLastN: 4, FrequencyPenalty: 2},
			priors: []int32{1, 1},
			logits: []float32{0, 5, 4},
			want:   2, // 5 - 2*count(1)=2*2=4 → 1, argmax shifts to 2
		},
		{
			name:   "top-k",
			opts:   Options{Temperature: 1, TopK: 1},
			logits: []float32{1, 5, 4},
			want:   1, // only argmax survives → deterministic even with temperature
		},
		{
			name:   "top-p",
			opts:   Options{Temperature: 1, TopP: 0.4},
			logits: []float32{logOf(0.5), logOf(0.3), logOf(0.2)},
			want:   0, // exclusive cumsum below 0.4 keeps only token 0
		},
		{
			name:   "min-p",
			opts:   Options{Temperature: 1, MinP: 0.7},
			logits: []float32{logOf(0.5), logOf(0.3), logOf(0.2)},
			want:   0, // threshold 0.5*0.7=0.35 drops all but the top token
		},
		{
			name:   "RepeatLastN=0 disables penalties",
			opts:   Options{RepeatLastN: 0, RepeatPenalty: 2, PresencePenalty: 10},
			priors: []int32{1},
			logits: []float32{0, 5, 4},
			want:   1, // 0 = disabled per API contract, argmax unchanged
		},
		{
			name:   "RepeatLastN=-1 resolves to num_ctx",
			opts:   Options{RepeatLastN: -1, PresencePenalty: 6},
			priors: []int32{1},
			logits: []float32{0, 5, 4},
			want:   2, // -1 → num_ctx (128); penalty applies, argmax shifts
		},
	}

	for _, tc := range cases {
		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
			if got := sampleOne(t, tc.opts, tc.priors, tc.logits); got != tc.want {
				t.Errorf("got %d, want %d", got, tc.want)
			}
		})
	}
}

func TestDistributionAppliesTopKBeforeTopP(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})
		s.Add(0, Options{Temperature: 1, TopK: 2, TopP: 0.7}, nil)

		dist := s.Distribution(0, slotLogits([]float32{logOf(0.6), logOf(0.2), logOf(0.2)}), nil)
		mlx.Eval(dist.Arrays()...)

		ids := dist.IDs.Ints()
		probs := dist.Probs.Floats()
		if len(ids) != 2 || len(probs) != 2 {
			t.Fatalf("support = ids %v probs %v, want 2 sparse entries", ids, probs)
		}

		foundTop := false
		for i, id := range ids {
			switch id {
			case 0:
				foundTop = true
				if math.Abs(float64(probs[i]-1)) > 1e-5 {
					t.Fatalf("top token prob = %v, want 1; ids=%v probs=%v", probs[i], ids, probs)
				}
			default:
				if math.Abs(float64(probs[i])) > 1e-5 {
					t.Fatalf("non-top token %d prob = %v, want 0; ids=%v probs=%v", id, probs[i], ids, probs)
				}
			}
		}
		if !foundTop {
			t.Fatalf("top-k support %v did not include token 0", ids)
		}
	})
}

func TestDistributionResidualUsesTargetSupport(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		target := Distribution{
			IDs:   mlx.NewArrayInt32([]int32{2, 5}, []int32{1, 2}),
			Probs: mlx.FromValues([]float32{0.7, 0.3}, 1, 2),
		}
		draft := Distribution{
			IDs:   mlx.NewArrayInt32([]int32{2, 4}, []int32{1, 2}),
			Probs: mlx.FromValues([]float32{0.2, 0.8}, 1, 2),
		}

		residual := target.ResidualAgainst(draft)
		mlx.Eval(residual.Arrays()...)

		ids := residual.IDs.Ints()
		probs := residual.Probs.Floats()
		want := map[int32]float64{2: 0.625, 5: 0.375}
		if len(ids) != 2 || len(probs) != 2 {
			t.Fatalf("residual = ids %v probs %v, want 2 sparse entries", ids, probs)
		}
		for i, id := range ids {
			w, ok := want[id]
			if !ok {
				t.Fatalf("residual includes token %d outside target support: ids=%v probs=%v", id, ids, probs)
			}
			if math.Abs(float64(probs[i])-w) > 1e-5 {
				t.Fatalf("residual token %d prob = %v, want %v; ids=%v probs=%v", id, probs[i], w, ids, probs)
			}
		}
	})
}

func TestSeededSamplingIsReproducible(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		seededSequence := func(seed int) []int32 {
			s := New(128)
			t.Cleanup(func() {
				s.Free()
			})
			s.Add(0, Options{Temperature: 1, TopK: 4, Seed: seed, UseSeed: true}, nil)

			logits := slotLogits([]float32{0, 0, 0, 0})
			out := make([]int32, 32)
			for i := range out {
				token := s.Sample([]int{0}, logits).Token
				mlx.Eval(token)
				out[i] = token.Int()
			}
			return out
		}

		a := seededSequence(1234)
		b := seededSequence(1234)
		if !slices.Equal(a, b) {
			t.Fatalf("same seed produced different sequences:\n%v\n%v", a, b)
		}

		c := seededSequence(5678)
		if slices.Equal(a, c) {
			t.Fatalf("different seeds produced the same sequence: %v", a)
		}
	})
}

func TestSeededBernoulliIsReproducible(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		seededMask := func() []int32 {
			s := New(128)
			t.Cleanup(func() {
				s.Free()
			})
			s.Add(0, Options{Seed: 99, UseSeed: true}, nil)

			mask := s.Bernoulli(0, mlx.FromValues([]float32{0.5, 0.5, 0.5, 0.5, 0.5, 0.5}, 6)).AsType(mlx.DTypeInt32)
			mlx.Eval(mask)
			return mask.Ints()
		}

		a := seededMask()
		b := seededMask()
		if !slices.Equal(a, b) {
			t.Fatalf("same seed produced different bernoulli masks:\n%v\n%v", a, b)
		}
	})
}

// TestSampleHistoryWindow verifies that penalty history respects the
// RepeatLastN window: priors longer than RepeatLastN are trimmed on Add,
// and once the ring wraps, tokens that rotate out no longer contribute
// to penalties.
func TestSampleHistoryWindow(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})

		// RepeatLastN=2 with priors {1, 2, 3}: makeHistoryRow keeps only
		// {2, 3}. Token 1 was trimmed — its penalty is NOT active.
		s.Add(0, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{1, 2, 3})

		// Step 1: logits favor token 1 (trimmed). If the trim were broken it
		// would be penalized and the argmax would move.
		step1 := s.Sample([]int{0}, slotLogits([]float32{0, 5, 0, 0, 0})).Token
		mlx.Eval(step1)
		if got := step1.Int(); got != 1 {
			t.Fatalf("step 1 = %d, want 1 (token 1 trimmed from priors)", got)
		}
		// After step 1 the ring holds {1, 3}; token 2 has rotated out.

		// Step 2: logits favor token 2 (rotated out). If the ring wrap were
		// wrong, token 2 would still be penalized.
		step2 := s.Sample([]int{0}, slotLogits([]float32{0, 0, 5, 0, 0})).Token
		mlx.Eval(step2)
		if got := step2.Int(); got != 2 {
			t.Fatalf("step 2 = %d, want 2 (token 2 rotated out of ring)", got)
		}
	})
}

func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})

		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{1, 2})
		draftTokens := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
		scores := s.SpeculativeScores(0, batchLogits(
			[]float32{0, 9, 9, 8, 0}, // history {1,2}; token 3 wins
			[]float32{0, 0, 9, 9, 8}, // history {2,3}; token 4 wins
			[]float32{0, 0, 9, 9, 8}, // history {3,4}; token 2 wins
		), draftTokens)
		tokens := scores.Argmax(-1, false).AsType(mlx.DTypeInt32)
		mlx.Eval(tokens)

		if got, want := tokens.Ints(), []int32{3, 4, 2}; len(got) != len(want) {
			t.Fatalf("tokens = %v, want %v", got, want)
		} else {
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("tokens = %v, want %v", got, want)
				}
			}
		}
		if s.byID[0].historyLen != 2 {
			t.Fatalf("historyLen = %d, want 2", s.byID[0].historyLen)
		}
	})
}

func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})

		// A proposal step passes one logits row with the chain's earlier drafts:
		// the single row is the chain's final step, so every draft belongs to
		// its history. Slot 0 exercises the batched history path (full ring),
		// slot 1 the serial path (ring not yet full).
		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{0, 1})
		s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{0, 1})
		prefix := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})

		for _, seqID := range []int{0, 1} {
			// Drafts 3 and 4 are penalized, so token 2 wins over the higher raw
			// scores; with the drafts absent from the history, token 3 would.
			dist := s.Distribution(seqID, batchLogits([]float32{0, 0, 9, 9, 8}), prefix)
			mlx.Eval(dist.IDs)
			if got := dist.IDs.Ints()[0]; got != 2 {
				t.Fatalf("seq %d token = %d, want 2 (drafts 3 and 4 penalized)", seqID, got)
			}
		}
	})
}

func TestDistributionMultiRowWithoutChain(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})

		// A block drafter's proposal batch samples every row from one call with
		// no draft chain: each row sees the slot history unchanged. Slot 0
		// exercises the batched history path (full ring), slot 1 the serial path
		// (ring not yet full).
		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{3, 4})
		s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{3, 4})

		for _, seqID := range []int{0, 1} {
			// Tokens 3 and 4 are penalized in every row alike; rows 1 and 3
			// share logits, so a chain alignment leaking between rows would
			// split their winners.
			dist := s.Distribution(seqID, batchLogits(
				[]float32{0, 0, 8, 9, 9},
				[]float32{0, 8, 0, 9, 9},
				[]float32{0, 0, 8, 9, 9},
			), nil)
			top := dist.IDs.Slice(mlx.Slice(), mlx.Slice(0, 1))
			mlx.Eval(top)
			if got, want := top.Ints(), []int32{2, 1, 2}; !slices.Equal(got, want) {
				t.Fatalf("seq %d top tokens = %v, want %v", seqID, got, want)
			}
		}
	})
}

func TestCommitBatchesRingWrites(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})

		s.Add(0, Options{RepeatLastN: 4, RepeatPenalty: 1.1}, []int32{10, 11, 12})
		s.Commit(0, []int32{20, 21, 22})
		s.Commit(0, []int32{30, 31, 32, 33, 34})
		mlx.Eval(s.history)

		got := s.history.Ints()
		want := []int32{32, 33, 34, 31}
		for i := range want {
			if got[i] != want[i] {
				t.Fatalf("history = %v, want %v", got, want)
			}
		}
		if s.byID[0].historyLen != 11 {
			t.Fatalf("historyLen = %d, want 11", s.byID[0].historyLen)
		}
	})
}

// TestBatchSamplingPreservesPerSlotBehavior is the core equivalence test:
// for every representative dispatch branch (uniform, serial on mixed opts,
// serial on partial ring, subset/out-of-order), a batched Sample call must
// produce the same token per row as running the same slot alone.
func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) {
	mlxtest.SkipIfUnavailable(t)

	type slot struct {
		id     int
		opts   Options
		priors []int32
	}

	cases := []struct {
		name   string
		slots  []slot
		sample []int
		rows   [][]float32
	}{
		{
			name: "uniform",
			slots: []slot{
				{10, Options{RepeatLastN: 2, PresencePenalty: 5}, []int32{1, 2}},
				{20, Options{RepeatLastN: 2, PresencePenalty: 5}, []int32{0, 2}},
			},
			sample: []int{10, 20},
			rows:   [][]float32{{0, 5, 4}, {3, 0, 0}},
		},
		{
			name: "serial — mixed opts",
			slots: []slot{
				{1, Options{RepeatLastN: 1, RepeatPenalty: 2}, []int32{1}},
				{2, Options{Temperature: 1, TopK: 1}, nil},
			},
			sample: []int{1, 2},
			rows:   [][]float32{{0, 5, 4, 1}, {2, 1, 5, 3}},
		},
		{
			name: "serial — partial ring",
			slots: []slot{
				{1, Options{RepeatLastN: 4, PresencePenalty: 5}, []int32{1, 1, 1, 1}},
				{2, Options{RepeatLastN: 4, PresencePenalty: 5}, []int32{2}},
			},
			sample: []int{1, 2},
			rows:   [][]float32{{0, 5, 4}, {0, 4, 5}},
		},
		{
			name: "subset out-of-order",
			slots: []slot{
				{10, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{1, 1}},
				{20, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{2, 2}},
				{30, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{3, 3}},
			},
			sample: []int{30, 10},
			rows:   [][]float32{{5, 5, 5, 0, 5, 5}, {5, 0, 5, 5, 0, 5}},
		},
	}

	for _, tc := range cases {
		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
			// Per-slot reference for each sampled seq.
			want := make([]int32, len(tc.sample))
			for i, id := range tc.sample {
				var spec slot
				for _, s := range tc.slots {
					if s.id == id {
						spec = s
						break
					}
				}
				want[i] = sampleOne(t, spec.opts, spec.priors, tc.rows[i])
			}

			// Batched call.
			s := New(128)
			t.Cleanup(func() {
				s.Free()
			})
			for _, spec := range tc.slots {
				s.Add(spec.id, spec.opts, spec.priors)
			}
			res := s.Sample(tc.sample, batchLogits(tc.rows...))
			mlx.Eval(res.Token)
			got := res.Token.Ints()

			for i, id := range tc.sample {
				if got[i] != want[i] {
					t.Errorf("seq %d: batched = %d, per-slot = %d", id, got[i], want[i])
				}
			}
		})
	}
}

// TestRemoveDoesNotLeakHistory: after Remove, a newly-added slot at the
// recycled row must start from its own priors only — no carryover from
// the removed slot's history.
func TestRemoveDoesNotLeakHistory(t *testing.T) {
	mlxtest.Run(t, func(t *mlxtest.T) {
		opts := Options{RepeatLastN: 1, PresencePenalty: 10}
		s := New(128)
		t.Cleanup(func() {
			s.Free()
		})
		s.Add(1, opts, []int32{1})
		s.Add(2, opts, []int32{2})
		s.Remove(1)
		s.Add(3, opts, []int32{0})

		// Slot 2 retains history {2}; slot 3 retains history {0}. With
		// equal logits and PresencePenalty=10 the argmax drops to the first
		// unpenalized token.
		res := s.Sample([]int{2, 3}, batchLogits(
			[]float32{3, 3, 0},
			[]float32{3, 3, 0},
		))
		mlx.Eval(res.Token)
		tokens := res.Token.Ints()
		if tokens[0] != 0 {
			t.Errorf("slot 2 = %d, want 0 (token 2 penalized)", tokens[0])
		}
		if tokens[1] != 1 {
			t.Errorf("slot 3 = %d, want 1 (token 0 penalized, no slot-1 carryover)", tokens[1])
		}
	})
}