llm: context shift allow shiftable prompts (#16764)

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"))
 }
 

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 {

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
 

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

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))
 		})
 	}
 }