Reuse committed prompt prefixes from active text requests

examples/llm_server/README.md

@@ -144,3 +144,17 @@ request stop strings remain active. Rebuild the multiplexed worker to include th
 support; no chat-template changes are needed. Generation continues after EOS, but
 the HTTP response only parses and displays the first assistant turn. Usage counts
 include all generated tokens.
+
+### Experimental active prompt reuse
+
+Pass `--active-prefix-reuse-min-tokens 1024` to the batching server launcher to
+reuse an already-computed prompt prefix from another active anonymous text
+request. The value is the minimum additional tokens saved over a cached prefix;
+`0` (the default) disables it. The native worker flag uses underscores:
+`--active_prefix_reuse_min_tokens=1024`.
+
+The recipient clones the committed matching prefix and computes its own suffix.
+It never waits for the donor to reach a future position. If the backend cannot
+restore that prefix, including overwritten sliding-window state, the request
+falls back to cached or full prefill. MLX and CUDA use the same serving policy;
+named sessions and multimodal requests keep their existing behavior.

examples/llm_server/cpp/module_worker_main.cpp

@@ -39,6 +39,10 @@ DEFINE_int32(
     prefix_capture_slots,
     0,
     "Transient capture rows; 0: size automatically from session/request limits");
+DEFINE_int32(
+    active_prefix_reuse_min_tokens,
+    0,
+    "Minimum extra tokens saved by active prompt reuse; zero disables");
 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(
@@ -51,6 +55,7 @@ int main(int argc, char** argv) {
   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 ||
+      FLAGS_active_prefix_reuse_min_tokens < 0 ||
       FLAGS_prefix_capture_slots < 0 || FLAGS_prefix_checkpoint_interval < 0 ||
       (FLAGS_prefix_checkpoint_interval > 0 &&
        FLAGS_prefix_cache_entries == 0)) {
@@ -62,6 +67,7 @@ int main(int argc, char** argv) {
   }
   executorch::extension::llm::serving::ServingRuntimeConfig config;
   config.max_sessions = FLAGS_max_sessions;
+  config.active_prefix_reuse_min_tokens = FLAGS_active_prefix_reuse_min_tokens;
   config.prefix_cache_capacity = FLAGS_prefix_cache_entries;
   config.max_prefix_captures = FLAGS_prefix_capture_slots;
   config.prefix_checkpoint_interval = FLAGS_prefix_checkpoint_interval;

examples/llm_server/cpp/multiplexed_worker.cpp

@@ -223,6 +223,20 @@ Json terminal_json(const serving::TerminalEvent& event) {
       {"decode_tok_s",
        s.decode_ms > 0 ? s.completion_tokens * 1000.0 / s.decode_ms : 0},
       {"session_reset_reason", s.session_reset_reason}};
+  if (s.first_token_ms)
+    out["first_token_ms"] = *s.first_token_ms;
+  if (s.active_prefix.candidates || s.active_prefix.attempted) {
+    const auto& active = s.active_prefix;
+    out["active_prefix"] = {
+        {"candidates", active.candidates},
+        {"insufficient_committed", active.insufficient_committed},
+        {"attempted", active.attempted},
+        {"reused", active.reused},
+        {"matched_tokens", active.matched_tokens},
+        {"shared_tokens", active.shared_tokens},
+        {"committed_tokens", active.committed_tokens},
+        {"clone_ms", active.clone_ms}};
+  }
   if (s.generated_token_ids)
     out["generated_token_ids"] = *s.generated_token_ids;
   return out;

examples/llm_server/python/serve.py

@@ -33,10 +33,13 @@ def _validate_limits(parser, args, *, prefix_capture_slots=0) -> int:
     ):
         if not 1 <= value <= _MAX_INT32:
             parser.error(f"{flag} must be between 1 and {_MAX_INT32}")
-    if not 0 <= args.prefix_cache_entries <= _MAX_INT32:
-        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}")
+    for flag, value in (
+        ("--prefix-cache-entries", args.prefix_cache_entries),
+        ("--prefix-checkpoint-interval", args.prefix_checkpoint_interval),
+        ("--active-prefix-reuse-min-tokens", args.active_prefix_reuse_min_tokens),
+    ):
+        if not 0 <= value <= _MAX_INT32:
+            parser.error(f"{flag} 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 not 0 <= prefix_capture_slots <= _MAX_INT32:
@@ -101,6 +104,12 @@ def _parse_args(argv=None):
         default=0,
         help="Checkpoint interval in tokens; 0: initial prompt only. Uses global LRU.",
     )
