Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 23 additions & 4 deletions tools/server/server-context.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -277,10 +277,13 @@ struct server_slot {

n_remaining = -1;

if (task->params.n_predict != -1) {
n_remaining = task->params.n_predict - n_decoded;
} else if (global_params.n_predict != -1) {
n_remaining = global_params.n_predict - n_decoded;
const int32_t n_predict = task->params.n_predict != -1 ? task->params.n_predict : global_params.n_predict;

if (n_predict == -2) {
// generate until context is full
n_remaining = n_ctx - prompt.n_tokens();
} else if (n_predict >= 0) {
n_remaining = n_predict - n_decoded;
}

return n_remaining > 0; // no budget
Expand Down Expand Up @@ -2175,6 +2178,22 @@ struct server_context_impl {
continue;
}

// n_predict == -2: context exhaustion is the limit, do not shift
{
const int32_t n_predict = slot.task->params.n_predict != -1 ? slot.task->params.n_predict : params_base.n_predict;
if (n_predict == -2) {
SLT_DBG(slot, "stopped by n_predict = -2 (context full), prompt.n_tokens() = %d, n_decoded = %d, n_ctx = %d\n",
slot.prompt.n_tokens(), slot.n_decoded, slot.n_ctx);
slot.stop = STOP_TYPE_LIMIT;
slot.truncated = true;
slot.print_timings();
send_final_response(slot);
metrics.on_prediction(slot);
slot.release();
continue;
}
}

if (mctx) {
// we should never reach this because params_base.ctx_shift is automatically disabled if mmproj is loaded
// we don't support ctx_shift because an image chunk may contains multiple tokens
Expand Down
41 changes: 41 additions & 0 deletions tools/server/tests/unit/test_ctx_shift.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,44 @@ def test_ctx_shift_disabled_stream():
else:
assert choice["finish_reason"] is None
content += choice["text"]


@pytest.mark.parametrize("endpoint,n_predict_key", [
("/completion", "n_predict"),
("/v1/completions", "max_tokens"),
])
def test_n_predict_minus_2(endpoint: str, n_predict_key: str):
"""n_predict == -2: generate until context is full, then stop (no ctx shift)."""
global server
server.n_predict = -1 # global: unlimited, request-level -2 overrides
server.enable_ctx_shift = True # enabled, but -2 should stop instead of shifting
server.start()
res = server.make_request("POST", endpoint, data={
n_predict_key: -2,
"prompt": "Hi how are you",
})
assert res.status_code == 200
# "Hi how are you" is 8 tokens, slot ctx = 512/2 = 256, expect ~248 generated
if "timings" in res.body:
n_predicted = res.body["timings"]["predicted_n"]
assert res.body["truncated"] is True
assert res.body["stopped_limit"] is True
else:
n_predicted = res.body["usage"]["completion_tokens"]
assert n_predicted == 248, f"n_predict=-2 should fill context (expected 248), got {n_predicted}"


def test_n_predict_minus_2_global():
"""n_predict == -2 set globally (via server CLI) should also work."""
global server
server.n_predict = -2
server.enable_ctx_shift = True
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "Hi how are you",
})
assert res.status_code == 200
n_predicted = res.body["timings"]["predicted_n"]
assert res.body["truncated"] is True
assert res.body["stopped_limit"] is True
assert n_predicted == 248, f"global n_predict=-2 should fill context (expected 248), got {n_predicted}"