CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-10-21 11:01:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

mlx: Decision model improvements (#18776)

Browse source 
* reject instead of truncate on overflow

* tune route logic for faster warm runner lookup
This commit is contained in:
Daniel Hiltgen 2026-10-03 17:48:08 -07:00 • committed by GitHub
parent a00d083bb5
commit 27a36513d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 156 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

13
mlxrunner/model/clef/encode.go
View file

@ -21,7 +21,7 @@ type encodedRecord struct {
}
// Clef's shared encoder supplies prefix, state, then schema and suffix segments.
// The prepared prefix may include images; only state tokens may be truncated.
// The prepared prefix includes image tokens, which count toward the context limit.
func encode(tok *tokenizer.Tokenizer, input llm.ScoreRequest, prefix []int32) (encodedRecord, error) {
var record encodedRecord
if len(input.Segments) < 3 {
@ -29,17 +29,14 @@ func encode(tok *tokenizer.Tokenizer, input llm.ScoreRequest, prefix []int32) (e
}
tokens := make([][]int32, len(input.Segments))
tokens[0] = prefix
fixed := len(prefix)
total := len(prefix)
for i := 1; i < len(tokens); i++ {
tokens[i] = tok.Encode(input.Segments[i], false)
if i != 1 {
fixed += len(tokens[i])
}
total += len(tokens[i])
}
if fixed > input.MaxTokens {
return record, fmt.Errorf("Clef schema exceeds the %d-token context", input.MaxTokens)
if total > input.MaxTokens {
return record, fmt.Errorf("Clef input has %d tokens; context limit is %d", total, input.MaxTokens)
}
tokens[1] = tokens[1][:min(len(tokens[1]), input.MaxTokens-fixed)]
offsets := make([]int, len(tokens)+1)
for i, segment := range tokens {
record.ids = append(record.ids, segment...)

27
mlxrunner/model/clef/encode_test.go
View file

@ -9,16 +9,16 @@ import (
"github.com/ollama/ollama/mlxrunner/tokenizer"
)
func TestEncodeTruncatesOnlyState(t *testing.T) {
func TestEncodeContextLimit(t *testing.T) {
tok, err := tokenizer.LoadFromBytes([]byte(`{"model":{"type":"BPE","vocab":{"x":0},"merges":[]}}`))
if err != nil {
t.Fatal(err)
}
// Fixed prompt tokens must survive truncation along with the schema.
// The prefix includes the expanded image tokens, not just the text prefix.
prefix := []int32{100, 101, 101, 101, 102}
input := llm.ScoreRequest{
MaxTokens: 12,
Segments: []string{"prefix", strings.Repeat("x", 100), "xx", "x", "x"},
Segments: []string{"prefix", "xxx", "xx", "x", "x"},
Fields: []llm.ScoreField{{Type: 0, Question: [2]int{2, 3}, Options: [][2]int{{3, 4}, {4, 5}}}},
}
got, err := encode(tok, input, prefix)
@ -32,9 +32,24 @@ func TestEncodeTruncatesOnlyState(t *testing.T) {
if q.instruction != (tokenSpan{8, 10}) || !slices.Equal(q.options, []tokenSpan{{10, 11}, {11, 12}}) {
t.Fatalf("schema spans shifted incorrectly: %+v", q)
}
input.MaxTokens = 8
if _, err := encode(tok, input, prefix); err == nil {
t.Fatal("prefix and schema exceeding the context must be rejected")
for _, tc := range []struct {
name string
state string
prefix []int32
limit int
}{
{"state one token over", "xxxx", prefix, 12},
{"long state", strings.Repeat("x", 100), prefix, 12},
{"prefix and schema over", "", prefix, 8},
{"expanded image prefix over", "xxx", []int32{100, 101, 101, 101, 101, 102}, 12},
} {
t.Run(tc.name, func(t *testing.T) {
input.Segments[1], input.MaxTokens = tc.state, tc.limit
got, err := encode(tok, input, tc.prefix)
if err == nil || !strings.Contains(err.Error(), "context limit is") || len(got.ids) != 0 || len(got.questions) != 0 {
t.Fatalf("encode = %+v, %v; want context rejection without a partial record", got, err)
}
})
}
}

1
mlxrunner/model/clef/score.go
View file

@ -48,7 +48,6 @@ func (m *Model) PrepareScore(ctx context.Context, input llm.ScoreRequest) (*Scor
if err != nil {
return nil, api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: err.Error()}
}
// Reserve the entire image prefix before truncating only the state.
record, err := encode(m.Tokenizer(), input, prepared.Tokens)
if err != nil {
return nil, api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: err.Error()}

49
mlxrunner/model/laya/score.go
View file

@ -37,24 +37,34 @@ func (m *Model) sequence(state string, q llm.ScoreQuestion, maxLen int) (ids, ma
if len(q.Options) < 2 || len(q.Options) > 26 {
return nil, nil, 0, fmt.Errorf("decision requires 2–26 options")
}
encode := func(s string) []int32 { return m.tok.Encode(strings.ReplaceAll(s, m.maskToken, " "), false) }
head := encode(q.Type + " question: " + q.Instructions)
if strings.Contains(state, m.maskToken) {
return nil, nil, 0, fmt.Errorf("state contains reserved token %q", m.maskToken)
}
if strings.Contains(q.Instructions, m.maskToken) {
return nil, nil, 0, fmt.Errorf("instructions contain reserved token %q", m.maskToken)
}
head := m.tok.Encode(q.Type+" question: "+q.Instructions, false)
opts := make([][]int32, len(q.Options))
budget := m.config.HeadMaxLen
for i, text := range q.Options {
tokens := encode(" " + text)
opts[i] = append([]int32{m.mask}, tokens[:min(48, len(tokens))]...)
if strings.Contains(text, m.maskToken) {
return nil, nil, 0, fmt.Errorf("option %d contains reserved token %q", i, m.maskToken)
}
tokens := m.tok.Encode(" "+text, false)
if len(tokens) > 48 {
return nil, nil, 0, fmt.Errorf("option %d has %d tokens; limit is 48", i, len(tokens))
}
opts[i] = append([]int32{m.mask}, tokens...)
budget -= len(opts[i])
}
// Keep the publisher's head limits, but reject inputs that would be clipped.
if budget < 16 {
per := max(4, (m.config.HeadMaxLen-16)/len(opts))
budget = m.config.HeadMaxLen
for i := range opts {
opts[i] = opts[i][:min(per, len(opts[i]))]
budget -= len(opts[i])
}
return nil, nil, 0, fmt.Errorf("decision options exceed the %d-token budget", m.config.HeadMaxLen-16)
}
ids = append([]int32{m.cls}, head[:min(len(head), max(8, budget))]...)
if len(head) > budget {
return nil, nil, 0, fmt.Errorf("instructions with the question type have %d tokens; limit is %d with these options", len(head), budget)
}
ids = append([]int32{m.cls}, head...)
ids = append(ids, m.sep)
for _, tokens := range opts {
markers = append(markers, int32(len(ids)))
@ -62,12 +72,13 @@ func (m *Model) sequence(state string, q llm.ScoreQuestion, maxLen int) (ids, ma
}
ids = append(ids, m.sep)
if len(ids)+1 > maxLen {
return nil, nil, 0, fmt.Errorf("decision options exceed the %d-token context", maxLen)
return nil, nil, 0, fmt.Errorf("decision question requires %d tokens; context limit is %d", len(ids)+1, maxLen)
}
stateTokens := encode(state)
// Match the publisher's question/state budget: options and instructions
// have their own cap, then state is truncated on the right to fit.
ids = append(ids, stateTokens[:min(len(stateTokens), maxLen-len(ids)-1)]...)
stateTokens := m.tok.Encode(state, false)
if limit := maxLen - len(ids) - 1; len(stateTokens) > limit {
return nil, nil, 0, fmt.Errorf("state has %d tokens; limit is %d with this question", len(stateTokens), limit)
}
ids = append(ids, stateTokens...)
ids = append(ids, m.sep)
return ids, markers, qtype, nil
}
@ -158,7 +169,11 @@ func (m *Model) Score(ctx context.Context, input llm.ScoreRequest) (llm.ScoreRes
SeqQueryLens: []int32{int32(len(row.Tokens))},
Layout: []any{row.Type},
}
hidden, _ := m.Forward(b, nil)
// Release forward intermediates before evaluating the decision head.
hidden := mlx.ScopedArrays(func() []*mlx.Array {
h, _ := m.Forward(b, nil)
return []*mlx.Array{h}
})[0]
logits = m.FinishScore(row, hidden)
})
result.Logits = append(result.Logits, logits)

79
mlxrunner/model/laya/score_test.go
View file

@ -5,12 +5,91 @@ import (
"errors"
"math"
"net/http"
"slices"
"strings"
"testing"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/llm"
"github.com/ollama/ollama/mlxrunner/tokenizer"
)
func TestPrepareScoreInputLimits(t *testing.T) {
// One token per character makes the field and context boundaries explicit.
tok, err := tokenizer.LoadFromBytes([]byte(`{"model":{"type":"BPE","vocab":{"x":0,"Ġ":1,"c":2,"h":3,"o":4,"i":5,"e":6,"q":7,"u":8,"s":9,"t":10,"n":11,":":12},"merges":[]}}`))
if err != nil {
t.Fatal(err)
}
m := Model{config: config{MaxLen: 512, HeadMaxLen: 192}, tok: tok, cls: 100, sep: 101, mask: 102, maskToken: "[MASK]"}
for _, tc := range []struct {
name string
state string
instructions string
options []string
limit int
wantTokens int
wantError string
}{
{"context boundary", strings.Repeat("x", 36), "x", []string{"x", "x"}, 64, 64, ""},
{"state overflow", strings.Repeat("x", 37), "x", []string{"x", "x"}, 64, 0, "state has 37 tokens; limit is 36"},
{"question overflow", "", "x", []string{"x", "x"}, 27, 0, "decision question requires 28 tokens; context limit is 27"},
{"option boundary", "x", "x", []string{strings.Repeat("x", 47), "x"}, 512, 75, ""},
{"option overflow", "x", "x", []string{strings.Repeat("x", 48), "x"}, 512, 0, "option 0 has 49 tokens; limit is 48"},
{"options and instructions fit", "x", "", []string{strings.Repeat("x", 42), strings.Repeat("x", 42), strings.Repeat("x", 42), strings.Repeat("x", 41)}, 512, 197, ""},
{"options leave too little instruction room", "x", "", []string{strings.Repeat("x", 42), strings.Repeat("x", 42), strings.Repeat("x", 42), strings.Repeat("x", 42)}, 512, 0, "instructions with the question type have 17 tokens; limit is 16"},
{"options budget overflow", "x", "", []string{strings.Repeat("x", 43), strings.Repeat("x", 42), strings.Repeat("x", 42), strings.Repeat("x", 42)}, 512, 0, "decision options exceed the 176-token budget"},
{"instruction boundary", "x", strings.Repeat("x", 169), []string{"x", "x"}, 512, 197, ""},
{"instruction overflow", "x", strings.Repeat("x", 170), []string{"x", "x"}, 512, 0, "instructions with the question type have 187 tokens; limit is 186"},
{"reserved state", "x[MASK]x", "x", []string{"x", "x"}, 512, 0, "state contains reserved token"},
{"reserved instructions", "x", "x[MASK]x", []string{"x", "x"}, 512, 0, "instructions contain reserved token"},
{"reserved option", "x", "x", []string{"x", "x[MASK]x"}, 512, 0, "option 1 contains reserved token"},
} {
t.Run(tc.name, func(t *testing.T) {
input := llm.ScoreRequest{State: tc.state, MaxTokens: tc.limit, Rows: []llm.ScoreRow{{Question: &llm.ScoreQuestion{Type: "choice", Instructions: tc.instructions, Options: tc.options}}}}
rows, err := m.PrepareScore(t.Context(), input)
if tc.wantError != "" {
var status api.StatusError
if !errors.As(err, &status) || status.StatusCode != http.StatusBadRequest || !strings.Contains(status.ErrorMessage, "question 0: "+tc.wantError) || rows != nil {
t.Fatalf("rows=%v error=%v, want HTTP 400 containing %q and no prepared rows", rows, err, tc.wantError)
}
return
}
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 || len(rows[0].Tokens) != tc.wantTokens || len(rows[0].Markers) != len(tc.options) {
t.Fatalf("prepared rows=%+v, want one %d-token row with %d options", rows, tc.wantTokens, len(tc.options))
}
row := rows[0]
if row.Tokens[0] != m.cls || row.Tokens[len(row.Tokens)-1] != m.sep {
t.Fatalf("missing sequence delimiters: %v", row.Tokens)
}
})
}
// A late invalid question must fail before Score can forward even the first
// valid row. This model has no weights, so any inference attempt also fails.
input := llm.ScoreRequest{State: "x", MaxTokens: 512, Rows: []llm.ScoreRow{
{Question: &llm.ScoreQuestion{Type: "choice", Options: []string{"x", "x"}}},
{Question: &llm.ScoreQuestion{Type: "choice", Options: []string{"x", strings.Repeat("x", 48)}}},
}}
result, err := m.Score(t.Context(), input)
var status api.StatusError
if !errors.As(err, &status) || status.StatusCode != http.StatusBadRequest || !strings.Contains(status.ErrorMessage, "question 1: option 1") || len(result.Logits) != 0 || result.InputTokens != 0 {
t.Fatalf("Score returned result=%+v error=%v, want rejection before inference", result, err)
}
// An ordinary request retains its complete input and marker positions.
rows, err := m.PrepareScore(t.Context(), llm.ScoreRequest{State: "xx", MaxTokens: 512, Rows: input.Rows[:1]})
if err != nil {
t.Fatal(err)
}
want := []int32{100, 2, 3, 4, 5, 2, 6, 1, 7, 8, 6, 9, 10, 5, 4, 11, 12, 1, 101, 102, 1, 0, 102, 1, 0, 101, 0, 0, 101}
if !slices.Equal(rows[0].Tokens, want) || !slices.Equal(rows[0].Markers, []int32{19, 22}) {
t.Fatalf("prepared row=%+v, want tokens=%v markers=[19 22]", rows[0], want)
}
}
func TestTemperature(t *testing.T) {
m := Model{config: config{Temperature: []float32{2, 3, 4}, TemperatureByOptions: map[string]float32{"choice:2": 1.5, "choice:11+": 0.1, "score:3-5": 7, "noul:2": float32(math.NaN())}}}
for _, tc := range []struct {

33
server/routes.go
View file

@ -242,15 +242,20 @@ func (s *Server) scheduleRunner(ctx context.Context, name, selectedRunner string
}
name = existingName.String()
}
selectedName := model.ParseName(name)
model, err := GetModelForRunner(name, selectedRunner)
if err != nil {
return nil, nil, nil, err
}
return s.scheduleRunnerForModel(ctx, model, caps, requestOpts, keepAlive, shift)
}
// scheduleRunnerForModel retains the request's resolved model, including its selected
// manifest-list child, without repeating model lookup and metadata loading.
func (s *Server) scheduleRunnerForModel(ctx context.Context, m *Model, caps []model.Capability, requestOpts map[string]any, keepAlive *api.Duration, shift *bool) (llm.LlamaServer, *Model, *api.Options, error) {
selectedName := model.ParseName(m.Name)
if !manifest.IsDigestReferenceName(selectedName) {
runner := model.Runner
if runner == "" && model.isGGUF() {
runner := m.Runner
if runner == "" && m.isGGUF() {
runner = manifest.RunnerGGML
}
if runner == manifest.RunnerGGML {
@ -258,23 +263,23 @@ func (s *Server) scheduleRunner(ctx context.Context, name, selectedRunner string
}
}
if slices.Contains(model.Config.ModelFamilies, "mllama") && len(model.ProjectorPaths) > 0 {
if slices.Contains(m.Config.ModelFamilies, "mllama") && len(m.ProjectorPaths) > 0 {
return nil, nil, nil, fmt.Errorf("'llama3.2-vision' is no longer compatible with your version of Ollama and has been replaced by a newer version. To re-download, run 'ollama pull llama3.2-vision'")
}
if err := model.CheckCapabilities(caps...); err != nil {
return nil, nil, nil, fmt.Errorf("%s %w", model.Name, err)
if err := m.CheckCapabilities(caps...); err != nil {
return nil, nil, nil, fmt.Errorf("%s %w", m.Name, err)
}
numCtxAuto := usesAutomaticNumCtx(model, requestOpts)
embeddingBatchDefault := shouldApplyEmbeddingBatchDefault(model, requestOpts)
numBatchAuto := usesAutomaticNumBatch(model, requestOpts) && !embeddingBatchDefault
opts, err := s.modelOptionsWithEmbeddingBatchDefault(model, requestOpts, embeddingBatchDefault)
numCtxAuto := usesAutomaticNumCtx(m, requestOpts)
embeddingBatchDefault := shouldApplyEmbeddingBatchDefault(m, requestOpts)
numBatchAuto := usesAutomaticNumBatch(m, requestOpts) && !embeddingBatchDefault
opts, err := s.modelOptionsWithEmbeddingBatchDefault(m, requestOpts, embeddingBatchDefault)
if err != nil {
return nil, nil, nil, err
}
runnerCh, errCh := s.sched.getRunner(ctx, model, opts, keepAlive, numCtxAuto, numBatchAuto, shift)
runnerCh, errCh := s.sched.getRunner(ctx, m, opts, keepAlive, numCtxAuto, numBatchAuto, shift)
var runner *runnerRef
select {
case runner = <-runnerCh:
@ -282,7 +287,7 @@ func (s *Server) scheduleRunner(ctx context.Context, name, selectedRunner string
return nil, nil, nil, err
}
return runner.llama, model, &opts, nil
return runner.llama, m, &opts, nil
}
func signinURL() (string, error) {
@ -958,7 +963,7 @@ func (s *Server) SystemOneHandler(c *gin.Context) {
if len(req.Images) > 0 {
caps = append(caps, model.CapabilityVision)
}
r, _, _, err := s.scheduleRunner(c.Request.Context(), name.String(), "", caps, nil, req.KeepAlive, nil)
r, _, _, err := s.scheduleRunnerForModel(c.Request.Context(), m, caps, nil, req.KeepAlive, nil)
if err != nil {
handleScheduleError(c, req.Model, err)
return