diff --git a/docs/PROMPT_CACHE_TAIL.md b/docs/PROMPT_CACHE_TAIL.md new file mode 100644 index 0000000000..3ef8304a4a --- /dev/null +++ b/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. diff --git a/src/program/generate.cpp b/src/program/generate.cpp index a7c77d0eff..5b5e4cf598 100644 --- a/src/program/generate.cpp +++ b/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(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 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 lk(part_mu); part_at.clear();