Add bounded opt-in capture at serving prompt boundaries

Author: mergennachin

examples/llm_server/cpp/module_worker_main.cpp

@@ -34,11 +34,15 @@ DEFINE_int32(
 DEFINE_int32(
     prefix_checkpoint_interval,
     0,
-    "Checkpoint interval in tokens; 0: initial prompt only. Uses global LRU.");
+    "Periodic checkpoint interval in tokens; 0 disables periodic capture.");
 DEFINE_int32(
     prefix_capture_slots,
     0,
     "Transient capture rows; 0: size automatically from session/request limits");
+DEFINE_bool(
+    prefix_capture_prompt_end,
+    false,
+    "Retain 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(
@@ -52,7 +56,8 @@ int main(int argc, char** argv) {
       FLAGS_max_session_tokens <= 0 || FLAGS_max_decode_sequences <= 0 ||
       FLAGS_max_inflight_requests <= 0 || FLAGS_prefix_cache_entries < 0 ||
       FLAGS_prefix_capture_slots < 0 || FLAGS_prefix_checkpoint_interval < 0 ||
-      (FLAGS_prefix_checkpoint_interval > 0 &&
+      ((FLAGS_prefix_checkpoint_interval > 0 ||
+        FLAGS_prefix_capture_prompt_end) &&
        FLAGS_prefix_cache_entries == 0)) {
     std::cerr
         << "--pte and --tokenizer are required; capacity limits must be positive"
@@ -65,6 +70,7 @@ int main(int argc, char** argv) {
   config.prefix_cache_capacity = FLAGS_prefix_cache_entries;
   config.max_prefix_captures = FLAGS_prefix_capture_slots;
   config.prefix_checkpoint_interval = FLAGS_prefix_checkpoint_interval;
+  config.prefix_capture_prompt_end = FLAGS_prefix_capture_prompt_end;
   config.max_context_length = FLAGS_max_session_tokens;
   config.max_requests = FLAGS_max_inflight_requests;
   config.max_pending_operations = FLAGS_max_inflight_requests;

examples/llm_server/python/serve.py

@@ -24,7 +24,9 @@
 _MAX_INT32 = (1 << 31) - 1
 
 
-def _validate_limits(parser, args, *, prefix_capture_slots=0) -> int:
+def _validate_limits(
+    parser, args, *, prefix_capture_slots=0, prefix_capture_prompt_end=False
+) -> int:
     for flag, value in (
         ("--max-context", args.max_context),
         ("--max-sessions", args.max_sessions),
@@ -37,14 +39,18 @@ def _validate_limits(parser, args, *, prefix_capture_slots=0) -> int:
         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 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:
+        parser.error(
+            "prefix checkpoints and prompt-end capture require --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:
+        if args.prefix_checkpoint_interval or prefix_capture_prompt_end:
             prefix_capture_slots = min(prefix_capture_slots, 2)
     if not args.prefix_cache_entries:
         prefix_capture_slots = 0
@@ -89,6 +95,11 @@ def _parse_args(argv=None):
     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-capture-prompt-end",
+        action="store_true",
+        help="Retain each request's prompt end in the bounded prefix cache",
+    )
     parser.add_argument(
         "--prefix-capture-slots",
         type=int,
@@ -99,13 +110,18 @@ 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("--host", default="127.0.0.1")
     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:
@@ -177,6 +193,11 @@ async def serving_factory():
                 if args.prefix_checkpoint_interval
                 else []
             )
+            + (
+                ["--prefix_capture_prompt_end=true"]
+                if args.prefix_capture_prompt_end
+                else []
+            )
         )
         runtime = None
         try:

examples/llm_server/python/tests/test_serve_launcher.py

@@ -157,6 +157,40 @@ 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):
+    assert not serve._parse_args(argv).prefix_capture_prompt_end
+    with pytest.raises(SystemExit):
+        serve._parse_args(argv + ["--prefix-capture-prompt-end"])
+    serve.main(argv + ["--prefix-cache-entries", "3", "--prefix-capture-prompt-end"])
+    assert (
+        "--prefix_capture_prompt_end=true" in launch.stages["spawn"].call_args.args[0]
+    )
+
+
+@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(
     "sessions,requests,interval,override,expected",
     [

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -143,6 +143,20 @@ 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.
+
 - Text and image prompts use the shared prefix-reuse path. Image rows count as
   decoder positions; matching uses token IDs and an image-content digest with
   content-relative row offsets, not patch token IDs alone.

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

@@ -39,11 +39,15 @@ DEFINE_int32(
 DEFINE_int32(
     prefix_checkpoint_interval,
     0,
-    "Checkpoint interval in tokens; 0: initial prompt only. Uses global LRU.");
+    "Periodic checkpoint interval in tokens; 0 disables periodic capture.");
 DEFINE_int32(
     prefix_capture_slots,
     0,
     "Transient capture rows; 0: size automatically from session/request limits");
+DEFINE_bool(
+    prefix_capture_prompt_end,
+    false,
+    "Retain 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,
@@ -67,7 +71,8 @@ create_muse_glimmer_batching_runtime() {
       FLAGS_max_inflight_requests <= 0 || FLAGS_max_decode_sequences <= 0 ||
       FLAGS_prefix_cache_entries < 0 || FLAGS_max_image_bytes == 0 ||
       FLAGS_prefix_checkpoint_interval < 0 || FLAGS_prefix_capture_slots < 0 ||
-      (FLAGS_prefix_checkpoint_interval > 0 &&
+      ((FLAGS_prefix_checkpoint_interval > 0 ||
+        FLAGS_prefix_capture_prompt_end) &&
        FLAGS_prefix_cache_entries == 0) ||
       FLAGS_max_image_bytes > 20 * 1024 * 1024 ||
       FLAGS_max_input_frame_bytes == 0 ||
@@ -84,6 +89,7 @@ create_muse_glimmer_batching_runtime() {
   serving.prefix_cache_capacity = FLAGS_prefix_cache_entries;
   serving.max_prefix_captures = FLAGS_prefix_capture_slots;
   serving.prefix_checkpoint_interval = FLAGS_prefix_checkpoint_interval;
+  serving.prefix_capture_prompt_end = FLAGS_prefix_capture_prompt_end;
   const int64_t physical = static_cast<int64_t>(FLAGS_max_sessions) +
       FLAGS_prefix_cache_entries +
       static_cast<int64_t>(serving.resolved_prefix_captures());

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

@@ -74,11 +74,16 @@ def _parse_args(argv=None):
         help="Transient capture rows; 0 (default): size from session/request limits",
     )
     parser.add_argument("--max-vision-patches", type=int, default=4096)
+    parser.add_argument(
+        "--prefix-capture-prompt-end",
+        action="store_true",
+        help="Retain each request's prompt end in the bounded prefix cache",
+    )
     parser.add_argument(
         "--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("--max-image-bytes", type=int, default=_MAX_IMAGE_BYTES)
     parser.add_argument(
@@ -93,7 +98,12 @@ 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)
-    _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,
+    )
     if not 0 <= args.dflash_n_draft < _INT32_MAX:
         parser.error("--dflash-n-draft must be nonnegative and below INT32_MAX")
     if args.dflash_n_draft and not args.dflash:
@@ -169,6 +179,10 @@ def _worker_command(args):
             args.prefix_capture_slots if args.prefix_cache_entries else None,
         ),
         ("prefix_checkpoint_interval", args.prefix_checkpoint_interval or None),
+        (
+            "prefix_capture_prompt_end",
+            "true" if args.prefix_capture_prompt_end else None,
+        ),
         ("max_vision_patches", args.max_vision_patches),
         ("max_image_bytes", args.max_image_bytes),
         ("max_input_frame_bytes", args.max_request_bytes),

examples/models/muse-glimmer/tests/test_serve_batching.py

@@ -51,6 +51,21 @@ def test_prefix_capture_override_reaches_worker(arguments, captures):
     assert f"--prefix_capture_slots={captures}" in serve_batching._worker_command(args)
 
 
+def test_prompt_end_capture_is_opt_in_and_reaches_worker(arguments):
+    default = serve_batching._parse_args(arguments + ["--prefix-cache-entries", "4"])
+    assert not default.prefix_capture_prompt_end
+    assert not any(
+        "prefix_capture_prompt_end" in flag
+        for flag in serve_batching._worker_command(default)
+    )
+    enabled = serve_batching._parse_args(
+        arguments + ["--prefix-cache-entries", "4", "--prefix-capture-prompt-end"]
+    )
+    assert "--prefix_capture_prompt_end=true" in serve_batching._worker_command(enabled)
+    with pytest.raises(SystemExit):
+        serve_batching._parse_args(arguments + ["--prefix-capture-prompt-end"])
+
+
 def test_worker_forwards_mg_options(arguments):
     args = serve_batching._parse_args(
         arguments

extension/llm/serving/serving_runtime.cpp

@@ -262,6 +262,7 @@ struct ServingRuntime::Impl {
                 static_cast<std::size_t>(
                     std::numeric_limits<batching::Position>::max()) &&
             ((!config.prefix_checkpoint_interval &&
+              !config.prefix_capture_prompt_end &&
               !config.prefix_checkpoint_policy) ||
              config.prefix_cache_capacity)),
         tokenizer_(tokenizer),
@@ -1159,6 +1160,7 @@ struct ServingRuntime::Impl {
       std::function<void(batching::Position)> on_prefill;
       if (request->text && request->text->prompt &&
           (config_.prefix_checkpoint_interval != 0 ||
+           config_.prefix_capture_prompt_end ||
            config_.prefix_checkpoint_policy)) {
 #if ET_HAS_EXCEPTIONS
         try {
@@ -1192,6 +1194,10 @@ struct ServingRuntime::Impl {
                 static_cast<batching::Position>(
                     config_.prefix_checkpoint_interval),
                 opening);
+            if (config_.prefix_capture_prompt_end && end > start &&
+                (breakpoints.empty() || breakpoints.back() != end)) {
+              breakpoints.push_back(end);
+            }
           }
           on_prefill = request->text->prefix_capture->observer(
               [this, request] {

extension/llm/serving/serving_runtime.h

@@ -60,9 +60,12 @@ struct ET_EXPERIMENTAL ServingRuntimeConfig {
   // Capture at absolute interval multiples and the initial prompt end.
   // Requires cache capacity; zero captures only the initial prompt.
   std::size_t prefix_checkpoint_interval = 0;
-  // Override the interval with increasing checkpoints in (start,
-  // prompt.size()]. Invalid or throwing policies skip capture. Requires cache
-  // capacity. Runs on control; must not block or reenter the runtime.
+  // Best-effort capture at each request's prompt end, including continuations
+  // and replacements. Shares checkpoint admission; requires cache capacity.
+  bool prefix_capture_prompt_end = false;
+  // Override the interval and prompt-end flag with increasing checkpoints in
+  // (start, prompt.size()]. Invalid or throwing policies skip capture. Requires
+  // cache capacity. Runs on control; must not block or reenter the runtime.
   std::function<std::vector<batching::PrefixCheckpoint>(
       const batching::PrefixIdentity& prompt,
       batching::Position start,
@@ -86,7 +89,8 @@ struct ET_EXPERIMENTAL ServingRuntimeConfig {
       return max_prefix_captures;
     }
     const auto concurrency = std::min(max_sessions, max_requests);
-    return prefix_checkpoint_interval || prefix_checkpoint_policy
+    return prefix_checkpoint_interval || prefix_capture_prompt_end ||
+            prefix_checkpoint_policy
         ? std::min(concurrency, std::size_t{2})
         : concurrency;
   }