server: context shift for context windows larger than 8k, add error when hitting context limit (#16712)

llm/llama_server.go

@@ -273,28 +273,15 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req
 		return nil, err
 	}
 
-	// llama-server rejects prompts that fill the entire slot context, while the
-	// old runner could accept exactly num_ctx prompt tokens. Keep one token of
-	// headroom so token-level truncation preserves old behavior as closely as
-	// llama-server allows.
-	limit := s.options.NumCtx - 1
-	if len(tokens) <= limit {
-		return prompt, nil
+	// 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 {
+		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",
+		}
 	}
-
-	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", s.options.NumCtx, "prompt", len(tokens), "keep", nKeep, "new", len(truncated))
-	return truncated, nil
+	return prompt, nil
 }
 
 func (s *llamaServerRunner) ContextLength() int {

llm/llama_server_test.go

@@ -488,14 +488,16 @@ func TestLlamaServerCompletionForwardsRepeatLastNZero(t *testing.T) {
 	}
 }
 
-func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
-	var completionReq llamaServerCompletionRequest
+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"
+
 	var tokenizeReq struct {
 		Content      string `json:"content"`
 		AddSpecial   bool   `json:"add_special"`
 		ParseSpecial *bool  `json:"parse_special"`
 	}
 
+	completionCalled := false
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/health":
@@ -507,10 +509,7 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 			}
 			fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`)
 		case "/completion":
-			if err := json.NewDecoder(r.Body).Decode(&completionReq); err != nil {
-				t.Errorf("invalid completion request body: %v", err)
-				return
-			}
+			completionCalled = true
 			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:
@@ -537,8 +536,15 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 		Options:  &opts,
 		Truncate: true,
 	}, func(cr CompletionResponse) {})
-	if err != nil {
-		t.Fatalf("Completion error: %v", err)
+	var statusErr api.StatusError
+	if !errors.As(err, &statusErr) {
+		t.Fatalf("Completion error = %T %v, want api.StatusError", err, err)
+	}
+	if statusErr.StatusCode != http.StatusBadRequest {
+		t.Fatalf("StatusCode = %d, want %d", statusErr.StatusCode, http.StatusBadRequest)
+	}
+	if statusErr.ErrorMessage != wantError {
+		t.Fatalf("ErrorMessage = %q, want %q", statusErr.ErrorMessage, wantError)
 	}
 
 	if tokenizeReq.Content != strings.Repeat("long prompt ", 2) {
@@ -547,20 +553,8 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 	if !tokenizeReq.AddSpecial {
 		t.Fatal("expected tokenize request to add special tokens")
 	}
-
-	got, ok := completionReq.Prompt.([]any)
-	if !ok {
-		t.Fatalf("completion prompt = %T, want token array", completionReq.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 completionCalled {
+		t.Fatal("completion endpoint was called")
 	}
 }
 

server/sched.go

@@ -132,14 +132,8 @@ func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, ses
 	return s.getRunner(c, m, opts, sessionDuration, false, false, nil)
 }
 
-const contextShiftSmallContextLimit = 8192
-
-func resolveContextShift(shift *bool, numCtx int) bool {
-	if shift != nil {
-		return *shift
-	}
-
-	return numCtx > 0 && numCtx < contextShiftSmallContextLimit
+func resolveContextShift(shift *bool) bool {
+	return shift == nil || *shift
 }
 
 func effectiveModelContext(numCtx int, f *ggml.GGML) int {
@@ -174,7 +168,7 @@ func (s *Scheduler) getRunner(c context.Context, m *Model, opts api.Options, ses
 
 	contextShift := false
 	if m.ModelPath != "" {
-		contextShift = resolveContextShift(shift, opts.NumCtx)
+		contextShift = resolveContextShift(shift)
 	}
 
 	req := &LlmRequest{
@@ -566,7 +560,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, effectiveModelContext(launchOpts.NumCtx, f))
+			req.contextShift = resolveContextShift(req.shift)
 
 			config := llamaServerConfigForModel(req.model)
 			config.ContextShift = req.contextShift
@@ -691,7 +685,7 @@ iGPUScan:
 	trainContext := modelTrainContext(f)
 	if effectiveNumCtx := llama.ContextLength(); req.model.ModelPath != "" && effectiveNumCtx > 0 {
 		req.opts.NumCtx = effectiveNumCtx
-		req.contextShift = resolveContextShift(req.shift, effectiveNumCtx)
+		req.contextShift = resolveContextShift(req.shift)
 	}
 	runner := &runnerRef{
 		model:           req.model,
@@ -1420,7 +1414,7 @@ func (runner *runnerRef) needsReload(ctx context.Context, req *LlmRequest) bool
 
 	contextShift := req.contextShift
 	if req.model.ModelPath != "" {
-		contextShift = resolveContextShift(req.shift, optsNew.NumCtx)
+		contextShift = resolveContextShift(req.shift)
 	}
 	if runner.contextShift != contextShift {
 		return true

server/sched_test.go

@@ -907,19 +907,16 @@ func TestResolveContextShift(t *testing.T) {
 	tests := []struct {
 		name  string
 		shift *bool
-		ctx   int
 		want  bool
 	}{
-		{name: "unset small context keeps legacy shift", ctx: 128, want: true},
-		{name: "unset large context disables shift", ctx: contextShiftSmallContextLimit, want: false},
-		{name: "unset invalid context disables shift", ctx: 0, want: false},
-		{name: "explicit false wins for small context", shift: &falseValue, ctx: 128, want: false},
-		{name: "explicit true wins for large context", shift: &trueValue, ctx: 32768, want: true},
+		{name: "unset defaults to shift", want: true},
+		{name: "explicit false disables shift", shift: &falseValue, want: false},
+		{name: "explicit true enables shift", shift: &trueValue, want: true},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.ctx))
+			require.Equal(t, tt.want, resolveContextShift(tt.shift))
 		})
 	}
 }