Retain prompt-end checkpoints per session using existing entry limit

Author: mergennachin

examples/llm_server/cpp/module_worker_main.cpp

@@ -30,7 +30,7 @@ DEFINE_int32(max_sessions, 4, "Maximum logical serving sessions");
 DEFINE_int32(
     prefix_cache_entries,
     0,
-    "Retained prefix snapshots; zero disables prefix caching");
+    "Retained snapshots: shared by default (0), per session with prompt-end checkpoints (8)");
 DEFINE_int32(
     prefix_checkpoint_interval,
     0,
@@ -42,7 +42,7 @@ DEFINE_int32(
 DEFINE_bool(
     prefix_checkpoint_prompt_end,
     false,
-    "Best-effort retention of each request's prompt end in the bounded prefix cache");
+    "Retain prompt-end checkpoints separately for each named session");
 DEFINE_int32(max_session_tokens, 2048, "Maximum retained tokens per session");
 DEFINE_int32(max_decode_sequences, 2, "Decode slots in one packed forward");
 DEFINE_int32(
@@ -52,11 +52,16 @@ DEFINE_int32(
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
+  if (FLAGS_prefix_checkpoint_prompt_end &&
+      gflags::GetCommandLineFlagInfoOrDie("prefix_cache_entries").is_default) {
+    FLAGS_prefix_cache_entries = 8;
+  }
   if (FLAGS_prefix_checkpoint_prompt_end && FLAGS_prefix_checkpoint_interval) {
     std::cerr << "--prefix_checkpoint_prompt_end cannot be combined with "
                  "--prefix_checkpoint_interval\n";
     return 1;
   }
+
   if (FLAGS_pte.empty() || FLAGS_tokenizer.empty() || FLAGS_max_sessions <= 0 ||
       FLAGS_max_session_tokens <= 0 || FLAGS_max_decode_sequences <= 0 ||
       FLAGS_max_inflight_requests <= 0 || FLAGS_prefix_cache_entries < 0 ||
@@ -79,10 +84,9 @@ int main(int argc, char** argv) {
   config.max_context_length = FLAGS_max_session_tokens;
   config.max_requests = FLAGS_max_inflight_requests;
   config.max_pending_operations = FLAGS_max_inflight_requests;
-  const auto physical_sessions = static_cast<std::int64_t>(FLAGS_max_sessions) +
-      FLAGS_prefix_cache_entries +
-      static_cast<int64_t>(config.resolved_prefix_captures());
-  if (physical_sessions > std::numeric_limits<int>::max()) {
+  const auto physical_sessions = config.required_executor_rows();
+  if (!physical_sessions ||
+      *physical_sessions > std::numeric_limits<int>::max()) {
     std::cerr
         << "session and prefix-cache capacity exceeds executor row limit\n";
     return 1;
@@ -113,7 +117,7 @@ int main(int argc, char** argv) {
     const auto eos = llm::get_eos_ids(tokenizer.get(), module.get());
     auto executor = llm::batching::ModuleExecutor::create(
         std::move(module),
-        static_cast<int>(physical_sessions),
+        static_cast<int>(*physical_sessions),
         FLAGS_max_session_tokens,
         static_cast<int>(*dtype),
         -1,

examples/llm_server/python/README.md

@@ -66,17 +66,23 @@ limit; choose a value supported by the export. Capacity defaults are
 `--prefix-cache-entries 0`. With prefix caching enabled, `--prefix-capture-slots 0`
 (the default) automatically allows one pending prompt capture per admitted request,
 bounded by `min(max_sessions, max_inflight_requests)`. A positive value overrides
-this allowance. Add `--prefix-checkpoint-prompt-end` to retain later request boundaries
-as well as initial prompts; it requires `--prefix-cache-entries > 0`. This mode
-keeps the same concurrent capture allowance and skips already-retained prompt ends.
-The snapshot becomes reusable and its capture slot is released when the clone is
-collected, without waiting for the source request to finish decoding.
+this allowance. Alternatively, `--prefix-checkpoint-prompt-end` retains initial,
+continuation, and replacement prompt boundaries separately for each named session.
+`--prefix-cache-entries` defaults to eight snapshots per session; duplicates
+are skipped and LRU eviction stays within the owning session. Close/reset clears
+that session's snapshots. Anonymous requests neither use nor retain them.
+This mode reserves `max_sessions * prefix_cache_entries` snapshot rows,
+plus working sessions and capture slots. Eight sessions with eight checkpoints
+and automatic capture sizing need 80 rows in total. Capture slots are released
+after clone collection, without waiting for decoding to finish.
 Capture is best effort: occupied transient rows can prevent a snapshot.
 With periodic checkpoints enabled, automatic sizing uses at most two rows:
 checkpoint collection is serialized, but one replacement can overlap. Explicit
 values above two add unused row capacity in that mode.
 `--prefix-checkpoint-prompt-end` cannot be combined with a nonzero
 `--prefix-checkpoint-interval`; the launcher and native worker reject that combination.
+Here `--prefix-cache-entries` sets the per-session limit and defaults to eight;
+without prompt-end mode it keeps its shared-cache meaning and defaults to zero.
 Logical sessions, packed decode width, and in-flight request capacity are independent
 limits. Python concurrency follows advertised request capacity, not the single
 native delivery thread. Python reads native async streams directly through bounded

examples/llm_server/python/serve.py

@@ -41,7 +41,11 @@ def _parse_preflight_timeout(value):
 
 
 def _validate_limits(
-    parser, args, *, prefix_capture_slots=0, prefix_checkpoint_prompt_end=False
+    parser,
+    args,
+    *,
+    prefix_capture_slots=0,
+    prefix_checkpoint_prompt_end=False,
 ) -> int:
     for flag, value in (
         ("--max-context", args.max_context),
@@ -60,15 +64,12 @@ def _validate_limits(
             "--prefix-checkpoint-prompt-end cannot be combined with "
             "--prefix-checkpoint-interval"
         )
-    if (
-        args.prefix_checkpoint_interval or prefix_checkpoint_prompt_end
-    ) and not args.prefix_cache_entries:
-        flag = (
-            "--prefix-checkpoint-prompt-end"
-            if prefix_checkpoint_prompt_end
-            else "--prefix-checkpoint-interval"
+    if prefix_checkpoint_prompt_end and not args.prefix_cache_entries:
+        parser.error(
+            "--prefix-checkpoint-prompt-end requires --prefix-cache-entries > 0"
         )
-        parser.error(f"{flag} requires --prefix-cache-entries > 0")
+    if args.prefix_checkpoint_interval and not args.prefix_cache_entries:
+        parser.error("--prefix-checkpoint-interval requires --prefix-cache-entries > 0")
     if not 0 <= prefix_capture_slots <= _MAX_INT32:
         parser.error(f"--prefix-capture-slots must be between 0 and {_MAX_INT32}")
     # Mirrors ServingRuntimeConfig::resolved_prefix_captures for row validation.
@@ -78,9 +79,12 @@ def _validate_limits(
             prefix_capture_slots = min(prefix_capture_slots, 2)
     if not args.prefix_cache_entries:
         prefix_capture_slots = 0
-    physical_sessions = (
-        args.max_sessions + args.prefix_cache_entries + prefix_capture_slots
+    snapshots = (
+        args.max_sessions * args.prefix_cache_entries
+        if prefix_checkpoint_prompt_end
+        else args.prefix_cache_entries
     )
+    physical_sessions = args.max_sessions + snapshots + prefix_capture_slots
     if physical_sessions > _MAX_INT32:
         parser.error("session and prefix-cache capacity exceeds the row limit")
     if not 1 <= args.port <= 65535:
@@ -118,11 +122,15 @@ def _parse_args(argv=None):
     parser.add_argument("--max-sessions", type=int, default=16)
     parser.add_argument("--max-decode-sequences", type=int, default=8)
     parser.add_argument("--max-inflight-requests", type=int, default=64)
-    parser.add_argument("--prefix-cache-entries", type=int, default=0)
+    parser.add_argument(
+        "--prefix-cache-entries",
+        type=int,
+        help="Retained snapshots: shared by default (0), per named session with prompt-end checkpoints (8)",
+    )
     parser.add_argument(
         "--prefix-checkpoint-prompt-end",
         action="store_true",
-        help="Best-effort retention of each request's prompt end in the bounded prefix cache",
+        help="Retain prompt-end checkpoints separately for each named session",
     )
     parser.add_argument(
         "--prefix-capture-slots",
@@ -148,6 +156,8 @@ def _parse_args(argv=None):
     parser.add_argument("--host", default="127.0.0.1")
     parser.add_argument("--port", type=int, default=8000)
     args = parser.parse_args(argv)
+    if args.prefix_cache_entries is None:
+        args.prefix_cache_entries = 8 if args.prefix_checkpoint_prompt_end else 0
 
     _validate_limits(
         parser,

examples/llm_server/python/tests/test_serve_launcher.py

@@ -157,24 +157,38 @@ def test_checkpoint_options_reach_worker(argv, launch):
     assert command[command.index("--prefix_checkpoint_interval") + 1] == "512"
 
 
-def test_prompt_end_capture_reaches_worker_and_requires_cache(argv, launch):
+@pytest.mark.parametrize("capacity", [8, 2])
+def test_prompt_end_capture_reaches_worker_with_session_capacity(
+    argv, launch, capacity
+):
     assert not serve._parse_args(argv).prefix_checkpoint_prompt_end
-    with pytest.raises(SystemExit):
-        serve._parse_args(argv + ["--prefix-checkpoint-prompt-end"])
-    serve.main(argv + ["--prefix-cache-entries", "3", "--prefix-checkpoint-prompt-end"])
+    assert (
+        serve._parse_args(
+            argv + ["--prefix-checkpoint-prompt-end"]
+        ).prefix_cache_entries
+        == 8
+    )
+    serve.main(
+        argv
+        + [
+            "--prefix-checkpoint-prompt-end",
+            "--prefix-cache-entries",
+            str(capacity),
+        ]
+    )
     assert (
         "--prefix_checkpoint_prompt_end=true"
         in launch.stages["spawn"].call_args.args[0]
     )
+    command = launch.stages["spawn"].call_args.args[0]
+    assert command[command.index("--prefix_cache_entries") + 1] == str(capacity)
 
 
 def test_prompt_end_rejects_periodic_capture(argv, capsys):
     with pytest.raises(SystemExit):
         serve._parse_args(
             argv
             + [
-                "--prefix-cache-entries",
-                "3",
                 "--prefix-checkpoint-prompt-end",
                 "--prefix-checkpoint-interval",
                 "64",
@@ -183,6 +197,42 @@ def test_prompt_end_rejects_periodic_capture(argv, capsys):
     assert "cannot be combined" in capsys.readouterr().err
 
 
+@pytest.mark.parametrize(
+    "options",
+    [
+        ["--prefix-cache-entries", "0"],
+        ["--prefix-cache-entries", "-1"],
+        ["--max-sessions", str(1 << 28)],
+        ["--prefix-cache-entries", str((1 << 31) - 1)],
+    ],
+)
+def test_session_checkpoint_limits_reject_invalid_combinations_and_row_overflow(
+    argv, options
+):
+    with pytest.raises(SystemExit):
+        serve._parse_args(argv + ["--prefix-checkpoint-prompt-end"] + options)
+
+
+def test_session_checkpoint_rows_include_every_sessions_budget(argv):
+    maximum = (1 << 31) - 1
+    # Two working rows + two checkpoint budgets + one transient row.
+    capacity = (maximum - 3) // 2
+    flags = [
+        "--prefix-checkpoint-prompt-end",
+        "--max-sessions",
+        "2",
+        "--prefix-capture-slots",
+        "1",
+        "--prefix-cache-entries",
+    ]
+    assert (
+        serve._parse_args(argv + flags + [str(capacity)]).prefix_cache_entries
+        == capacity
+    )
+    with pytest.raises(SystemExit):
+        serve._parse_args(argv + flags + [str(capacity + 1)])
+
+
 @pytest.mark.parametrize(
     "sessions,requests,interval,override,expected,prompt_end",
     [

examples/llm_server/spec/README.md

@@ -392,15 +392,21 @@ sampled requests. With `--prefix_capture_slots=C`, it adds `N + C` physical rows
 for retained snapshots and transient captures; logical `--max_sessions` remains
 unchanged. The default `--prefix_capture_slots=0` resolves `C` automatically to
 `min(max_sessions, max_inflight_requests)`, so overlapping requests can each capture
-their initial prompt. `--prefix_checkpoint_prompt_end=true` additionally requests
-continuation and replacement snapshots with the same concurrent allowance.
-It requires nonzero cache capacity and skips exact prompt ends already retained.
-Snapshots are published and their transient slots released after clone collection,
-without waiting for the request to finish decoding.
+their initial prompt. Alternatively, `--prefix_checkpoint_prompt_end=true` retains
+initial, continuation, and replacement snapshots separately for each named session.
+`--prefix_cache_entries=8` is the default per-session limit. This mode adds
+`max_sessions * prefix_cache_entries + C` physical rows. The existing
+`--prefix_cache_entries` flag now sets the per-session limit when prompt-end mode
+is enabled; otherwise it keeps the shared-cache meaning and default of zero.
+Exact prompt identities already retained
+within the session are not cloned again. Close/reset clears the session's snapshots,
+and anonymous requests neither reuse nor retain them. LRU eviction stays within
+each session. Transient slots are released after clone collection without waiting
+for the request to finish decoding.
 With periodic checkpoints enabled, automatic sizing uses at most two
 rows because checkpoint collection remains serialized and one replacement can
 overlap. A positive value overrides automatic sizing; values above two add unused
-row capacity in checkpoint mode. Disabled prefix caching reserves no capture rows.
+row capacity in checkpoint mode. Disabling both cache modes reserves no capture rows.
 Capture is best effort. Prompt-end capture cannot be combined with a nonzero
 checkpoint interval or a custom checkpoint policy; invalid combinations are
 rejected. A C++ `prefix_checkpoint_policy` overrides the interval. Periodic/custom

examples/llm_server/spec/batching.md

@@ -35,13 +35,17 @@ speculative drafting/verification remain inside the runner and executor.
 of N and at the initial prompt end. It requires `--prefix-cache-entries > 0`.
 Zero (the default) disables periodic capture. With prefix caching enabled and no
 other capture options, only initial prompts are retained.
-`--prefix-checkpoint-prompt-end` also requests snapshots of continuation and replacement
-prompt ends. It preserves per-request capture concurrency and cannot be combined
-with a nonzero interval or custom checkpoint policy. Snapshots become reusable and
-capture slots are released when their clones are collected, before decode completes.
+`--prefix-checkpoint-prompt-end` instead retains prompt ends within each named
+session, including continuations and replacements. `--prefix-cache-entries`
+defaults to eight per session. Checkpoints cannot be reused or evicted by another
+session, and close/reset clears them. Anonymous requests do not use this cache.
+This mode cannot be combined with a nonzero interval or custom checkpoint policy.
+In this mode `--prefix-cache-entries` sets the per-session limit (default eight);
+otherwise it sets the shared-cache limit (default zero). It reserves `max_sessions * prefix_cache_entries`
+snapshot rows in addition to working sessions and transient captures.
+Capture slots are released after clone collection, before decode completes.
 Capture is best effort: occupied transient rows can skip a boundary. Already-retained
-prompt ends are not cloned again. Retention uses global LRU; early checkpoints are
-not protected.
+prompt ends are not cloned again. Default retention uses LRU within each cache.
 
 C++ applications can customize two independent policies in
 [`ServingRuntimeConfig`](../../../extension/llm/serving/serving_runtime.h):

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -224,23 +224,29 @@ apply. Text/image prefix snapshots are opt-in via `--prefix-cache-entries`
 
 ## Session Behavior
 
-Prefix caching is disabled unless `--prefix-cache-entries` is nonzero. When
-configured without other capture options, it retains initial prompts. Add
-`--prefix-checkpoint-prompt-end` to request snapshots of subsequent boundaries,
-including tool-result continuations and replacements. Repeated or branched requests
-can reuse those snapshots; new tool-result text still needs prefill. Exact prompt
-ends already in the cache are not cloned again.
-Snapshots become reusable and their capture slots are released after clone
-collection, without waiting for decoding to finish.
+Without prompt-end mode, nonzero `--prefix-cache-entries` enables shared prefix
+caching; without other capture options it retains initial prompts. `--prefix-checkpoint-prompt-end` instead
+retains prompt boundaries separately for each named session, including tool-result
+continuations and replacements. `--prefix-cache-entries` defaults to eight
+checkpoints per session. Repeated or branched requests using the same session ID can
+reuse its snapshots; other sessions and anonymous requests cannot. New tool-result
+text still needs prefill. Exact prompt ends already retained in that session are
+not cloned again. Close/reset, including idle-session eviction, clears its snapshots.
 
 Prompt-end-only capture preserves per-request capture concurrency and the default
-transient-row budget. All snapshots share the bounded cache and global LRU, so new
-boundaries can evict earlier prompts. Capture is best effort: concurrent captures
+transient-row budget. LRU eviction stays within each session, so one session cannot
+evict another's checkpoints. Capture slots are released after clone collection,
+without waiting for decoding to finish. Capture is best effort: concurrent captures
 or session replacements can occupy the available rows. Measure capture overhead
 and reuse for your workload before enabling it.
 
 `--prefix-checkpoint-prompt-end` cannot be combined with a nonzero
-`--prefix-checkpoint-interval`; both the launcher and native worker reject it.
+`--prefix-checkpoint-interval`; both the launcher and native worker reject that combination.
+With prompt-end mode, `--prefix-cache-entries` is the per-session budget (default eight);
+without it, the same option is the shared-cache budget (default zero). This mode reserves
+`max_sessions * prefix_cache_entries` snapshot rows plus working sessions
+and capture slots: eight sessions, eight checkpoints each, and eight automatic
+capture slots require 80 physical rows. The feature remains disabled by default.
 Periodic checkpoints use one serialized lane and may skip a boundary while an
 earlier clone is pending. Automatic transient-row sizing is capped at two in
 periodic mode; prompt-end capture retains the ordinary concurrency budget.
@@ -339,7 +345,8 @@ python -m executorch.examples.models.muse_glimmer.serving.serve_batching \
 
 `--data-path` is required. Image artifacts read `pos_embed.bin` beside
 `model.pte`, or `--pos-embed-path`. Every resident session reserves `--max-context` cells of the
-pool: requests, `--prefix-cache-entries` snapshots, and the resolved number of
+pool: requests, shared `--prefix-cache-entries` snapshots or
+`max_sessions * prefix_cache_entries` per-session snapshots, and the resolved number of
 capture slots when snapshots are enabled. `--prefix-capture-slots 0` (the default)
 matches request concurrency, including prompt-end-only mode; periodic checkpoints
 cap the automatic count at two; a