+    parser.add_argument(
+        "--active-prefix-reuse-min-tokens",
+        type=int,
+        default=0,
+        help="Minimum additional tokens saved by active prompt reuse; 0 disables",
+    )
     parser.add_argument("--host", default="127.0.0.1")
     parser.add_argument("--port", type=int, default=8000)
     args = parser.parse_args(argv)
@@ -172,6 +181,14 @@ async def serving_factory():
                 if args.prefix_cache_entries
                 else []
             )
+            + (
+                [
+                    "--active_prefix_reuse_min_tokens",
+                    str(args.active_prefix_reuse_min_tokens),
+                ]
+                if args.active_prefix_reuse_min_tokens
+                else []
+            )
             + (
                 ["--prefix_checkpoint_interval", str(args.prefix_checkpoint_interval)]
                 if args.prefix_checkpoint_interval

examples/llm_server/python/tests/test_serve_launcher.py

@@ -195,6 +195,7 @@ def test_capture_capacity_validation_and_worker_override(
         max_inflight_requests=requests,
         prefix_cache_entries=1,
         prefix_checkpoint_interval=interval,
+        active_prefix_reuse_min_tokens=0,
         port=8000,
     )
     assert (
@@ -771,9 +772,7 @@ def test_launcher_rejects_legacy_readiness_and_reaps_real_child(
     real_spawn = asyncio.create_subprocess_exec
     readiness = json.dumps({"ready": True, **capability})
     script = (
-        "import sys\n"
-        f"print({readiness!r}, flush=True)\n"
-        "for line in sys.stdin: pass\n"
+        f"import sys\nprint({readiness!r}, flush=True)\nfor line in sys.stdin: pass\n"
     )
 
     async def spawn(*command, **kwargs):
@@ -810,3 +809,16 @@ async def scenario(app):
 
     launch.stages["uvicorn"].side_effect = lambda app, **kwargs: _run(scenario(app))
     serve.main(argv)
+
+
+def test_active_prefix_reuse_flag_reaches_worker(argv, launch):
+    serve.main(argv + ["--active-prefix-reuse-min-tokens", "1024"])
+    command = launch.stages["spawn"].await_args.args[0]
+    index = command.index("--active_prefix_reuse_min_tokens")
+    assert command[index + 1] == "1024"
+
+
+@pytest.mark.parametrize("value", ["-1", str(2**31)])
+def test_active_prefix_reuse_rejects_invalid_limits(argv, value):
+    with pytest.raises(SystemExit):
+        serve._parse_args(argv + ["--active-prefix-reuse-min-tokens", value])

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -240,3 +240,17 @@ and `--max-draft-owners`), then add `--dflash` to the serve command above. The
 draft ABI and `--dflash-n-draft` behave as on MLX. Drafts are greedy; the target
 accepts or resamples them on the device under each request's sampling. Decode
 drafts of up to four owners run in CUDA graphs.
+
+### Experimental active prompt reuse
+
+Pass `--active-prefix-reuse-min-tokens 1024` to the batching server launcher to
+reuse an already-computed prompt prefix from another active anonymous text
+request. The value is the minimum additional tokens saved over a cached prefix;
+`0` (the default) disables it. The native worker flag uses underscores:
+`--active_prefix_reuse_min_tokens=1024`.
+
+The recipient clones the committed matching prefix and computes its own suffix.
+It never waits for the donor to reach a future position. If the backend cannot
+restore that prefix, including overwritten sliding-window state, the request
+falls back to cached or full prefill. MLX and CUDA use the same serving policy;
+named sessions and multimodal requests keep their existing behavior.

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

@@ -25,6 +25,10 @@ DEFINE_string(backend, "cuda", "Executor backend");
 #else
 DEFINE_string(backend, "mlx", "Executor backend");
 #endif
+DEFINE_int32(
+    active_prefix_reuse_min_tokens,
+    0,
+    "Minimum extra tokens saved by active prompt reuse; zero disables");
 DEFINE_int32(max_sessions, 4, "Maximum logical sessions");
 DEFINE_int32(
     max_session_tokens,
@@ -65,6 +69,7 @@ std::unique_ptr<MuseGlimmerBatchingRuntime>
 create_muse_glimmer_batching_runtime() {
   if (FLAGS_pte.empty() || FLAGS_tokenizer.empty() || FLAGS_max_sessions <= 0 ||
       FLAGS_max_inflight_requests <= 0 || FLAGS_max_decode_sequences <= 0 ||
+      FLAGS_active_prefix_reuse_min_tokens < 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 &&
@@ -79,6 +84,7 @@ create_muse_glimmer_batching_runtime() {
   }
   serving::ServingRuntimeConfig serving;
   serving.max_sessions = FLAGS_max_sessions;
+  serving.active_prefix_reuse_min_tokens = FLAGS_active_prefix_reuse_min_tokens;
   serving.max_requests = FLAGS_max_inflight_requests;
   serving.max_pending_operations = FLAGS_max_inflight_requests;
   serving.prefix_cache_capacity = FLAGS_prefix_cache_entries;

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

@@ -73,6 +73,12 @@ def _parse_args(argv=None):
         default=0,
         help="Transient capture rows; 0 (default): size from session/request limits",
     )
+    parser.add_argument(
+        "--active-prefix-reuse-min-tokens",
+        type=int,
+        default=0,
+        help="Minimum extra tokens saved by active text prompt reuse; 0 disables",
+    )
     parser.add_argument("--max-vision-patches", type=int, default=4096)
     parser.add_argument(
         "--prefix-checkpoint-interval",
@@ -164,6 +170,7 @@ def _worker_command(args):
         ("max_decode_sequences", args.max_decode_sequences),
         ("max_inflight_requests", args.max_inflight_requests),
         ("prefix_cache_entries", args.prefix_cache_entries),
+        ("active_prefix_reuse_min_tokens", args.active_prefix_reuse_min_tokens or None),
         (
             "prefix_capture_slots",
             args.prefix_capture_slots if args.prefix_cache_entries else None,

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

@@ -242,3 +242,28 @@ async def scenario():
         asyncio.run(scenario())
     assert worker.closed
     assert captured["command"] == serve_batching._worker_command(args)
+
+
+@pytest.mark.parametrize("backend", ["mlx", "cuda"])
+def test_active_prefix_setting_is_shared(arguments, backend, monkeypatch):
+    monkeypatch.setattr(serve_batching, "_validate_backend", lambda *args: None)
+    args = serve_batching._parse_args(
+        arguments
+        + [
+            "--backend",
+            backend,
+            "--active-prefix-reuse-min-tokens",
+            "1024",
+        ]
+    )
+    assert "--active_prefix_reuse_min_tokens=1024" in serve_batching._worker_command(
+        args
+    )
+
+
+@pytest.mark.parametrize("value", ["-1", str(2**31)])
+def test_invalid_active_prefix_setting(arguments, value):
+    with pytest.raises(SystemExit):
+        serve_batching._parse_args(
+            arguments + ["--active-prefix-reuse-min-tokens", value]
+        )

extension/llm/batching/prefix_cache.h

@@ -346,6 +346,15 @@ class ET_EXPERIMENTAL PrefixCache {
     return lookup_identity(copy_tokens(request));
   }
 
+  // Token match only: restoration can still fail. Does not clone or touch LRU.
+  std::size_t best_match_length(const PrefixIdentity& request) const {
+    std::size_t best = 0;
+    for (const auto& entry : entries_) {
+      best = std::max(best, match_length(entry.tokens, request));
+    }
+    return best;
+  }
+
   std::optional<PrefixMatch> lookup_identity(const PrefixIdentity& request) {
     if (request.size() < 2 || entries_.empty()) {
       return std::nullopt;
@@ -362,8 +371,7 @@ class ET_EXPERIMENTAL PrefixCache {
       candidates.reserve(entries_.size());
       std::size_t rank = 0;
       for (auto it = entries_.begin(); it != entries_.end(); ++it, ++rank) {
-        const auto matched =
-            std::min(common_prefix(it->tokens, request), request.size() - 1);
+        const auto matched = match_length(it->tokens, request);
         if (matched > 0) {
           candidates.push_back(Candidate{it, matched, rank});
         }
@@ -404,6 +412,14 @@ class ET_EXPERIMENTAL PrefixCache {
   }
 
  private:
+  static std::size_t match_length(
+      const PrefixIdentity& prefix,
+      const PrefixIdentity& request) {
+    return request.size() < 2
+        ? 0
+        : std::min(common_prefix(prefix, request), request.size() - 1);
+  }
+
   PrefixIdentity copy_tokens(const std::vector<Token>& tokens) const {
     if (!max_entries_ || tokens.empty()) {
       return {};