From 083e88cd3da17edc612d1f32c56f87f7bd04a973 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Mon, 2 Nov 2026 17:45:14 -0500 Subject: [PATCH 1/2] Reuse committed prompt prefixes from active text requests Allow anonymous text requests to clone a matching committed prompt prefix from an active request when it saves enough prefill work. Complete clones through the serving control loop, retain cached and cold fallback paths, and keep the feature disabled by default while evaluating performance. Authored with assistance from Codex. --- examples/llm_server/README.md | 14 + .../llm_server/cpp/module_worker_main.cpp | 6 + .../llm_server/cpp/multiplexed_worker.cpp | 14 + examples/llm_server/python/serve.py | 18 + .../python/tests/test_serve_launcher.py | 18 +- .../models/muse-glimmer/BATCHING_SERVING.md | 14 + .../runtime/runners/batching_bootstrap.cpp | 6 + .../muse-glimmer/serving/serve_batching.py | 7 + .../muse-glimmer/tests/test_serve_batching.py | 25 ++ extension/llm/batching/prefix_cache.h | 20 +- extension/llm/batching/runner.cpp | 92 +++-- extension/llm/batching/runner.h | 7 + .../llm/batching/test/prefix_cache_test.cpp | 20 + extension/llm/batching/test/runner_test.cpp | 37 ++ extension/llm/serving/serving_runtime.cpp | 273 ++++++++++++-- extension/llm/serving/serving_runtime.h | 4 + .../llm/serving/test/prefix_reuse_test.cpp | 349 +++++++++++++++++- extension/llm/serving/types.h | 15 + 18 files changed, 877 insertions(+), 62 deletions(-) diff --git a/examples/llm_server/README.md b/examples/llm_server/README.md index 960710df3d2..3e7fb8369e4 100644 --- a/examples/llm_server/README.md +++ b/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. diff --git a/examples/llm_server/cpp/module_worker_main.cpp b/examples/llm_server/cpp/module_worker_main.cpp index a3f2f8b2d75..a6e016a7027 100644 --- a/examples/llm_server/cpp/module_worker_main.cpp +++ b/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; diff --git a/examples/llm_server/cpp/multiplexed_worker.cpp b/examples/llm_server/cpp/multiplexed_worker.cpp index 4418a52a52c..5876ba7a387 100644 --- a/examples/llm_server/cpp/multiplexed_worker.cpp +++ b/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; diff --git a/examples/llm_server/python/serve.py b/examples/llm_server/python/serve.py index a9a0a3a6f1e..f6877ca4bbe 100644 --- a/examples/llm_server/python/serve.py +++ b/examples/llm_server/python/serve.py @@ -37,6 +37,10 @@ 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 not 0 <= args.active_prefix_reuse_min_tokens <= _MAX_INT32: + parser.error( + f"--active-prefix-reuse-min-tokens 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 +105,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 +182,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 diff --git a/examples/llm_server/python/tests/test_serve_launcher.py b/examples/llm_server/python/tests/test_serve_launcher.py index d1725069c1d..56dd331f8c5 100644 --- a/examples/llm_server/python/tests/test_serve_launcher.py +++ b/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]) diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index e20422bba73..a2d89c2f57d 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/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. diff --git a/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp b/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp index afa9b93a9fd..30d4ba4a992 100644 --- a/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp +++ b/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 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; diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index e5d13f10461..f4529b2f997 100644 --- a/examples/models/muse-glimmer/serving/serve_batching.py +++ b/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, diff --git a/examples/models/muse-glimmer/tests/test_serve_batching.py b/examples/models/muse-glimmer/tests/test_serve_batching.py index 9b1c767c6de..eb174f73663 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/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] + ) diff --git a/extension/llm/batching/prefix_cache.h b/extension/llm/batching/prefix_cache.h index 7bf5940b28f..63102f575ae 100644 --- a/extension/llm/batching/prefix_cache.h +++ b/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 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& tokens) const { if (!max_entries_ || tokens.empty()) { return {}; diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index bb504fec540..8569779942d 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -31,6 +31,22 @@ namespace batching { namespace { +void complete_clone( + std::function)> callback, + std::optional result) noexcept { +#if ET_HAS_EXCEPTIONS + try { +#endif + if (callback) { + callback(std::move(result)); + } +#if ET_HAS_EXCEPTIONS + } catch (...) { + // Delivery is attempted once. Unclaimed sessions close through RAII. + } +#endif +} + // Everything the runner needs from one answered input, checked together so a // broken executor is reported the same way whether or not this step ends the // generation. `room` is how far the session can still advance, so a run that @@ -301,6 +317,10 @@ class RunnerImpl : public std::enable_shared_from_this { std::future> clone_async( SessionId source, Position upto); + void clone_async( + SessionId source, + Position upto, + std::function)> on_complete); void request_close(SessionId session) noexcept; GenerationHandle generate_async( SessionId session, @@ -434,7 +454,7 @@ class RunnerImpl : public std::enable_shared_from_this { struct CloneCommand { SessionId source; Position upto; - std::shared_ptr>> ack; + std::function)> on_complete; }; struct StartCommand { @@ -613,13 +633,22 @@ Position Session::position() const noexcept { } std::future> Session::clone_async(Position upto) const { + auto ack = std::make_shared>>(); + auto future = ack->get_future(); + clone_async(upto, [ack](std::optional result) { + ack->set_value(std::move(result)); + }); + return future; +} + +void Session::clone_async( + Position upto, + std::function)> on_complete) const { if (!valid()) { - std::promise> ack; - auto future = ack.get_future(); - ack.set_value(std::nullopt); - return future; + complete_clone(std::move(on_complete), std::nullopt); + return; } - return state_->impl->clone_async(state_->session, upto); + state_->impl->clone_async(state_->session, upto, std::move(on_complete)); } std::function>(Position)> @@ -931,14 +960,14 @@ void RunnerImpl::process_command_(CloneCommand command) { // admission still gives this clone access until that command retires source. if (!is_running_() || source == sessions_.end() || source->second.poisoned || command.upto < 0 || command.upto > source->second.position()) { - command.ack->set_value(std::nullopt); + complete_clone(std::move(command.on_complete), std::nullopt); return; } std::optional sid; + std::optional result; bool owned = false; bool registered = false; - bool published = false; #if ET_HAS_EXCEPTIONS try { #endif @@ -956,27 +985,25 @@ void RunnerImpl::process_command_(CloneCommand command) { record.status = status; sessions_.emplace(*sid, std::move(record)); registered = true; - command.ack->set_value(Session(std::make_unique( - shared_from_this(), *sid, std::move(status)))); - published = true; + result = Session(std::make_unique( + shared_from_this(), *sid, std::move(status))); } } } #if ET_HAS_EXCEPTIONS } catch (...) { - // Cloning is optional. A failed allocation or backend clone must not stop - // the engine or strand the caller's future. + // Cloning is optional. Backend or allocation failure must not stop work. } #endif - if (!published) { + if (!result) { if (registered) { sessions_.erase(*sid); } if (owned) { executor_.close_session(*sid); } - command.ack->set_value(std::nullopt); } + complete_clone(std::move(command.on_complete), std::move(result)); } void RunnerImpl::process_command_(StartCommand command) { @@ -1371,16 +1398,35 @@ std::future> RunnerImpl::clone_async( Position upto) { auto ack = std::make_shared>>(); auto future = ack->get_future(); - { - std::lock_guard lock(control_mutex_); - if (lifecycle_.load(std::memory_order_relaxed) != Lifecycle::Running) { - ack->set_value(std::nullopt); - return future; + clone_async(source, upto, [ack](std::optional result) { + ack->set_value(std::move(result)); + }); + return future; +} + +void RunnerImpl::clone_async( + SessionId source, + Position upto, + std::function)> on_complete) { +#if ET_HAS_EXCEPTIONS + try { +#endif + std::unique_lock lock(control_mutex_); + if (lifecycle_.load(std::memory_order_relaxed) == Lifecycle::Running) { + // Allocate the command before moving its completion, so allocation + // failure can still settle the caller. + inbox_.emplace_back(CloneCommand{source, upto, {}}); + std::get(inbox_.back()).on_complete = + std::move(on_complete); + lock.unlock(); + notify_engine_(); + return; } - inbox_.emplace_back(CloneCommand{source, upto, std::move(ack)}); +#if ET_HAS_EXCEPTIONS + } catch (...) { } - notify_engine_(); - return future; +#endif + complete_clone(std::move(on_complete), std::nullopt); } void RunnerImpl::request_close(SessionId session) noexcept { diff --git a/extension/llm/batching/runner.h b/extension/llm/batching/runner.h index 435e2efd115..7b0b05c1f8d 100644 --- a/extension/llm/batching/runner.h +++ b/extension/llm/batching/runner.h @@ -220,6 +220,13 @@ class ET_EXPERIMENTAL Session { // Never wait for the returned future in a callback serviced by this runner. std::future> clone_async(Position upto) const; + // Same operation, with completion outside runner locks. Admission failure + // may complete inline; accepted commands complete on the engine thread. + // Must not block or wait for runner work. Exceptions are contained. + void clone_async( + Position upto, + std::function)> on_complete) const; + // `delta` is the caller-resolved suffix to append. The session tracks its // committed position and any pending generated token, so consecutive // generations continue where prior executor work left them. Retain this diff --git a/extension/llm/batching/test/prefix_cache_test.cpp b/extension/llm/batching/test/prefix_cache_test.cpp index 54c96937406..979ea0b484d 100644 --- a/extension/llm/batching/test/prefix_cache_test.cpp +++ b/extension/llm/batching/test/prefix_cache_test.cpp @@ -307,6 +307,26 @@ class PrefixCacheTest : public ::testing::Test { std::unique_ptr runner; }; +TEST_F(PrefixCacheTest, MatchLengthIsMetadataOnlyAndLeavesLastToken) { + PrefixCache cache(2); + ASSERT_TRUE(cache.insert({1, 2}, snapshot({1, 2}))); + ASSERT_TRUE(cache.insert({1, 2, 3, 4}, snapshot({1, 2, 3, 4}))); + const auto before = executor.clones().size(); + auto match_length = [&](Tokens tokens) { + return cache.best_match_length( + token_identity(std::make_shared(std::move(tokens)))); + }; + EXPECT_EQ(match_length({1, 2, 3, 4}), 3u); + EXPECT_EQ(match_length({1, 2, 3, 4, 5}), 4u); + EXPECT_EQ(match_length({1, 8}), 1u); + EXPECT_EQ(match_length({1}), 0u); + EXPECT_EQ(match_length({}), 0u); + EXPECT_EQ(executor.clones().size(), before); + executor.reject_clones = true; + EXPECT_EQ(match_length({1, 2, 3, 4}), 3u); + EXPECT_FALSE(cache.lookup({1, 2, 3, 4})); +} + class PrefillBreakpointTest : public PrefixCacheTest, public ::testing::WithParamInterface {}; diff --git a/extension/llm/batching/test/runner_test.cpp b/extension/llm/batching/test/runner_test.cpp index d85118ba689..f40a8b3151b 100644 --- a/extension/llm/batching/test/runner_test.cpp +++ b/extension/llm/batching/test/runner_test.cpp @@ -1198,6 +1198,43 @@ TEST(CloneTest, DefaultMovedFromAndClosedSessionsRefuseSynchronously) { EXPECT_TRUE(executor.clones().empty()); } +TEST(CloneTest, CallbackRejectsInlineAndContainsExceptions) { + Session empty; + int calls = 0; + empty.clone_async(0, [&](std::optional result) { + ++calls; + EXPECT_FALSE(result); +#if ET_HAS_EXCEPTIONS + throw std::runtime_error("callback failed"); +#endif + }); + EXPECT_EQ(calls, 1); +} + +TEST(CloneTest, CallbackCanReenterRunnerAndDropResult) { + CloneExecutor executor; + Fixture fixture(executor); + auto source = open(fixture.runner); + auto updates = std::make_shared(); + auto generation = generate(source, tokens(3), config(1), updates); + generation.wait(); + std::promise>> reentered; + source.clone_async(2, [&](std::optional result) { + EXPECT_TRUE(result); + reentered.set_value(fixture.runner.open_session_async()); +#if ET_HAS_EXCEPTIONS + throw std::runtime_error("discard result"); +#endif + }); + auto ready = reentered.get_future(); + ASSERT_EQ(ready.wait_for(kTimeout), std::future_status::ready); + auto opened = ready.get(); + ASSERT_EQ(opened.wait_for(kTimeout), std::future_status::ready); + EXPECT_TRUE(opened.get()); + fixture.runner.shutdown(); + EXPECT_EQ(executor.open_count(), 0); +} + TEST(CloneTest, DefaultExecutorRefusesWithoutPoisoningSource) { FakeExecutor executor; Fixture fixture(executor); diff --git a/extension/llm/serving/serving_runtime.cpp b/extension/llm/serving/serving_runtime.cpp index be98d5b9366..9bab7fd843f 100644 --- a/extension/llm/serving/serving_runtime.cpp +++ b/extension/llm/serving/serving_runtime.cpp @@ -224,7 +224,13 @@ struct ServingRuntime::Impl { using SessionKey = std::variant; enum class Operation { Open, Close, Reset, Generate }; enum class Lifecycle { Running, Stopping, Stopped }; - enum class SessionPhase { Opening, Open, Reopening, Unavailable }; + enum class SessionPhase { + Opening, + ClonePending, + Open, + Reopening, + Unavailable + }; struct Slot { SessionPhase phase = SessionPhase::Opening; @@ -235,6 +241,23 @@ struct ServingRuntime::Impl { bool dirty = false; }; + struct ActivePrefix { + const batching::Session* source = + nullptr; // Used only within this control turn. + RequestId donor = 0; + std::size_t matched = 0; + std::size_t shared = 0; + std::size_t committed = 0; + }; + + struct CloneCompletion { + Request request; + std::uint64_t incarnation; + std::optional session; + bool cache_tried; + double elapsed_ms; + }; + struct Command { Operation operation; std::string key; @@ -270,6 +293,10 @@ struct ServingRuntime::Impl { config.prefix_cache_capacity, config.prefix_cache_retention) { if (valid_config_) { + if (config_.active_prefix_reuse_min_tokens) { + // Each admitted recipient can own at most one undelivered result. + clone_completions_.reserve(config_.max_requests); + } runner_ = std::make_unique(executor, std::move(scheduler)); #if ET_HAS_EXCEPTIONS @@ -623,6 +650,12 @@ struct ServingRuntime::Impl { text->emit(std::move(text->terminal)); return; } + if (!update.tokens.empty() && !text->terminal.stats.first_token_ms) { + text->terminal.stats.first_token_ms = + std::chrono::duration( + batching::MetricsClock::now() - text->submitted) + .count(); + } text->raw_tokens.insert( text->raw_tokens.end(), update.tokens.begin(), update.tokens.end()); if (text->output && !text->render_error) { @@ -1063,26 +1096,148 @@ struct ServingRuntime::Impl { } it->second.active_request = request->id; request->completion.incarnation = it->second.incarnation; + lock.unlock(); if (opening) { - lock.unlock(); - std::optional opened; - LifecycleResult preparation_error; + open_generation(request); + } else { + start_generation(request, false); + } + } + + static bool active_prefix_eligible(const Request& request) { + return !request->request.key && request->text && request->text->prompt && + request->text->prepared && + request->text->prepared->kind() == detail::TokenPreparedInput::type(); + } + + ActivePrefix select_active_prefix(const Request& request) { + ActivePrefix best; + if (!config_.active_prefix_reuse_min_tokens || + !active_prefix_eligible(request) || request->text->prompt->size() < 2) { + return best; + } + auto& stats = request->text->terminal.stats.active_prefix; + std::lock_guard lock(mutex_); + for (const auto& entry : sessions_) { + const auto& slot = entry.second; + if (slot.phase != SessionPhase::Open || !slot.session || + !slot.session->valid() || !slot.active_request || + slot.active_request == request->id) { + continue; + } + auto it = requests_.find(slot.active_request); + if (it == requests_.end()) { + continue; + } + const auto& donor = it->second; + if (!active_prefix_eligible(donor) || !donor->text->started || + donor->cancelled.load()) { + continue; + } + { + std::lock_guard donor_lock(donor->mutex); + if (donor->settled) { + continue; + } + } + const auto shared = + batching::common_prefix(*donor->text->prompt, *request->text->prompt); + if (!shared) { + continue; + } + ++stats.candidates; + const auto committed = static_cast(slot.session->position()); + const auto matched = + std::min({shared, committed, request->text->prompt->size() - 1}); + if (matched < config_.active_prefix_reuse_min_tokens) { + ++stats.insufficient_committed; + continue; + } + if (!best.source || matched > best.matched || + (matched == best.matched && donor->id < best.donor)) { + best = {&*slot.session, donor->id, matched, shared, committed}; + } + } + return best; + } + + void clone_active_prefix( + const Request& request, + const ActivePrefix& candidate, + bool cache_tried) { + auto& stats = request->text->terminal.stats.active_prefix; + stats.attempted = true; + stats.matched_tokens = candidate.matched; + stats.shared_tokens = candidate.shared; + stats.committed_tokens = candidate.committed; + const auto incarnation = request->completion.incarnation; + const auto started = batching::MetricsClock::now(); + // Build the callback before publishing pending state. Runner contains all + // failures after admission and invokes it exactly once, possibly inline. + std::function)> callback = + [this, request, incarnation, cache_tried, started](auto session) { + const auto elapsed = std::chrono::duration( + batching::MetricsClock::now() - started) + .count(); + { + std::lock_guard lock(mutex_); + clone_completions_.push_back( + {request, + incarnation, + std::move(session), + cache_tried, + elapsed}); + } + cv_.notify_one(); + }; + { + std::lock_guard lock(mutex_); + sessions_.at(SessionKey(request->id)).phase = SessionPhase::ClonePending; + } + candidate.source->clone_async( + static_cast(candidate.matched), + std::move(callback)); + } + + void open_generation( + const Request& request, + bool cache_tried = false, + std::optional match = std::nullopt, + bool active_tried = false) { + std::optional opened; + LifecycleResult preparation_error; #if ET_HAS_EXCEPTIONS - try { + try { #endif - auto match = lookup_prefix(request); + if (!request->cancelled.load()) { + const auto candidate = + active_tried ? ActivePrefix{} : select_active_prefix(request); + const auto cached = candidate.source + ? prefix_cache_.best_match_length(*request->text->prompt) + : 0; + if (candidate.source && candidate.matched >= cached && + candidate.matched - cached >= + config_.active_prefix_reuse_min_tokens) { + clone_active_prefix(request, candidate, false); + return; + } + if (!match && !cache_tried) { + match = lookup_prefix(request); + cache_tried = true; + } + if (!match && candidate.source) { + clone_active_prefix(request, candidate, cache_tried); + return; + } Selection outcome = Selection::Selected; if (match) { opened = std::move(match->session); - // Snapshot positions are committed; named history also includes - // pending. Keep the accepted full delta until suffix success. auto selection = select_input(request, match->matched_tokens); if (auto* error = std::get_if(&selection)) { preparation_error = std::move(*error); } else { outcome = std::get(selection); if (outcome == Selection::Unavailable) { - // Queue clone closure outside mutex_ before the cold open. opened.reset(); } } @@ -1091,33 +1246,76 @@ struct ServingRuntime::Impl { !request->cancelled.load()) { opened = runner_->open_session_async().get(); } -#if ET_HAS_EXCEPTIONS - } catch (...) { - preparation_error = - ServingError{ErrorCode::Internal, "session preparation failed"}; } +#if ET_HAS_EXCEPTIONS + } catch (...) { + preparation_error = + ServingError{ErrorCode::Internal, "session preparation failed"}; + } #endif - lock.lock(); - const bool cancelled = - lifecycle_ != Lifecycle::Running || request->cancelled.load(); - if (!opened || cancelled || preparation_error) { + const SessionKey key = request->request.key + ? SessionKey(*request->request.key) + : SessionKey(request->id); + std::unique_lock lock(mutex_); + auto it = sessions_.find(key); + const bool current = it != sessions_.end() && + it->second.incarnation == request->completion.incarnation && + it->second.active_request == request->id; + const bool cancelled = !current || lifecycle_ != Lifecycle::Running || + request->cancelled.load(); + if (!opened || cancelled || preparation_error) { + if (current) { sessions_.erase(it); - const auto code = runner_->initialization_state() == - batching::InitializationState::Failed - ? ErrorCode::NotReady - : ErrorCode::CapacityExceeded; - lock.unlock(); - opened.reset(); - request->reject( - cancelled ? std::nullopt - : preparation_error - ? std::move(preparation_error) - : std::optional({code, "session open failed"})); + } + const auto code = runner_->initialization_state() == + batching::InitializationState::Failed + ? ErrorCode::NotReady + : ErrorCode::CapacityExceeded; + lock.unlock(); + opened.reset(); + request->reject( + cancelled ? std::nullopt + : preparation_error + ? std::move(preparation_error) + : std::optional({code, "session open failed"})); + return; + } + it->second.session = std::move(opened); + it->second.phase = SessionPhase::Open; + lock.unlock(); + start_generation(request, true); + } + + void resume_clone(CloneCompletion completion) { + const auto& request = completion.request; + { + std::lock_guard lock(mutex_); + auto it = sessions_.find(SessionKey(request->id)); + if (it == sessions_.end() || + it->second.incarnation != completion.incarnation || + it->second.active_request != request->id || + it->second.phase != SessionPhase::ClonePending) { return; } - it->second.session = std::move(opened); - it->second.phase = SessionPhase::Open; } + auto& stats = request->text->terminal.stats.active_prefix; + stats.clone_ms = completion.elapsed_ms; + stats.reused = completion.session.has_value() && !request->cancelled.load(); + std::optional match; + if (completion.session) { + match = batching::PrefixMatch{ + std::move(*completion.session), stats.matched_tokens}; + } + open_generation(request, completion.cache_tried, std::move(match), true); + schedule(request); + } + + void start_generation(const Request& request, bool opening) { + const SessionKey key = request->request.key + ? SessionKey(*request->request.key) + : SessionKey(request->id); + std::unique_lock lock(mutex_); + auto it = sessions_.find(key); auto& session = *it->second.session; if (!request->text) { it->second.dirty = true; @@ -1549,15 +1747,19 @@ struct ServingRuntime::Impl { std::size_t command_index = 0; Request finalizer; Request checkpoint; + std::optional clone; { std::unique_lock lock(mutex_); cv_.wait(lock, [this] { - return pending_checkpoint_ || !finalizers_.empty() || - runnable_command() != inbox_.end() || + return !clone_completions_.empty() || pending_checkpoint_ || + !finalizers_.empty() || runnable_command() != inbox_.end() || (lifecycle_ != Lifecycle::Running && inbox_.empty() && requests_.empty() && !retiring_); }); - if (pending_checkpoint_) { + if (!clone_completions_.empty()) { + clone = std::move(clone_completions_.front()); + clone_completions_.erase(clone_completions_.begin()); + } else if (pending_checkpoint_) { checkpoint = std::move(pending_checkpoint_); } else if (!finalizers_.empty()) { finalizer = std::move(finalizers_.front()); @@ -1570,6 +1772,10 @@ struct ServingRuntime::Impl { break; } } + if (clone) { + resume_clone(std::move(*clone)); + continue; + } if (checkpoint) { collect_capture(checkpoint); capture_lanes_used_.fetch_sub(1); @@ -1709,6 +1915,7 @@ struct ServingRuntime::Impl { Lifecycle lifecycle_ = Lifecycle::Running; std::size_t outstanding_ = 0; std::deque inbox_; + std::vector clone_completions_; std::unordered_map sessions_; std::uint64_t next_incarnation_ = 1; RequestId next_request_id_ = 1; diff --git a/extension/llm/serving/serving_runtime.h b/extension/llm/serving/serving_runtime.h index 907d1d8db5a..c48b3950895 100644 --- a/extension/llm/serving/serving_runtime.h +++ b/extension/llm/serving/serving_runtime.h @@ -78,6 +78,10 @@ struct ET_EXPERIMENTAL ServingRuntimeConfig { // serialized; one replacement may overlap. Positive values override sizing. std::size_t max_prefix_captures = 0; + // Experimental: minimum additional prompt tokens saved by cloning an active + // anonymous text request. Zero disables; never waits for future progress. + std::size_t active_prefix_reuse_min_tokens = 0; + std::size_t resolved_prefix_captures() const { if (prefix_cache_capacity == 0) { return 0; diff --git a/extension/llm/serving/test/prefix_reuse_test.cpp b/extension/llm/serving/test/prefix_reuse_test.cpp index 8170c8e6e73..c1cdbb8cb37 100644 --- a/extension/llm/serving/test/prefix_reuse_test.cpp +++ b/extension/llm/serving/test/prefix_reuse_test.cpp @@ -164,6 +164,7 @@ class Executor : public batching::testing::FakeExecutor { Gate cloning; Gate decoding; std::atomic refuse_clone{false}; + std::atomic refused_source{0}; std::atomic clone_calls{0}; std::atomic peak_rows{0}; bool reject_after_decode = false; @@ -205,7 +206,7 @@ class Executor : public batching::testing::FakeExecutor { std::optional clone(SessionId source, Position upto) override { ++clone_calls; cloning.enter(); - if (refuse_clone.load()) { + if (refuse_clone.load() || source == refused_source.load()) { return std::nullopt; } std::vector prefix; @@ -476,6 +477,352 @@ class PrefixReuseTest : public ::testing::Test { } }; +class ActivePrefixTest : public PrefixReuseTest { + protected: + void SetUp() override { + PrefixReuseTest::SetUp(); + config.prefix_cache_capacity = 0; + config.active_prefix_reuse_min_tokens = 2; + watch_callbacks(); + } + + // This request fails in preparation without consulting the engine. Its + // completion proves that control got past a preceding clone submission. + void control_barrier() { + auto event = output(); + auto request = submit(event, {}, std::nullopt); + request.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->finish_reason, FinishReason::Failed); + } + + RequestHandle donor(std::optional key = std::nullopt) { + executor.decoding.hold(); + auto request = submit(output(), {1, 2, 3, 4}, key, 4); + EXPECT_TRUE(executor.decoding.wait()); + return request; + } +}; + +TEST_F(ActivePrefixTest, BranchesHaveOwnSuffixSamplingAndBudgets) { + start(); + auto source = donor(); + auto left = output(); + auto right = output(); + GenerationOptions options; + options.max_new_tokens = 2; + options.sampling.temperature = 0.8f; + options.seed = 99; + auto b = submit(left, {1, 2, 3, 9}, std::nullopt, options); + options.max_new_tokens = 1; + options.sampling.temperature = 0; + options.seed = 100; + auto c = submit(right, {1, 2, 3, 8}, std::nullopt, options); + control_barrier(); + executor.decoding.release(); + source.wait(); + b.wait(); + c.wait(); + ASSERT_EQ(executor.clones().size(), 2u); + for (const auto& clone : executor.clones()) { + EXPECT_EQ(clone.prefix, (std::vector{1, 2, 3})); + } + for (const auto& event : {left, right}) { + ASSERT_TRUE(event->terminal); + const auto& stats = event->terminal->stats; + EXPECT_TRUE(stats.active_prefix.reused); + EXPECT_EQ(stats.reused_prompt_tokens, 3u); + EXPECT_EQ(stats.prefilled_prompt_tokens, 1u); + EXPECT_EQ(stats.active_prefix.committed_tokens, 4u); + EXPECT_TRUE(stats.first_token_ms); + EXPECT_GE(*stats.first_token_ms, stats.active_prefix.clone_ms); + } + EXPECT_EQ(left->terminal->stats.completion_tokens, 2u); + EXPECT_EQ(right->terminal->stats.completion_tokens, 1u); + for (const auto& feed : executor.feeds()) { + if (feed.position == 3) { + EXPECT_TRUE( + feed.tokens == std::vector{9} || + feed.tokens == std::vector{8}); + ASSERT_TRUE(feed.sampling); + EXPECT_FLOAT_EQ(feed.sampling->temperature, feed.seed == 99 ? 0.8f : 0); + } + } +} + +TEST_F(ActivePrefixTest, IdenticalPromptLeavesFinalTokenForForward) { + start(); + auto source = donor(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 4}, std::nullopt); + control_barrier(); + executor.decoding.release(); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->stats.active_prefix.shared_tokens, 4u); + EXPECT_EQ(event->terminal->stats.reused_prompt_tokens, 3u); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 1u); +} + +TEST_F(ActivePrefixTest, DoesNotMatchGeneratedOutput) { + start(); + auto source = donor(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 4, 100, 101, 9}, std::nullopt); + control_barrier(); + executor.decoding.release(); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->stats.active_prefix.shared_tokens, 4u); + EXPECT_EQ(event->terminal->stats.reused_prompt_tokens, 4u); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 3u); +} + +TEST_F(ActivePrefixTest, RefusedCloneFallsBackOnceWithFullPrompt) { + start(); + executor.refuse_clone = true; + auto source = donor(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt, 2); + control_barrier(); + executor.decoding.release(); + branch.wait(); + source.wait(); + EXPECT_EQ(executor.clone_calls.load(), 1); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->finish_reason, FinishReason::Length); + EXPECT_TRUE(event->terminal->stats.active_prefix.attempted); + EXPECT_FALSE(event->terminal->stats.active_prefix.reused); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 4u); + EXPECT_EQ(event->terminal->stats.completion_tokens, 2u); +} + +TEST_F(ActivePrefixTest, AdvancingBeyondRetainedBoundaryFallsBack) { + executor.retained_tail = 1; + start(); + auto source = donor(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + executor.decoding.release(); + branch.wait(); + source.wait(); + EXPECT_EQ(executor.clone_calls.load(), 1); + EXPECT_TRUE(executor.clones().empty()); + ASSERT_TRUE(event->terminal); + EXPECT_FALSE(event->terminal->stats.active_prefix.reused); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 4u); +} + +TEST_F( + ActivePrefixTest, + CancellationDuringCloneDiscardsResultAndPreservesDonor) { + start(); + auto source = donor(); + executor.cloning.hold(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + executor.decoding.release(); + ASSERT_TRUE(executor.cloning.wait()); + branch.cancel(); + control_barrier(); + executor.cloning.release(); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminals, 1); + EXPECT_EQ(event->terminal->finish_reason, FinishReason::Cancelled); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 0u); + EXPECT_EQ(events.front()->terminal->finish_reason, FinishReason::Length); +} + +TEST_F(ActivePrefixTest, ShutdownDrainsInFlightClone) { + start(); + auto source = donor(); + executor.cloning.hold(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + executor.decoding.release(); + ASSERT_TRUE(executor.cloning.wait()); + branch.cancel(); + source.cancel(); + auto stopped = std::async(std::launch::async, [&] { runtime->shutdown(); }); + executor.cloning.release(); + ASSERT_EQ(stopped.wait_for(5s), std::future_status::ready); + stopped.get(); + branch.wait(); + source.wait(); + EXPECT_EQ(event->terminals, 1); + EXPECT_EQ(executor.open_count(), 0); +} + +class ActivePrefixExclusionTest : public ActivePrefixTest, + public ::testing::WithParamInterface {}; +TEST_P( + ActivePrefixExclusionTest, + LeavesUnsupportedOrUnhelpfulRequestsUnchanged) { + const auto mode = GetParam(); + if (mode == 0) + config.active_prefix_reuse_min_tokens = 0; + start(); + auto source = + donor(mode == 1 ? std::optional{"named"} : std::nullopt); + auto event = output(); + auto prompt = mode == 3 ? std::vector{8, 9} + : mode == 4 ? std::vector{1} + : std::vector{1, 2, 3, 9}; + auto branch = submit( + event, + prompt, + mode == 2 ? std::optional{"recipient"} : std::nullopt); + // The ordinary cold-open path can await the engine, so do not use a control + // barrier while holding its decode here. + executor.decoding.release(); + branch.wait(); + source.wait(); + EXPECT_EQ(executor.clone_calls.load(), 0); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, prompt.size()); +} +INSTANTIATE_TEST_SUITE_P( + Scope, + ActivePrefixExclusionTest, + ::testing::Range(0, 5)); + +TEST_F(ActivePrefixTest, PartialPrefillUsesCommittedPositionWithoutWaiting) { + auto scheduler = std::make_unique(); + scheduler->inner = batching::DecodeFirstScheduler::create(12, 4, 8); + auto* paused = scheduler.get(); + executor.executing.hold(); + start(true, std::move(scheduler)); + std::vector prompt(20); + std::iota(prompt.begin(), prompt.end(), 1); + auto source = submit(output(), prompt, std::nullopt, 2); + ASSERT_TRUE(executor.executing.wait()); + paused->paused = true; + executor.executing.release(); + // The engine processes this open after the first eight-token forward. + ASSERT_FALSE(runtime->open_session_async("barrier").get()); + auto event = output(); + prompt.back() = 99; + auto branch = submit(event, prompt, std::nullopt); + ASSERT_TRUE(executor.cloning.wait()); + control_barrier(); + paused->paused = false; + // Wake the engine if it went idle while the scheduler was paused. + ASSERT_FALSE(runtime->close_session_async("barrier").get()); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + const auto& stats = event->terminal->stats; + EXPECT_EQ(stats.active_prefix.shared_tokens, 19u); + EXPECT_EQ(stats.active_prefix.committed_tokens, 8u); + EXPECT_EQ(stats.reused_prompt_tokens, 8u); + EXPECT_EQ(stats.prefilled_prompt_tokens, 12u); + ASSERT_EQ(executor.clones().size(), 1u); + EXPECT_EQ(executor.clones()[0].prefix.size(), 8u); +} + +TEST_F(ActivePrefixTest, NoCommittedPrefixDoesNotWaitForDonor) { + auto scheduler = std::make_unique(); + scheduler->paused = true; + auto* paused = scheduler.get(); + start(true, std::move(scheduler)); + auto source = submit(output(), {1, 2, 3, 4}, std::nullopt); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + paused->paused = false; + ASSERT_FALSE(runtime->open_session_async("wake").get()); + branch.wait(); + source.wait(); + EXPECT_EQ(executor.clone_calls.load(), 0); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->stats.active_prefix.insufficient_committed, 1u); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 4u); +} + +class ActivePrefixCacheTest : public ActivePrefixTest, + public ::testing::WithParamInterface {}; +TEST_P(ActivePrefixCacheTest, ChoosesCacheOrActiveByAdditionalSavings) { + config.prefix_cache_capacity = 1; + start(); + const bool tiny = GetParam(); + submit( + output(), + tiny ? std::vector{1} : std::vector{1, 2, 3}, + std::nullopt) + .wait(); + // Keep the initial cache entry while the donor starts from cold input. + executor.refuse_clone = true; + auto source = donor(); + executor.refuse_clone = false; + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + if (tiny) + control_barrier(); + executor.decoding.release(); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_EQ(event->terminal->stats.active_prefix.attempted, tiny); + EXPECT_EQ(event->terminal->stats.active_prefix.reused, tiny); + EXPECT_EQ(event->terminal->stats.reused_prompt_tokens, 3u); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 1u); +} +INSTANTIATE_TEST_SUITE_P(CacheSizes, ActivePrefixCacheTest, ::testing::Bool()); + +TEST_F(ActivePrefixTest, FailedCacheLookupStillTriesActiveDonor) { + config.prefix_cache_capacity = 1; + auto scheduler = std::make_unique(); + auto* paused = scheduler.get(); + start(true, std::move(scheduler)); + submit(output(), {1, 2, 3}, std::nullopt).wait(); + ASSERT_EQ(executor.clones().size(), 1u); + executor.refused_source = executor.clones().back().destination; + executor.refuse_clone = true; + auto source = donor(); + paused->paused = true; + executor.decoding.release(); + ASSERT_FALSE(runtime->open_session_async("barrier").get()); + executor.refuse_clone = false; + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + paused->paused = false; + ASSERT_FALSE(runtime->close_session_async("barrier").get()); + branch.wait(); + source.wait(); + ASSERT_TRUE(event->terminal); + EXPECT_TRUE(event->terminal->stats.active_prefix.reused); + EXPECT_EQ(event->terminal->stats.reused_prompt_tokens, 3u); +} + +TEST_F(ActivePrefixTest, RefusalDoesNotRepeatPreparation) { + std::atomic prepared{0}; + model_preparer = [&](const auto&, + PromptInput input) -> ModelPreparationResult { + ++prepared; + return make_token_prepared_input(input.segments[0].get_tokens()); + }; + start(); + executor.refuse_clone = true; + auto source = donor(); + auto event = output(); + auto branch = submit(event, {1, 2, 3, 9}, std::nullopt); + control_barrier(); + const auto before = prepared.load(); + executor.decoding.release(); + branch.wait(); + source.wait(); + EXPECT_EQ(prepared.load(), before); + EXPECT_EQ(event->terminal->stats.prefilled_prompt_tokens, 4u); +} + class PrefixReuseSamplingTest : public PrefixReuseTest, public ::testing::WithParamInterface> {}; diff --git a/extension/llm/serving/types.h b/extension/llm/serving/types.h index 64949b4d7ba..6f4b05c1658 100644 --- a/extension/llm/serving/types.h +++ b/extension/llm/serving/types.h @@ -109,6 +109,18 @@ using ModelPreparer ET_EXPERIMENTAL = std::function< ModelPreparationResult make_token_prepared_input( std::vector tokens); +struct ET_EXPERIMENTAL ActivePrefixStats { + std::size_t candidates = 0; + std::size_t insufficient_committed = 0; + bool attempted = false; + bool reused = false; + std::size_t matched_tokens = 0; + std::size_t shared_tokens = 0; + std::size_t committed_tokens = 0; + // Includes engine queueing; not isolated GPU copy time. + double clone_ms = 0.0; +}; + struct ET_EXPERIMENTAL GenerationStats { // Full prompt size in decoder positions, including any reused prefix. std::size_t prompt_tokens = 0; @@ -125,6 +137,9 @@ struct ET_EXPERIMENTAL GenerationStats { double prefill_ms = 0.0; double decode_ms = 0.0; double total_ms = 0.0; + // Submission through first generated token delivery, including admission. + std::optional first_token_ms; + ActivePrefixStats active_prefix; std::string session_reset_reason; // Original output IDs counted by completion_tokens, excluding EOS/stop // tokens; not a retokenization of visible text or the full logical history From cf4a0cf6332ceae38cf456bc95984122e7f3ed06 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Mon, 2 Nov 2026 17:56:47 -0500 Subject: [PATCH 2/2] Consolidate prefix option range validation Validate nonnegative prefix options in one loop to keep the launcher within the repository complexity limit without changing the accepted ranges or error messages. The two launcher suites pass all 118 tests. Authored with assistance from Codex. --- examples/llm_server/python/serve.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/examples/llm_server/python/serve.py b/examples/llm_server/python/serve.py index f6877ca4bbe..6831dd86f37 100644 --- a/examples/llm_server/python/serve.py +++ b/examples/llm_server/python/serve.py @@ -33,14 +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}") - if not 0 <= args.active_prefix_reuse_min_tokens <= _MAX_INT32: - parser.error( - f"--active-prefix-reuse-min-tokens 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: