server : improve context checkpoint logic (#19408)

diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index b71d496eeb80357f65c2be69e951199d46623707..8ec84513399155fafb11639b42bc0bf6281d560f 100644 (file)
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -2507,7 +2507,8 @@ private:
                         slot.n_prompt_tokens_processed++;
 
                         // process the last few tokens of the prompt separately in order to allow for a checkpoint to be created.
-                        if (do_checkpoint && slot.task->n_tokens() - slot.prompt.n_tokens() == 64) {
+                        const int n_last = std::min(n_batch, 512);
+                        if (do_checkpoint && slot.task->n_tokens() == slot.prompt.n_tokens() + n_last) {
                             break;
                         }
                     }