mlx: apply repeat penalties in sampler (#15631)

2026-07-23 09:10:53 -05:00 · 2026-04-18 07:49:38 -07:00
parent 123b300af6
commit ff23dd343f
7 changed files with 164 additions and 85 deletions
--- a/x/mlxrunner/client.go
+++ b/x/mlxrunner/client.go
@@ -158,13 +158,15 @@ type completionRequest struct {
 }

 type completionOpts struct {
-	Temperature     float32 `json:"temperature,omitempty"`
-	TopP            float32 `json:"top_p,omitempty"`
-	MinP            float32 `json:"min_p,omitempty"`
-	TopK            int     `json:"top_k,omitempty"`
-	RepeatLastN     int     `json:"repeat_last_n,omitempty"`
-	PresencePenalty float32 `json:"presence_penalty,omitempty"`
-	NumPredict      int     `json:"num_predict,omitempty"`
+	Temperature      float32 `json:"temperature,omitempty"`
+	TopP             float32 `json:"top_p,omitempty"`
+	MinP             float32 `json:"min_p,omitempty"`
+	TopK             int     `json:"top_k,omitempty"`
+	RepeatLastN      int     `json:"repeat_last_n,omitempty"`
+	RepeatPenalty    float32 `json:"repeat_penalty,omitempty"`
+	PresencePenalty  float32 `json:"presence_penalty,omitempty"`
+	FrequencyPenalty float32 `json:"frequency_penalty,omitempty"`
+	NumPredict       int     `json:"num_predict,omitempty"`
 }

 type CompletionResponse struct {
@@ -206,13 +208,15 @@ func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn f
 	}
 	if req.Options != nil {
 		creq.Options = &completionOpts{
-			Temperature:     req.Options.Temperature,
-			TopP:            req.Options.TopP,
-			MinP:            req.Options.MinP,
-			TopK:            req.Options.TopK,
-			RepeatLastN:     req.Options.RepeatLastN,
-			PresencePenalty: req.Options.PresencePenalty,
-			NumPredict:      req.Options.NumPredict,
+			Temperature:      req.Options.Temperature,
+			TopP:             req.Options.TopP,
+			MinP:             req.Options.MinP,
+			TopK:             req.Options.TopK,
+			RepeatLastN:      req.Options.RepeatLastN,
+			RepeatPenalty:    req.Options.RepeatPenalty,
+			PresencePenalty:  req.Options.PresencePenalty,
+			FrequencyPenalty: req.Options.FrequencyPenalty,
+			NumPredict:       req.Options.NumPredict,
 		}
 	}

--- a/x/mlxrunner/mlx/ops.go
+++ b/x/mlxrunner/mlx/ops.go
@@ -169,6 +169,12 @@ func (t *Array) PutAlongAxis(indices, values *Array, axis int) *Array {
 	return out
 }

+func (t *Array) ScatterAddAxis(indices, values *Array, axis int) *Array {
+	out := New("SCATTER_ADD_AXIS")
+	C.mlx_scatter_add_axis(&out.ctx, t.ctx, indices.ctx, values.ctx, C.int(axis), DefaultStream().ctx)
+	return out
+}
+
 func (t *Array) Reshape(axes ...int) *Array {
 	cAxes := make([]C.int, len(axes))
 	for i := range axes {
--- a/x/mlxrunner/pipeline.go
+++ b/x/mlxrunner/pipeline.go
@@ -26,16 +26,16 @@ func (r *Runner) TextGenerationPipeline(request Request) error {
 	mlx.ResetPeakMemory()
 	ctx := request.Ctx
 	var (
-		sample, logprobs         *mlx.Array
-		nextSample, nextLogprobs *mlx.Array
+		sample     *mlx.Array
+		nextSample *mlx.Array
 	)

 	defer func() {
 		if request.Sampler != nil {
 			request.Sampler.Free()
 		}
-		mlx.Unpin(sample, logprobs)
-		mlx.Unpin(nextSample, nextLogprobs)
+		mlx.Unpin(sample)
+		mlx.Unpin(nextSample)
 		mlx.Sweep()
 		mlx.ClearCache()

@@ -135,22 +135,21 @@ func (r *Runner) TextGenerationPipeline(request Request) error {
 		mlx.ClearCache()
 	}

-	step := func(token *mlx.Array) (*mlx.Array, *mlx.Array) {
+	step := func(token *mlx.Array) *mlx.Array {
 		fwd := r.Model.Forward(token.ExpandDims(0), caches)
 		logits := r.Model.Unembed(fwd)
 		logits = logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1)

-		logprobs := logits.Subtract(logits.Logsumexp(true))
-		sample := request.Sampler.Sample(logprobs)
+		sample := request.Sampler.Sample(logits)

-		mlx.Pin(sample, logprobs)
+		mlx.Pin(sample)
 		mlx.Sweep()
-		mlx.AsyncEval(sample, logprobs)
+		mlx.AsyncEval(sample)

-		return sample, logprobs
+		return sample
 	}

-	sample, logprobs = step(mlx.FromValues(tokens[processed:], total-processed))
+	sample = step(mlx.FromValues(tokens[processed:], total-processed))

 	var b bytes.Buffer

@@ -161,7 +160,7 @@ func (r *Runner) TextGenerationPipeline(request Request) error {
 		}

 		request.Sampler.AppendToken(sample)
-		nextSample, nextLogprobs = step(sample)
+		nextSample = step(sample)

 		if i == 0 {
 			mlx.Eval(sample)
@@ -186,9 +185,9 @@ func (r *Runner) TextGenerationPipeline(request Request) error {
 		}:
 		}

-		mlx.Unpin(sample, logprobs)
-		sample, logprobs = nextSample, nextLogprobs
-		nextSample, nextLogprobs = nil, nil
+		mlx.Unpin(sample)
+		sample = nextSample
+		nextSample = nil

 		if i%256 == 0 {
 			mlx.ClearCache()
--- a/x/mlxrunner/runner.go
+++ b/x/mlxrunner/runner.go
@@ -31,13 +31,15 @@ type Request struct {
 type TextCompletionsRequest struct {
 	Prompt  string `json:"prompt"`
 	Options struct {
-		Temperature     float32 `json:"temperature"`
-		TopP            float32 `json:"top_p"`
-		MinP            float32 `json:"min_p"`
-		TopK            int     `json:"top_k"`
-		RepeatLastN     int     `json:"repeat_last_n"`
-		PresencePenalty float32 `json:"presence_penalty"`
-		MaxTokens       int     `json:"max_tokens"`
+		Temperature      float32 `json:"temperature"`
+		TopP             float32 `json:"top_p"`
+		MinP             float32 `json:"min_p"`
+		TopK             int     `json:"top_k"`
+		RepeatLastN      int     `json:"repeat_last_n"`
+		RepeatPenalty    float32 `json:"repeat_penalty"`
+		PresencePenalty  float32 `json:"presence_penalty"`
+		FrequencyPenalty float32 `json:"frequency_penalty"`
+		MaxTokens        int     `json:"max_tokens"`

 		// Deprecated: use MaxTokens instead
 		NumPredict int `json:"num_predict"`
--- a/x/mlxrunner/sample/sample.go
+++ b/x/mlxrunner/sample/sample.go
@@ -9,30 +9,38 @@ import (
 type Transform func(*Sampler, *mlx.Array) *mlx.Array

 type Sampler struct {
-	Temperature     float32
-	TopP            float32
-	MinP            float32
-	TopK            int
-	RepeatLastN     int
-	PresencePenalty float32
+	Temperature      float32
+	TopP             float32
+	MinP             float32
+	TopK             int
+	RepeatLastN      int
+	RepeatPenalty    float32
+	PresencePenalty  float32
+	FrequencyPenalty float32

 	history    *mlx.Array
 	historyLen int
 	transforms []Transform
 }

-func New(temp, top_p, min_p float32, top_k, repeatLastN int, presencePenalty float32) *Sampler {
+func New(temp, top_p, min_p float32, top_k, repeatLastN int, repeatPenalty, presencePenalty, frequencyPenalty float32) *Sampler {
+	if repeatPenalty <= 0 {
+		repeatPenalty = 1
+	}
+
 	s := &Sampler{
-		Temperature:     temp,
-		TopP:            top_p,
-		MinP:            min_p,
-		TopK:            top_k,
-		RepeatLastN:     repeatLastN,
-		PresencePenalty: presencePenalty,
+		Temperature:      temp,
+		TopP:             top_p,
+		MinP:             min_p,
+		TopK:             top_k,
+		RepeatLastN:      repeatLastN,
+		RepeatPenalty:    repeatPenalty,
+		PresencePenalty:  presencePenalty,
+		FrequencyPenalty: frequencyPenalty,
 	}

 	var transforms []Transform
-	if presencePenalty != 0 {
+	if s.usesHistory() {
 		transforms = append(transforms, penalty)
 	}

@@ -59,7 +67,7 @@ func New(temp, top_p, min_p float32, top_k, repeatLastN int, presencePenalty flo
 }

 func (s *Sampler) usesHistory() bool {
-	return s.PresencePenalty != 0
+	return s.RepeatPenalty != 1 || s.PresencePenalty != 0 || s.FrequencyPenalty != 0
 }

 func (s *Sampler) setHistory(history *mlx.Array, historyLen int) {
@@ -130,60 +138,78 @@ func temperature(s *Sampler, logits *mlx.Array) *mlx.Array {
 	return mlx.DivScalar(logits, s.Temperature).Categorical(-1)
 }

-func topP(s *Sampler, logprobs *mlx.Array) *mlx.Array {
+func topP(s *Sampler, logits *mlx.Array) *mlx.Array {
 	if s.TopP <= 0 || s.TopP >= 1 {
-		return logprobs
+		return logits
 	}

-	order := logprobs.Negative().ArgsortAxis(-1)
-	sortedLogprobs := logprobs.TakeAlongAxis(order, -1)
-	sortedProbs := mlx.SoftmaxAxis(sortedLogprobs, -1, true)
+	order := logits.Negative().ArgsortAxis(-1)
+	sortedLogits := logits.TakeAlongAxis(order, -1)
+	sortedProbs := mlx.SoftmaxAxis(sortedLogits, -1, true)
 	prevCumProbs := sortedProbs.Cumsum(-1, false, true).Subtract(sortedProbs)
 	keep := prevCumProbs.Less(mlx.FromValue(s.TopP))
-	filtered := mlx.Where(keep, sortedLogprobs, mlx.FromValue(float32(math.Inf(-1))))
-	return logprobs.PutAlongAxis(order, filtered, -1)
+	filtered := mlx.Where(keep, sortedLogits, mlx.FromValue(float32(math.Inf(-1))))
+	return logits.PutAlongAxis(order, filtered, -1)
 }

-func minP(s *Sampler, logprobs *mlx.Array) *mlx.Array {
+func minP(s *Sampler, logits *mlx.Array) *mlx.Array {
 	if s.MinP <= 0 || s.MinP > 1 {
-		return logprobs
+		return logits
 	}

-	maxLogprobs := logprobs.TakeAlongAxis(logprobs.Argmax(-1, true), -1)
-	minLogprobs := mlx.AddScalar(maxLogprobs, float32(math.Log(float64(s.MinP))))
+	maxLogits := logits.TakeAlongAxis(logits.Argmax(-1, true), -1)
+	minLogits := mlx.AddScalar(maxLogits, float32(math.Log(float64(s.MinP))))

 	return mlx.Where(
-		logprobs.Less(minLogprobs),
+		logits.Less(minLogits),
 		mlx.FromValue(float32(math.Inf(-1))),
-		logprobs,
+		logits,
 	)
 }

-func topK(s *Sampler, logprobs *mlx.Array) *mlx.Array {
+func topK(s *Sampler, logits *mlx.Array) *mlx.Array {
 	if s.TopK <= 0 {
-		return logprobs
+		return logits
 	}

-	vocab := logprobs.Dim(logprobs.NumDims() - 1)
+	vocab := logits.Dim(logits.NumDims() - 1)
 	if s.TopK >= vocab {
-		return logprobs
+		return logits
 	}

-	mask := logprobs.Negative().ArgpartitionAxis(s.TopK-1, -1).Slice(mlx.Slice(), mlx.Slice(s.TopK, mlx.End))
-	return logprobs.PutAlongAxis(mask, mlx.FromValue(float32(math.Inf(-1))), -1)
+	mask := logits.Negative().ArgpartitionAxis(s.TopK-1, -1).Slice(mlx.Slice(), mlx.Slice(s.TopK, mlx.End))
+	return logits.PutAlongAxis(mask, mlx.FromValue(float32(math.Inf(-1))), -1)
 }

-func penalty(s *Sampler, logprobs *mlx.Array) *mlx.Array {
-	if s.history == nil || s.historyLen == 0 || s.PresencePenalty == 0 {
-		return logprobs
+func penalty(s *Sampler, logits *mlx.Array) *mlx.Array {
+	if s.historyLen == 0 {
+		return logits
 	}

 	tokenIndices := s.history
-	if logprobs.NumDims() > 1 {
+	if logits.NumDims() > 1 {
 		tokenIndices = tokenIndices.ExpandDims(0)
 	}

-	selected := logprobs.TakeAlongAxis(tokenIndices, -1)
-	adjusted := mlx.AddScalar(selected, -s.PresencePenalty)
-	return logprobs.PutAlongAxis(tokenIndices, adjusted, -1)
+	if s.RepeatPenalty != 1 || s.PresencePenalty != 0 {
+		adjusted := logits.TakeAlongAxis(tokenIndices, -1)
+		if s.RepeatPenalty != 1 {
+			factor := mlx.Where(
+				adjusted.Less(mlx.FromValue(float32(0))),
+				mlx.FromValue(s.RepeatPenalty),
+				mlx.FromValue(1/s.RepeatPenalty),
+			)
+			adjusted = adjusted.Multiply(factor)
+		}
+		if s.PresencePenalty != 0 {
+			adjusted = mlx.AddScalar(adjusted, -s.PresencePenalty)
+		}
+		logits = logits.PutAlongAxis(tokenIndices, adjusted, -1)
+	}
+
+	if s.FrequencyPenalty != 0 {
+		logits = logits.ScatterAddAxis(tokenIndices, mlx.FromValue(-s.FrequencyPenalty), -1)
+	}
+
+	return logits
 }
--- a/x/mlxrunner/sample/sample_test.go
+++ b/x/mlxrunner/sample/sample_test.go
@@ -11,7 +11,7 @@ import (

 func TestPresencePenaltyUsesAppendedTokenImmediately(t *testing.T) {
 	// RepeatLastN = 1, PresencePenalty = 6
-	s := New(0, 0, 0, 0, 1, 6)
+	s := New(0, 0, 0, 0, 1, 1, 6, 0)
 	defer func() {
 		s.Free()
 		mlx.Sweep()
@@ -20,11 +20,11 @@ func TestPresencePenaltyUsesAppendedTokenImmediately(t *testing.T) {
 	s.ResetHistory([]int32{0})
 	s.AppendToken(mlx.NewArrayInt32([]int32{1}, []int32{1}))

-	logprobs := mlx.FromValues([]float32{0, 5, 4}, 3)
-	got := s.Sample(logprobs)
+	logits := mlx.FromValues([]float32{0, 5, 4}, 3)
+	got := s.Sample(logits)
 	mlx.Eval(got)

-	// logprobs will be [0, -1, 4] after the penalty
+	// logits will be [0, -1, 4] after the penalty
 	// and then (index) 2 after the greedy sampler
 	gotInt := got.Int()
 	if gotInt != 2 {
@@ -32,19 +32,59 @@ func TestPresencePenaltyUsesAppendedTokenImmediately(t *testing.T) {
 	}
 }

-func TestMinPMasksTokensBelowThreshold(t *testing.T) {
-	s := New(0, 0, 0.5, 0, 0, 0)
+func TestRepeatPenaltyUsesHistoryWithoutPresencePenalty(t *testing.T) {
+	s := New(0, 0, 0, 0, 1, 2, 0, 0)
 	defer func() {
 		s.Free()
 		mlx.Sweep()
 	}()

-	logprobs := mlx.FromValues([]float32{
+	s.ResetHistory([]int32{1})
+
+	logits := mlx.FromValues([]float32{0, 5, 4}, 3)
+	got := s.Sample(logits)
+	mlx.Eval(got)
+
+	// token 1 is repeated and positive, so 5 / 2 falls below token 2.
+	gotInt := got.Int()
+	if gotInt != 2 {
+		t.Fatalf("got %d, want 2", gotInt)
+	}
+}
+
+func TestFrequencyPenaltyUsesTokenCounts(t *testing.T) {
+	s := New(0, 0, 0, 0, 4, 1, 0, 2)
+	defer func() {
+		s.Free()
+		mlx.Sweep()
+	}()
+
+	s.ResetHistory([]int32{1, 1})
+
+	logits := mlx.FromValues([]float32{0, 5, 4}, 3)
+	got := s.Sample(logits)
+	mlx.Eval(got)
+
+	// token 1 appears twice, so 5 - (2 * 2) falls below token 2.
+	gotInt := got.Int()
+	if gotInt != 2 {
+		t.Fatalf("got %d, want 2", gotInt)
+	}
+}
+
+func TestMinPMasksTokensBelowThreshold(t *testing.T) {
+	s := New(0, 0, 0.5, 0, 0, 1, 0, 0)
+	defer func() {
+		s.Free()
+		mlx.Sweep()
+	}()
+
+	logits := mlx.FromValues([]float32{
 		float32(math.Log(0.5)),
 		float32(math.Log(0.3)),
 		float32(math.Log(0.2)),
 	}, 3)
-	got := minP(s, logprobs)
+	got := minP(s, logits)
 	mlx.Eval(got)

 	gotFloats := got.Floats()
--- a/x/mlxrunner/server.go
+++ b/x/mlxrunner/server.go
@@ -102,7 +102,9 @@ func Execute(args []string) error {
 			request.Options.MinP,
 			request.Options.TopK,
 			request.Options.RepeatLastN,
+			request.Options.RepeatPenalty,
 			request.Options.PresencePenalty,
+			request.Options.FrequencyPenalty,
 		)

 		var cancel context.CancelFunc