diff --git a/mlxrunner/model/clef/encode.go b/mlxrunner/model/clef/encode.go index 0d9c9a015..3beef81a7 100644 --- a/mlxrunner/model/clef/encode.go +++ b/mlxrunner/model/clef/encode.go @@ -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...) diff --git a/mlxrunner/model/clef/encode_test.go b/mlxrunner/model/clef/encode_test.go index 5d841c758..ca55dd938 100644 --- a/mlxrunner/model/clef/encode_test.go +++ b/mlxrunner/model/clef/encode_test.go @@ -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) + } + }) } } diff --git a/mlxrunner/model/clef/score.go b/mlxrunner/model/clef/score.go index 133623815..6a4fc690b 100644 --- a/mlxrunner/model/clef/score.go +++ b/mlxrunner/model/clef/score.go @@ -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()} diff --git a/mlxrunner/model/laya/score.go b/mlxrunner/model/laya/score.go index d9506d182..67b464a05 100644 --- a/mlxrunner/model/laya/score.go +++ b/mlxrunner/model/laya/score.go @@ -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) diff --git a/mlxrunner/model/laya/score_test.go b/mlxrunner/model/laya/score_test.go index 7047f036d..5a9815e08 100644 --- a/mlxrunner/model/laya/score_test.go +++ b/mlxrunner/model/laya/score_test.go @@ -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 { diff --git a/server/routes.go b/server/routes.go index e7a70afe9..73b9100f1 100644 --- a/server/routes.go +++ b/server/routes.go @@ -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