Preserve concurrent prompt captures and avoid duplicate snapshots

Author: mergennachin

examples/llm_server/cpp/module_worker_main.cpp

@@ -42,7 +42,7 @@ DEFINE_int32(
 DEFINE_bool(
     prefix_capture_prompt_end,
     false,
-    "Retain each request's prompt end in the bounded prefix cache");
+    "Best-effort retention of each request's prompt end in the bounded prefix cache");
 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(
@@ -62,7 +62,7 @@ int main(int argc, char** argv) {
     std::cerr
         << "--pte and --tokenizer are required; capacity limits must be positive"
            "; prefix-cache limits must be nonnegative; "
-           "checkpoints require prefix_cache_entries > 0\n";
+           "--prefix_checkpoint_interval and --prefix_capture_prompt_end require --prefix_cache_entries > 0\n";
     return 1;
   }
   executorch::extension::llm::serving::ServingRuntimeConfig config;

examples/llm_server/python/README.md

@@ -66,9 +66,16 @@ 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. With checkpoints enabled, automatic sizing uses at most two rows:
+this allowance. Add `--prefix-capture-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.
+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.
+Combining `--prefix-checkpoint-interval` and `--prefix-capture-prompt-end` requests
+both periodic and exact-end snapshots through the shared checkpoint lane. A pending
+checkpoint can prevent capture at a later boundary, including the prompt end.
 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

@@ -55,14 +55,21 @@ def _validate_limits(
         parser.error(f"--prefix-cache-entries must be between 0 and {_MAX_INT32}")
     if not 0 <= args.prefix_checkpoint_interval <= _MAX_INT32:
         parser.error(f"--prefix-checkpoint-interval must be between 0 and {_MAX_INT32}")
-    if (args.prefix_checkpoint_interval or prefix_capture_prompt_end) and not args.prefix_cache_entries:
-        parser.error("--prefix-checkpoint-interval requires --prefix-cache-entries > 0")
+    if (
+        args.prefix_checkpoint_interval or prefix_capture_prompt_end
+    ) and not args.prefix_cache_entries:
+        flag = (
+            "--prefix-capture-prompt-end"
+            if prefix_capture_prompt_end
+            else "--prefix-checkpoint-interval"
+        )
+        parser.error(f"{flag} 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.
     if prefix_capture_slots == 0:
         prefix_capture_slots = min(args.max_sessions, args.max_inflight_requests)
-        if args.prefix_checkpoint_interval or prefix_capture_prompt_end:
+        if args.prefix_checkpoint_interval:
             prefix_capture_slots = min(prefix_capture_slots, 2)
     if not args.prefix_cache_entries:
         prefix_capture_slots = 0
@@ -108,8 +115,9 @@ def _parse_args(argv=None):
     parser.add_argument("--max-inflight-requests", type=int, default=64)
     parser.add_argument("--prefix-cache-entries", type=int, default=0)
     parser.add_argument(
-        "--prefix-capture-prompt-end", action="store_true",
-        help="Retain each request's prompt end in the bounded prefix cache",
+        "--prefix-capture-prompt-end",
+        action="store_true",
+        help="Best-effort retention of each request's prompt end in the bounded prefix cache",
     )
     parser.add_argument(
         "--prefix-capture-slots",
@@ -121,7 +129,7 @@ def _parse_args(argv=None):
         "--prefix-checkpoint-interval",
         type=int,
         default=0,
-        help="Checkpoint interval in tokens; 0: initial prompt only. Uses global LRU.",
+        help="Periodic checkpoint interval in tokens; 0 disables periodic capture.",
     )
     parser.add_argument(
         "--api-dialect", choices=("openai", "openrouter"), default="openai"
@@ -136,7 +144,12 @@ def _parse_args(argv=None):
     parser.add_argument("--port", type=int, default=8000)
     args = parser.parse_args(argv)
 
-    _validate_limits(parser, args, prefix_capture_slots=args.prefix_capture_slots)
+    _validate_limits(
+        parser,
+        args,
+        prefix_capture_slots=args.prefix_capture_slots,
+        prefix_capture_prompt_end=args.prefix_capture_prompt_end,
+    )
 
     args.worker_bin = os.path.expanduser(args.worker_bin)
     if shutil.which(args.worker_bin) is None:
@@ -212,7 +225,11 @@ async def serving_factory():
                     if args.prefix_checkpoint_interval
                     else []
                 )
-                + (["--prefix_capture_prompt_end=true"] if args.prefix_capture_prompt_end else []),
+                + (
+                    ["--prefix_capture_prompt_end=true"]
+                    if args.prefix_capture_prompt_end
+                    else []
+                ),
                 stderr=stderr,
             )
             runtime = None

examples/llm_server/python/tests/test_serve_launcher.py

@@ -167,30 +167,7 @@ def test_prompt_end_capture_reaches_worker_and_requires_cache(argv, launch):
     )
 
 
-@pytest.mark.parametrize(
-    "sessions,override,expected", [(8, 0, 2), (1, 0, 1), (8, 3, 3)]
-)
-def test_prompt_end_capture_bounds_automatic_rows(argv, sessions, override, expected):
-    args = serve._parse_args(
-        argv
-        + [
-            "--prefix-cache-entries",
-            "3",
-            "--prefix-capture-prompt-end",
-            "--max-sessions",
-            str(sessions),
-            "--prefix-capture-slots",
-            str(override),
-        ]
-    )
-    assert (
-        serve._validate_limits(
-            Mock(), args, prefix_capture_slots=override, prefix_capture_prompt_end=True
-        )
-        == expected
-    )
-
-
+@pytest.mark.parametrize("prompt_end", [False, True])
 @pytest.mark.parametrize(
     "sessions,requests,interval,override,expected",
     [
@@ -203,7 +180,7 @@ def test_prompt_end_capture_bounds_automatic_rows(argv, sessions, override, expe
     ],
 )
 def test_capture_capacity_validation_and_worker_override(
-    argv, launch, sessions, requests, interval, override, expected
+    argv, launch, sessions, requests, interval, override, expected, prompt_end
 ):
     serve.main(
         argv
@@ -219,9 +196,11 @@ def test_capture_capacity_validation_and_worker_override(
             "--prefix-capture-slots",
             str(override),
         ]
+        + (["--prefix-capture-prompt-end"] if prompt_end else [])
     )
     command = launch.stages["spawn"].call_args.args[0]
     assert command[command.index("--prefix_capture_slots") + 1] == str(override)
+    assert ("--prefix_capture_prompt_end=true" in command) == prompt_end
     args = SimpleNamespace(
         max_context=1024,
         max_sessions=sessions,
@@ -232,7 +211,13 @@ def test_capture_capacity_validation_and_worker_override(
         port=8000,
     )
     assert (
-        serve._validate_limits(Mock(), args, prefix_capture_slots=override) == expected
+        serve._validate_limits(
+            Mock(),
+            args,
+            prefix_capture_slots=override,
+            prefix_capture_prompt_end=prompt_end,
+        )
+        == expected
     )
 
 

examples/llm_server/spec/README.md

@@ -387,15 +387,21 @@ batched cache via `cache::kind::kBatched`, derives scheduler width from
 `preferred_batch_tokens()`, and constructs
 one `ServingRuntime`. It does not preload `forward`, run token-step loops, or use
 Glimmer's legacy session implementation. Prefix caching is disabled by default.
-`--prefix_cache_entries=N` enables creation-only reuse for both greedy and
+`--prefix_cache_entries=N` enables retained-prefix reuse for both greedy and
 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. With checkpoints enabled, automatic sizing uses at most two
+their initial prompt. `--prefix_capture_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.
+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.
+Capture is best effort. Periodic checkpoints and prompt ends requested together
+share a serialized checkpoint lane; a pending clone can cause another boundary to
+be skipped. A C++ `prefix_checkpoint_policy` overrides both capture flags.
 
 With an existing MLX-enabled ExecuTorch installation and matching gflags package,
 configure the MLX LLM examples and build `llm_worker`. Launch with `--pte`,

examples/llm_server/spec/batching.md

@@ -33,13 +33,20 @@ speculative drafting/verification remain inside the runner and executor.
 
 `--prefix-checkpoint-interval N` captures prompt snapshots at absolute multiples
 of N and at the initial prompt end. It requires `--prefix-cache-entries > 0`.
-Zero (the default) captures only the initial prompt. Capture is best effort.
-Retention uses global LRU; early checkpoints are not protected.
+Zero (the default) disables periodic capture. With prefix caching enabled and no
+other capture options, only initial prompts are retained.
+`--prefix-capture-prompt-end` also requests snapshots of continuation and replacement
+prompt ends. Alone, it preserves per-request capture concurrency; combined with an
+interval, it shares the serialized checkpoint lane. Capture is best effort: an
+occupied row or checkpoint lane can skip a boundary. Already-retained prompt ends
+are not cloned again. Retention uses global LRU; early checkpoints are not protected.
 
 C++ applications can customize two independent policies in
 [`ServingRuntimeConfig`](../../../extension/llm/serving/serving_runtime.h):
 
-- `prefix_checkpoint_policy` chooses capture positions and optional application metadata.
+- `prefix_checkpoint_policy` chooses capture positions and optional application metadata,
+  overriding both `prefix_checkpoint_interval` and `prefix_capture_prompt_end`.
+  This C++ callback is not a launcher option.
 - `prefix_cache_retention` chooses which snapshots to retain, including early anchors.
 
 The [retention interface and default LRU policy](../../../extension/llm/batching/prefix_cache_retention.h)

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -224,19 +224,25 @@ apply. Text/image prefix snapshots are opt-in via `--prefix-cache-entries`
 
 ## Session Behavior
 
-Prefix caching retains the initial prompt by default. Add
-`--prefix-capture-prompt-end` with `--prefix-cache-entries` to also retain the
-prompt boundary of subsequent requests, including tool-result continuations
-and replacements. Repeated or branched requests can then reuse those snapshots.
-New tool-result text still needs prefill. This is best-effort capture within the
-same bounded cache and transient-row limits, with global LRU retention; it can
-evict earlier prompts. Benchmark capture cost and memory for your workload
-before enabling it. It also works with `--prefix-checkpoint-interval` and adds
-the exact prompt end when it falls between periodic checkpoints. A pending
-checkpoint or another request can occupy the capture lane and skip a boundary;
-use prompt-end capture without periodic checkpoints when turn reuse is the
-priority. Automatic transient-row sizing is capped at two in either mode. The C++
-`prefix_checkpoint_policy`, when supplied, overrides both options.
+Prefix caching is disabled unless `--prefix-cache-entries` is nonzero. When
+configured without other capture options, it retains initial prompts. Add
+`--prefix-capture-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.
+
+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
+or session replacements can occupy the available rows. Measure capture overhead
+and reuse for your workload before enabling it.
+
+With `--prefix-checkpoint-interval`, periodic points and the exact prompt end use
+one serialized checkpoint lane. A pending checkpoint can cause a later boundary,
+including the prompt end, to be skipped. Automatic transient-row sizing is capped
+at two in periodic mode; prompt-end-only mode retains the ordinary concurrency
+budget. Use prompt-end capture without periodic checkpoints when turn reuse is the
+priority.
 
 - Named sessions persist by default. `--evict-idle-sessions` allows the server to
   close the least recently used idle session when logical session slots are exhausted, clearing
@@ -334,7 +340,8 @@ python -m executorch.examples.models.muse_glimmer.serving.serve_batching \
 `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
 capture slots when snapshots are enabled. `--prefix-capture-slots 0` (the default)
-matches request concurrency, capped at two when checkpoints are enabled; a
+matches request concurrency, including prompt-end-only mode; periodic checkpoints
+cap the automatic count at two; a
 positive value sets an explicit limit. Their total reservation must leave one
 cell for padded writes within the exported `--max-cells`. More capture slots
 reduce the context available per session in the native worker's automatic

examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp

@@ -47,7 +47,7 @@ DEFINE_int32(
 DEFINE_bool(
     prefix_capture_prompt_end,
     false,
-    "Retain each request's prompt end in the bounded prefix cache");
+    "Best-effort retention of each request's prompt end in the bounded prefix cache");
 DEFINE_uint64(max_image_bytes, 20 * 1024 * 1024, "Maximum encoded image bytes");
 DEFINE_uint64(
     max_input_frame_bytes,

examples/models/muse-glimmer/serving/serve_batching.py

@@ -90,7 +90,7 @@ def _parse_args(argv=None):
     parser.add_argument(
         "--prefix-capture-prompt-end",
         action="store_true",
-        help="Retain each request's prompt end in the bounded prefix cache",
+        help="Best-effort retention of each request's prompt end in the bounded prefix cache",
     )
     parser.add_argument(
         "--prefix-checkpoint-interval",

extension/llm/batching/prefix_cache.h

@@ -48,14 +48,16 @@ inline std::vector<Position> periodic_prefill_breakpoints(
     Position interval,
     bool include_end = false) {
   std::vector<Position> positions;
-  if (start < 0 || end <= start || interval <= 0) {
+  if (start < 0 || end <= start || interval < 0) {
     return positions;
   }
-  for (std::int64_t next =
-           (static_cast<std::int64_t>(start) / interval + 1) * interval;
-       next <= end;
-       next += interval) {
-    positions.push_back(static_cast<Position>(next));
+  if (interval > 0) {
+    for (std::int64_t next =
+             (static_cast<std::int64_t>(start) / interval + 1) * interval;
+         next <= end;
+         next += interval) {
+      positions.push_back(static_cast<Position>(next));
+    }
   }
   if (include_end && (positions.empty() || positions.back() != end)) {
     positions.push_back(end);
@@ -348,6 +350,15 @@ class ET_EXPERIMENTAL PrefixCache {
     return lookup_identity(copy_tokens(request));
   }
 
+  // Exact retained identity, without cloning or refreshing recency.
+  bool contains_identity(const PrefixIdentity& request) const {
+    return std::any_of(
+        entries_.begin(), entries_.end(), [&](const Entry& entry) {
+          return entry.tokens.size() == request.size() &&
+              common_prefix(entry.tokens, request) == request.size();
+        });
+  }
+
   std::optional<PrefixMatch> lookup_identity(const PrefixIdentity& request) {
     auto match = find_identity(request);
     if (match) {