diff --git a/integration/context_test.go b/integration/context_test.go index 5c2549c86..82f932997 100644 --- a/integration/context_test.go +++ b/integration/context_test.go @@ -17,9 +17,10 @@ import ( func TestLongInputContext(t *testing.T) { // Setting NUM_PARALLEL to 1 ensures the allocated context is exactly what // we asked for and there is nothing extra that we could spill over into. - // Older runners silently truncate oversized prompts, while llama-server - // rejects them with a client error. Accept both behaviors so this test can - // run against main and the llama-server branch. + // Context shift happens after a prompt has been admitted to a slot. Initial + // prompts that fill or exceed the slot are still rejected by llama-server. + // Accept a context-limit error here because older runners may truncate this + // prompt while llama-server reports it as too large to admit. t.Setenv("OLLAMA_NUM_PARALLEL", "1") ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) @@ -74,6 +75,7 @@ func isContextLimitError(err string) bool { return strings.Contains(err, "context") && (strings.Contains(err, "exceed") || strings.Contains(err, "too large") || + strings.Contains(err, "longer") || strings.Contains(err, "too long")) } diff --git a/llm/llama_server.go b/llm/llama_server.go index 0be725c63..7cf70cf28 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -273,15 +273,31 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req return nil, err } - // s.options.NumCtx is the runner's effective per-slot context length. It - // has already been capped by the model's training context during launch. - if len(tokens) >= s.options.NumCtx { + limit := s.options.NumCtx - 1 + if len(tokens) <= limit { + return prompt, nil + } + + if !s.launch.config.ContextShift { return nil, api.StatusError{ StatusCode: http.StatusBadRequest, - ErrorMessage: "the prompt is longer than the context length currently available to the model; shorten the prompt or adjust the context length in settings", + ErrorMessage: "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length", } } - return prompt, nil + + nKeep := req.Options.NumKeep + if nKeep < 0 { + nKeep = len(tokens) + } + nKeep = min(nKeep, limit) + + discard := len(tokens) - limit + truncated := make([]int, 0, limit) + truncated = append(truncated, tokens[:nKeep]...) + truncated = append(truncated, tokens[nKeep+discard:]...) + + slog.Warn("truncating input prompt", "limit", limit, "prompt", len(tokens), "keep", nKeep, "new", len(truncated)) + return truncated, nil } func (s *llamaServerRunner) ContextLength() int { diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index def5fca2e..93094e4d6 100644 --- a/llm/llama_server_test.go +++ b/llm/llama_server_test.go @@ -489,7 +489,7 @@ func TestLlamaServerCompletionForwardsRepeatLastNZero(t *testing.T) { } func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) { - const wantError = "the prompt is longer than the context length currently available to the model; shorten the prompt or adjust the context length in settings" + const wantError = "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length" var tokenizeReq struct { Content string `json:"content"` @@ -558,6 +558,160 @@ func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) { } } +func TestLlamaServerCompletionContextShiftAllowsPromptWithHeadroom(t *testing.T) { + var capturedReq llamaServerCompletionRequest + var tokenizeReq struct { + Content string `json:"content"` + AddSpecial bool `json:"add_special"` + ParseSpecial *bool `json:"parse_special"` + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/health": + fmt.Fprint(w, `{"status":"ok"}`) + case "/tokenize": + if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil { + t.Errorf("invalid tokenize request body: %v", err) + return + } + fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6]}`) + case "/completion": + if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil { + t.Errorf("invalid completion request body: %v", err) + return + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":10,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + })) + defer srv.Close() + + parts := strings.Split(srv.URL, ":") + var portInt int + fmt.Sscanf(parts[len(parts)-1], "%d", &portInt) + + runner := &llamaServerRunner{ + port: portInt, + cmd: fakeRunningCmd(), + sem: semaphore.NewWeighted(1), + options: api.Options{Runner: api.Runner{NumCtx: 8}}, + launch: llamaServerLaunchConfig{ + config: LlamaServerConfig{ContextShift: true}, + }, + } + + opts := api.DefaultOptions() + opts.NumKeep = 3 + prompt := strings.Repeat("long prompt ", 2) + err := runner.Completion(t.Context(), CompletionRequest{ + Prompt: prompt, + Options: &opts, + Truncate: true, + }, func(cr CompletionResponse) {}) + if err != nil { + t.Fatalf("Completion error: %v", err) + } + + if tokenizeReq.Content != prompt { + t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt) + } + if !tokenizeReq.AddSpecial { + t.Fatal("expected tokenize request to add special tokens") + } + if capturedReq.Prompt != prompt { + t.Fatalf("prompt = %q, want %q", capturedReq.Prompt, prompt) + } + if capturedReq.NKeep != opts.NumKeep { + t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep) + } +} + +func TestLlamaServerCompletionContextShiftTruncatesPromptOverContext(t *testing.T) { + var capturedReq llamaServerCompletionRequest + var tokenizeReq struct { + Content string `json:"content"` + AddSpecial bool `json:"add_special"` + ParseSpecial *bool `json:"parse_special"` + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/health": + fmt.Fprint(w, `{"status":"ok"}`) + case "/tokenize": + if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil { + t.Errorf("invalid tokenize request body: %v", err) + return + } + fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`) + case "/completion": + if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil { + t.Errorf("invalid completion request body: %v", err) + return + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":7,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + })) + defer srv.Close() + + parts := strings.Split(srv.URL, ":") + var portInt int + fmt.Sscanf(parts[len(parts)-1], "%d", &portInt) + + runner := &llamaServerRunner{ + port: portInt, + cmd: fakeRunningCmd(), + sem: semaphore.NewWeighted(1), + options: api.Options{Runner: api.Runner{NumCtx: 8}}, + launch: llamaServerLaunchConfig{ + config: LlamaServerConfig{ContextShift: true}, + }, + } + + opts := api.DefaultOptions() + opts.NumKeep = 3 + prompt := strings.Repeat("long prompt ", 2) + err := runner.Completion(t.Context(), CompletionRequest{ + Prompt: prompt, + Options: &opts, + Truncate: true, + }, func(cr CompletionResponse) {}) + if err != nil { + t.Fatalf("Completion error: %v", err) + } + + if tokenizeReq.Content != prompt { + t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt) + } + if !tokenizeReq.AddSpecial { + t.Fatal("expected tokenize request to add special tokens") + } + + got, ok := capturedReq.Prompt.([]any) + if !ok { + t.Fatalf("completion prompt = %T, want token array", capturedReq.Prompt) + } + want := []int{0, 1, 2, 6, 7, 8, 9} + if len(got) != len(want) { + t.Fatalf("token prompt len = %d, want %d: %#v", len(got), len(want), got) + } + for i, wantToken := range want { + gotToken, ok := got[i].(float64) + if !ok || int(gotToken) != wantToken { + t.Fatalf("token prompt[%d] = %#v, want %d", i, got[i], wantToken) + } + } + if capturedReq.NKeep != opts.NumKeep { + t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep) + } +} + func TestLlamaServerCompletionWithMediaUsesRunnerMarker(t *testing.T) { var capturedReq llamaServerCompletionRequest diff --git a/server/sched.go b/server/sched.go index 58c7c2d67..980ad3584 100644 --- a/server/sched.go +++ b/server/sched.go @@ -132,8 +132,24 @@ func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, ses return s.getRunner(c, m, opts, sessionDuration, false, false, nil) } -func resolveContextShift(shift *bool) bool { - return shift == nil || *shift +func resolveContextShift(shift *bool, m *Model) bool { + if shift != nil { + return *shift + } + + return supportsContextShift(m) +} + +func supportsContextShift(m *Model) bool { + if m == nil { + return true + } + + if m.Config.ModelFamily == "deepseek2" || slices.Contains(m.Config.ModelFamilies, "deepseek2") { + return false + } + + return true } func effectiveModelContext(numCtx int, f *ggml.GGML) int { @@ -168,7 +184,7 @@ func (s *Scheduler) getRunner(c context.Context, m *Model, opts api.Options, ses contextShift := false if m.ModelPath != "" { - contextShift = resolveContextShift(shift) + contextShift = resolveContextShift(shift, m) } req := &LlmRequest{ @@ -560,7 +576,7 @@ func (s *Scheduler) load(req *LlmRequest, systemInfo ml.SystemInfo, gpus []ml.De } launchOpts = s.applyLlamaServerMmapDefaults(req, launchOpts, systemInfo, loadGpus, f, numParallel) - req.contextShift = resolveContextShift(req.shift) + req.contextShift = resolveContextShift(req.shift, req.model) config := llamaServerConfigForModel(req.model) config.ContextShift = req.contextShift @@ -685,7 +701,7 @@ iGPUScan: trainContext := modelTrainContext(f) if effectiveNumCtx := llama.ContextLength(); req.model.ModelPath != "" && effectiveNumCtx > 0 { req.opts.NumCtx = effectiveNumCtx - req.contextShift = resolveContextShift(req.shift) + req.contextShift = resolveContextShift(req.shift, req.model) } runner := &runnerRef{ model: req.model, @@ -1414,7 +1430,7 @@ func (runner *runnerRef) needsReload(ctx context.Context, req *LlmRequest) bool contextShift := req.contextShift if req.model.ModelPath != "" { - contextShift = resolveContextShift(req.shift) + contextShift = resolveContextShift(req.shift, req.model) } if runner.contextShift != contextShift { return true diff --git a/server/sched_test.go b/server/sched_test.go index e26a66c45..23a8acd98 100644 --- a/server/sched_test.go +++ b/server/sched_test.go @@ -907,16 +907,21 @@ func TestResolveContextShift(t *testing.T) { tests := []struct { name string shift *bool + model *Model want bool }{ {name: "unset defaults to shift", want: true}, + {name: "unset deepseek2 disables shift", model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false}, + {name: "unset deepseek2 family disables shift", model: &Model{Config: model.ConfigV2{ModelFamilies: []string{"llama", "deepseek2"}}}, want: false}, {name: "explicit false disables shift", shift: &falseValue, want: false}, + {name: "explicit false disables shift for deepseek2", shift: &falseValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false}, {name: "explicit true enables shift", shift: &trueValue, want: true}, + {name: "explicit true enables shift for deepseek2", shift: &trueValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - require.Equal(t, tt.want, resolveContextShift(tt.shift)) + require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.model)) }) } }