serve: optionally checkpoint an existing chunk near the prompt tail

docs/PROMPT_CACHE_TAIL.md

@@ -0,0 +1,29 @@
+# Optional checkpoint near the prompt tail
+
+`--prompt-cache-tail` saves at most one extra mid-prompt checkpoint near the end
+of a request, at a boundary the batched prefill already reaches. It is off by
+default and applies to single-GPU serve with positive `--prompt-cache` and
+`--prompt-cache-every` values.
+
+A branch can share most of a prompt but diverge before its final cached state.
+If the periodic checkpoint before that branch is far away, the engine has to
+read the gap again. The extra checkpoint can reduce that gap. It uses the
+existing bounded checkpoint cache; it does not add another unbounded history.
+
+The option does not split prefill chunks or alter the periodic checkpoint
+schedule. It selects the first completed batched chunk within one configured
+chunk of the prompt's end, excluding the short final suffix controlled by
+`--short-read`. If a periodic checkpoint already falls there, the same state is
+saved once. Reread diagnostics and layer splits keep their existing behavior.
+
+For example, with a 6144-token chunk and a 16384-token checkpoint interval,
+the first periodic checkpoint is at 18432 and the next would be at 36864.
+A branch sharing 32769 tokens cannot use the latter. A near-tail checkpoint
+at 30720 can reuse more of that prefix,
+without changing how the original prompt was computed. The benefit depends
+on the request and the branch position; it is not a universal speed increase.
+
+This is a separate conservative alternative to the Strata-2080Ti Daily's
+checkpoint-grid patch. The grid changes chunk boundaries to reach exact
+checkpoint intervals and failed its local paired numerical gate. This option
+does not claim that the original grid patch has been fixed or validated.

src/program/generate.cpp

@@ -407,6 +407,7 @@ struct Options {
     int64_t conversation_cache_min_free_mib = 2560;
     /// --serve: also keep a checkpoint every N freshly read prompt tokens (0 = only at the last turn boundary)
     int64_t prompt_cache_every = 16384;
+    bool prompt_cache_tail = false;   // optional extra checkpoint at an existing near-tail chunk boundary
     /// --serve: a prompt read from token 0 is also checkpointed at its first turn boundary - the end of the system
     /// prompt, which every chat of the same client shares - when that is at least N tokens (0 = never)
     int64_t prompt_cache_root = 2048;
@@ -516,6 +517,8 @@ void usage() {
                  "  --conversation-cache-slots N  --serve: at most N parked conversations (default 4)\n"
                  "  --conversation-cache-min-free-mib N  --serve: physical RAM floor when parking (default 2560)\n"
                  "  --prompt-cache-every N  --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n"
+                 "  --prompt-cache-tail  --serve, single GPU: one extra checkpoint near the prompt's end, at an existing\n"
+                 "                       chunk boundary (default off; requires --prompt-cache and --prompt-cache-every > 0)\n"
                  "  --turn-token ID      --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n"
                  "  --short-read N       --serve: read at most N fresh text tokens through the decode windows instead\n"
                  "                       of the batched prompt path (default 64, 0 = off)\n"
@@ -1245,6 +1248,7 @@ int main(int argc, char** argv) {
             else if (a == "--conversation-cache-min-free-mib") o.conversation_cache_min_free_mib = number;
             else o.conversation_cache_slots = (int) number;
         }
+        else if (a == "--prompt-cache-tail") o.prompt_cache_tail = true;
         else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every")));
         else if (a == "--prompt-cache-root") o.prompt_cache_root = std::max(0LL, std::atoll(next("--prompt-cache-root")));
         else if (a == "--turn-token") o.turn_token = std::atoll(next("--turn-token"));
@@ -4628,6 +4632,7 @@ int main(int argc, char** argv) {
             return true;
         };
         int64_t pp_total = 0, pp_from = 0, pp_next_check = 0;
+        bool pp_tail_saved = false;
         // #471: the position the prompt pass has read up to (a chunk's or a window's end): what a request cancelled
         // mid-read reports as read, instead of the whole prompt
         int64_t pp_reached = 0;
@@ -4701,7 +4706,13 @@ int main(int argc, char** argv) {
             strata::core::progress_at("reading the prompt (batched), done up to token", done);
             strata::core::progress_beat();
             std::fflush(stdout);
-            if (o.prompt_cache_every > 0 && done >= pp_next_check) {
+            const bool periodic_checkpoint = o.prompt_cache_every > 0 && done >= pp_next_check;
+            // Keep the prefill's existing chunk geometry. One extra near-tail checkpoint
+            // can serve a branch whose shared prefix ends before the final cached state.
+            const bool tail_checkpoint = o.prompt_cache_tail && o.prompt_cache > 0 && o.prompt_cache_every > 0 &&
+                !multi_gpu && !pp_tail_saved && pp_total - done > std::max<int64_t>(1, o.short_read) &&
+                pp_total - done <= o.prefill_chunk;
+            if (periodic_checkpoint || tail_checkpoint) {
                 bool saved = false;
                 if (multi_gpu) {   // the stages' parts, saved when each of them read this chunk
                     std::vector<ConvCheckpoint> parts;
@@ -4718,7 +4729,8 @@ int main(int argc, char** argv) {
                     saved = checkpoint_at(done);
                 }
                 if (!saved) { e = "saving a conversation checkpoint failed"; return false; }
-                pp_next_check = done + o.prompt_cache_every;
+                if (periodic_checkpoint) pp_next_check = done + o.prompt_cache_every;
+                if (tail_checkpoint) pp_tail_saved = true;
             }
             return true;
         };
@@ -5378,6 +5390,7 @@ int main(int argc, char** argv) {
             pp_reached = read_from;
             pp_t0 = r0;
             pp_next_check = reread_to > 0 ? INT64_MAX : resume + o.prompt_cache_every;
+            pp_tail_saved = reread_to > 0;
             {
                 std::lock_guard<std::mutex> lk(part_mu);
                 part_at.clear();