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..6831dd86f37 100644 --- a/examples/llm_server/python/serve.py +++ b/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 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