From 9335ab47289db1f1d4a186d68f8ca2257d7918cf Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 28 Aug 2026 14:03:44 -0700 Subject: [PATCH 1/2] up --- extension/llm/batching/runner.cpp | 253 ++++++++++++++++++-- extension/llm/batching/runner.h | 10 + extension/llm/batching/test/fake_executor.h | 9 +- extension/llm/batching/test/runner_test.cpp | 21 ++ 4 files changed, 274 insertions(+), 19 deletions(-) diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 7a21c3910bd..3fd07a2e142 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -65,6 +65,7 @@ struct GenerationHandleState { std::condition_variable cv; bool done = false; FinishReason reason = FinishReason::Failed; + GenerationMetrics metrics; std::atomic cancelled{false}; }; @@ -79,12 +80,14 @@ struct SessionStatus { bool publish_terminal_state( const std::shared_ptr& state, - FinishReason reason) { + FinishReason reason, + const GenerationMetrics& metrics) { std::lock_guard lock(state->mutex); if (state->done) { return false; } state->reason = reason; + state->metrics = metrics; state->done = true; return true; } @@ -94,7 +97,8 @@ void complete_terminal( const GenerationCallback& on_update, const std::vector& tokens, FinishReason reason) { - if (!publish_terminal_state(state, reason)) { + // Rejected before admission, so there is no timeline to report. + if (!publish_terminal_state(state, reason, GenerationMetrics{})) { return; } state->cv.notify_all(); @@ -135,6 +139,14 @@ FinishReason GenerationHandle::finish_reason() const { return state_->reason; } +GenerationMetrics GenerationHandle::metrics() const { + if (!state_) { + return {}; + } + std::lock_guard lock(state_->mutex); + return state_->metrics; +} + // Everything the runner owns. Held by shared_ptr from both Runner and every // Session, so a Session outliving its Runner finds a stopped object rather // than a dangling one. @@ -169,6 +181,11 @@ class RunnerImpl : public std::enable_shared_from_this { GenConfig config, GenerationCallback on_update); + // Engine-thread data, so only stable once that thread is joined. + EngineMetrics metrics() const { + return metrics_; + } + private: enum class Lifecycle { Running, Stopping, Stopped }; @@ -180,6 +197,10 @@ class RunnerImpl : public std::enable_shared_from_this { // Shared with every handle, so cancelling needs no route back to the // runner and works after it is gone. std::shared_ptr state; + GenerationMetrics m; + // Engine-side only. An inter-token gap needs the previous delivery, and + // the published metrics keep only the summary, not the last timestamp. + MetricsTime last_token_at{}; }; // Start-only data. The sampling policy is installed on the executor at @@ -270,11 +291,22 @@ class RunnerImpl : public std::enable_shared_from_this { std::optional reason); // Publish terminal state and invoke the callback for a detached generation. + // + // `on_engine_thread` is false only on the post-shutdown path out of + // generate_async(), which runs on the caller's thread. The engine may still + // be draining there, so that path publishes to the handle but must leave the + // engine's own counters alone. void complete_generation_( Generation generation, FinishReason reason, + bool on_engine_thread, std::vector final_tokens = {}); - void complete_request_(GenerationRequest request, FinishReason reason); + void complete_request_( + GenerationRequest request, + FinishReason reason, + bool on_engine_thread); + // Engine thread only: rolls one finished generation into metrics_. + void record_completion_(const GenerationMetrics& m, FinishReason reason); void complete_active_generation_( SessionId session, FinishReason reason, @@ -303,6 +335,7 @@ class RunnerImpl : public std::enable_shared_from_this { // Kept after records close to enforce executor IDs are lifetime-unique. std::unordered_set issued_session_ids_; TaskId next_tid_ = 1; + EngineMetrics metrics_; std::thread engine_; }; @@ -387,6 +420,10 @@ void Runner::shutdown() { impl_->shutdown(); } +EngineMetrics Runner::metrics() const { + return impl_->metrics(); +} + std::future> Runner::open_session() { return impl_->open_session(); } @@ -485,7 +522,10 @@ void RunnerImpl::run_() { for (auto& entry : open_sessions) { if (entry.second) { - complete_generation_(std::move(*entry.second), FinishReason::Cancelled); + complete_generation_( + std::move(*entry.second), + FinishReason::Cancelled, + /*on_engine_thread=*/true); } executor_.close_session(entry.first); } @@ -562,7 +602,10 @@ void RunnerImpl::process_pending_commands_() { // left would run against a session the executor has released. (void)scheduler_->cancel(cmd.session); if (active) { - complete_generation_(std::move(*active), FinishReason::Cancelled); + complete_generation_( + std::move(*active), + FinishReason::Cancelled, + /*on_engine_thread=*/true); } executor_.close_session(cmd.session); break; @@ -588,9 +631,104 @@ bool RunnerImpl::execute_one_batch_() { return false; } + // Composition is read here, before to_batch_input moves the Inputs out and + // drops is_decode with the rest of the scheduling fields. + // + // Decode tasks are one sequence each. A session's prefill can arrive as + // several chunks which are not necessarily adjacent -- DecodeFirstScheduler + // rotates, taking one chunk per session per pass, so two sessions prefilling + // together interleave as A, B, A, B. Track every session seen rather than + // comparing with the previous one, which would count each chunk as a new + // sequence and charge the step to the generation repeatedly. + std::uint64_t decode_seqs = 0; + std::uint64_t prefill_seqs = 0; + std::uint64_t decode_tokens = 0; + std::uint64_t prefill_tokens = 0; + std::vector prefilling; // small: bounded by the batch width + const MetricsTime step_start = MetricsClock::now(); + for (const Task& task : tasks) { + bool first_chunk = false; + if (task.is_decode) { + ++decode_seqs; + decode_tokens += task.input.size; + } else { + prefill_tokens += task.input.size; + first_chunk = std::find( + prefilling.begin(), + prefilling.end(), + task.input.sid) == prefilling.end(); + if (first_chunk) { + ++prefill_seqs; + prefilling.push_back(task.input.sid); + } + } + // Charge the step to the generation once, however many chunks it brought. + if (!task.is_decode && !first_chunk) { + continue; + } + auto session = sessions_.find(task.input.sid); + if (session == sessions_.end() || !session->second.active_generation) { + continue; + } + GenerationMetrics& m = session->second.active_generation->m; + if (!stamped(m.t_first_step)) { + m.t_first_step = step_start; + } + if (task.is_decode) { + ++m.n_decode_steps; + } else { + ++m.n_prefill_steps; + } + } + // Generations eligible for a decode slot, admitted or not. Past their first + // token, so a generation still prefilling is not counted as held back when + // it is simply busy elsewhere. Against decode_seqs this is what separates a + // scheduler holding work back from there being no work. + for (const auto& entry : sessions_) { + const auto& generation = entry.second.active_generation; + if (generation && stamped(generation->m.t_first_token)) { + ++metrics_.ready_total; + } + } + BatchInput batch = to_batch_input(tasks); BatchOutput out; const bool ok = executor_.execute(batch, out); + const MetricsTime step_end = MetricsClock::now(); + + const std::int64_t latency = us_between(step_start, step_end); + ++metrics_.steps; + metrics_.decode_seqs_total += decode_seqs; + metrics_.prefill_seqs_total += prefill_seqs; + // Only what the model is known to have taken in. A failed execute leaves + // what it processed unknown -- that is why the batch is condemned and its + // sessions poisoned -- so counting the attempt as throughput would credit + // work that may never have happened. The time is still counted below, + // because it was really spent, and steps_failed records the attempt. + if (ok) { + metrics_.decode_tokens_total += decode_tokens; + metrics_.prefill_tokens_total += prefill_tokens; + } + metrics_.step_latency_sum_us += latency; + metrics_.step_latency_max_us = + std::max(metrics_.step_latency_max_us, latency); + if (!stamped(metrics_.t_first_step)) { + metrics_.t_first_step = step_start; + } + metrics_.t_last_step = step_end; + if (decode_tokens > 0) { + ++metrics_.steps_with_decode; + // Charged once per sequence: each of them waited this whole step. + metrics_.decode_wait_sum_us += + latency * static_cast(decode_seqs); + } + if (prefill_tokens > 0) { + ++metrics_.steps_with_prefill; + } + if (!ok) { + ++metrics_.steps_failed; + } + if (!is_running_()) { return true; // discard an in-flight result after the stop boundary } @@ -713,6 +851,10 @@ GenerationHandle RunnerImpl::generate_async( request.generation.stop_tokens = std::move(config.stop_tokens); request.generation.on_update = std::move(on_update); request.generation.state = state; + // The caller's thread, before the request is queued: the wait a caller sees + // starts here, not when the engine gets round to it. + request.generation.m.sid = session; + request.generation.m.t_submit = MetricsClock::now(); GenerationHandle handle(state); bool admitted = false; @@ -733,18 +875,25 @@ GenerationHandle RunnerImpl::generate_async( // After shutdown nothing drains the inbox, so complete synchronously instead // of admitting a start that can never report completion. - complete_request_(std::move(request), FinishReason::Cancelled); + complete_request_( + std::move(request), FinishReason::Cancelled, /*on_engine_thread=*/false); return handle; } void RunnerImpl::start_generation_(GenerationRequest request) { + // Counted on arrival at the engine, not on successful install: every path + // below ends in complete_generation_, which counts a completion, so + // deferring this would let completions exceed starts. + ++metrics_.generations_started; if (!is_running_() || request.generation.state->cancelled.load()) { - complete_request_(std::move(request), FinishReason::Cancelled); + complete_request_( + std::move(request), FinishReason::Cancelled, /*on_engine_thread=*/true); return; } auto session = sessions_.find(request.session); if (session == sessions_.end()) { - complete_request_(std::move(request), FinishReason::Failed); + complete_request_( + std::move(request), FinishReason::Failed, /*on_engine_thread=*/true); return; } SessionRecord& record = session->second; @@ -755,18 +904,25 @@ void RunnerImpl::start_generation_(GenerationRequest request) { if (request.generation.remaining_tokens <= 0 || !valid_positioned_tokens( start_position, request.delta, record.pending ? 1u : 0u)) { - complete_request_(std::move(request), FinishReason::Failed); + complete_request_( + std::move(request), FinishReason::Failed, /*on_engine_thread=*/true); return; } // A step on this session failed mid-execute, so what the executor holds for // it is unknown. Anything built on that would be silently wrong. if (record.poisoned || record.active_generation) { - complete_request_(std::move(request), FinishReason::Failed); + complete_request_( + std::move(request), FinishReason::Failed, /*on_engine_thread=*/true); return; } executor_.set_sampling(request.session, request.sampling, request.seed); + // The prompt this generation was asked to process. Set before the move + // below, which leaves request.generation empty. + request.generation.m.n_prompt_tokens = + static_cast(request.delta->size()); + bool installed = false; { std::lock_guard lock(control_mutex_); @@ -777,7 +933,8 @@ void RunnerImpl::start_generation_(GenerationRequest request) { } if (!installed) { // Sampling began before the stop transition, but no task was submitted. - complete_request_(std::move(request), FinishReason::Cancelled); + complete_request_( + std::move(request), FinishReason::Cancelled, /*on_engine_thread=*/true); return; } @@ -916,6 +1073,30 @@ void RunnerImpl::handle_output_( record.pending = emit.back(); } + if (!emit.empty()) { + const MetricsTime now = MetricsClock::now(); + generation.m.n_generated_tokens += static_cast(emit.size()); + if (!stamped(generation.m.t_first_token)) { + generation.m.t_first_token = now; + // The first token's wait is TTFT. Further tokens in the same callback + // have zero caller-visible latency between them. + const std::int64_t intra_burst = + static_cast(emit.size()) - 1; + if (intra_burst > 0) { + generation.m.itl_count += intra_burst; + generation.m.itl_min_us = 0; + } + } else { + const std::int64_t gap = us_between(generation.last_token_at, now); + generation.m.itl_count += static_cast(emit.size()); + generation.m.itl_sum_us += gap; + generation.m.itl_min_us = std::min( + generation.m.itl_min_us, emit.size() > 1 ? std::int64_t{0} : gap); + generation.m.itl_max_us = std::max(generation.m.itl_max_us, gap); + } + generation.last_token_at = now; + } + if (ends) { complete_active_generation_(session_id, reason, std::move(emit)); return; @@ -964,7 +1145,9 @@ void RunnerImpl::deliver_update_( void RunnerImpl::complete_generation_( Generation generation, FinishReason reason, + bool on_engine_thread, std::vector final_tokens) { + generation.m.t_end = MetricsClock::now(); bool published = false; { // The stop transition and terminal publication have a total order. User @@ -981,20 +1164,56 @@ void RunnerImpl::complete_generation_( reason = FinishReason::Cancelled; final_tokens.clear(); } - published = publish_terminal_state(generation.state, reason); + published = publish_terminal_state(generation.state, reason, generation.m); } if (!published) { return; } + if (on_engine_thread) { + record_completion_(generation.m, reason); + } generation.state->cv.notify_all(); // Handle state is visible before user code inspects it from the callback. deliver_update_(generation, final_tokens, reason); } +void RunnerImpl::record_completion_( + const GenerationMetrics& m, + FinishReason reason) { + ++metrics_.generations_completed; + switch (reason) { + case FinishReason::StopToken: + ++metrics_.finished_stop_token; + break; + case FinishReason::NewTokenLimit: + ++metrics_.finished_token_limit; + break; + case FinishReason::Cancelled: + ++metrics_.finished_cancelled; + break; + case FinishReason::Failed: + ++metrics_.finished_failed; + break; + } + metrics_.total_prompt_tokens += m.n_prompt_tokens; + metrics_.total_generated_tokens += m.n_generated_tokens; + // Zero for a generation that never reached a first token. Counted + // separately from completions so the mean divides by the samples it has, + // and so the minimum stays untouched when there are none. + const std::int64_t ttft = m.ttft_us(); + if (ttft > 0) { + ++metrics_.ttft_count; + metrics_.ttft_sum_us += ttft; + metrics_.ttft_min_us = std::min(metrics_.ttft_min_us, ttft); + metrics_.ttft_max_us = std::max(metrics_.ttft_max_us, ttft); + } +} + void RunnerImpl::complete_request_( GenerationRequest request, - FinishReason reason) { - complete_generation_(std::move(request.generation), reason); + FinishReason reason, + bool on_engine_thread) { + complete_generation_(std::move(request.generation), reason, on_engine_thread); } void RunnerImpl::complete_active_generation_( @@ -1014,7 +1233,11 @@ void RunnerImpl::complete_active_generation_( for (Task& task : scheduler_->cancel(session_id)) { (void)task; } - complete_generation_(std::move(active), reason, std::move(final_tokens)); + complete_generation_( + std::move(active), + reason, + /*on_engine_thread=*/true, + std::move(final_tokens)); } } // namespace batching diff --git a/extension/llm/batching/runner.h b/extension/llm/batching/runner.h index 535984b0955..926dedb2bae 100644 --- a/extension/llm/batching/runner.h +++ b/extension/llm/batching/runner.h @@ -39,6 +39,7 @@ #include #include +#include #include #include @@ -114,6 +115,10 @@ class GenerationHandle { // Meaningful once done(). FinishReason finish_reason() const; + // This generation's timeline and counts, complete once done(). Empty on a + // default-constructed handle. + GenerationMetrics metrics() const; + private: friend class RunnerImpl; friend class Session; @@ -206,6 +211,11 @@ class Runner { // from its callback. void shutdown(); + // What the engine measured. Read it after shutdown(): the counters are the + // engine thread's, so joining it is what makes them stable and visible. A + // call before then returns a torn snapshot. + EngineMetrics metrics() const; + private: std::shared_ptr impl_; }; diff --git a/extension/llm/batching/test/fake_executor.h b/extension/llm/batching/test/fake_executor.h index f464a16b5ea..48d57d62a19 100644 --- a/extension/llm/batching/test/fake_executor.h +++ b/extension/llm/batching/test/fake_executor.h @@ -129,9 +129,9 @@ class FakeExecutor : public Executor { // placed part way into a multi-token decode. Unset disables this. std::optional stop_token; int emit_before_stop = 0; - // Tokens a decode step produces. 1 is a plain executor; more simulates a - // speculative one answering with the run it accepted plus the model's own - // next token. Prefill always produces one whatever this is. + // Tokens an output-producing step returns. Values above 1 simulate a + // speculative executor answering with an accepted run plus the next token. + std::size_t tokens_per_prefill = 1; std::size_t tokens_per_decode = 1; // Malformed answers. An Output carries only the tokens an input produced, so // the only ways to break the contract are to produce none, or to answer for @@ -211,7 +211,8 @@ class FakeExecutor : public Executor { // Task::is_decode. Good enough for a fake: the runner only ever feeds one // token to continue. std::vector produce(const Input& input) { - const std::size_t n = input.size == 1 ? tokens_per_decode : 1; + const std::size_t n = + input.size == 1 ? tokens_per_decode : tokens_per_prefill; std::vector produced; produced.reserve(n); for (std::size_t i = 0; i < n; ++i) { diff --git a/extension/llm/batching/test/runner_test.cpp b/extension/llm/batching/test/runner_test.cpp index b9df5f5e471..92dba42afba 100644 --- a/extension/llm/batching/test/runner_test.cpp +++ b/extension/llm/batching/test/runner_test.cpp @@ -713,6 +713,27 @@ TEST(GenerationTest, StopTokenAndBudgetAreAppliedToUpdates) { EXPECT_EQ(std::find(emitted.begin(), emitted.end(), 999), emitted.end()); } +TEST(SpeculativeTest, FirstBurstRecordsIntraBurstTokenLatencies) { + FakeExecutor executor; + executor.tokens_per_prefill = 3; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(2), config(3), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + + const auto metrics = handle.metrics(); + EXPECT_EQ(metrics.n_generated_tokens, 3); + EXPECT_EQ(metrics.itl_count, 2); + EXPECT_EQ(metrics.itl_sum_us, 0); + EXPECT_EQ(metrics.itl_min_us, 0); + EXPECT_EQ(metrics.itl_max_us, 0); + EXPECT_EQ( + metrics.decode_tokens_per_sec(), std::numeric_limits::infinity()); +} + // Regression: the runner counts a step's produced tokens as far as the // executor committed them, so the next turn resumes past them rather than on // top of them. The last token of a run is not committed until it is fed back, From 0d4e6383ad61149789121c96267c2c3d899e8a98 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 28 Aug 2026 15:09:29 -0700 Subject: [PATCH 2/2] up --- extension/llm/batching/metrics.h | 402 +++++++++++++++++++++++++++++++ 1 file changed, 402 insertions(+) create mode 100644 extension/llm/batching/metrics.h diff --git a/extension/llm/batching/metrics.h b/extension/llm/batching/metrics.h new file mode 100644 index 00000000000..f1562349c7d --- /dev/null +++ b/extension/llm/batching/metrics.h @@ -0,0 +1,402 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// What a batched run measured, in two tiers. +// +// GenerationMetrics covers one generation and is published on its handle, so a +// caller reads it beside finish_reason(). EngineMetrics covers the engine and +// is owned by the runner. +// +// There is deliberately no session tier: Session::position() already reports a +// session's context, and a session's generations are recovered by grouping on +// GenerationMetrics::sid. Keeping generations separate is the point -- a second +// generation on a warm session has a different profile from a cold one, and an +// average over the two hides exactly that. +// +// Free of ExecuTorch runtime types, like the rest of these headers: no ET_LOG, +// no exceptions, and nothing here allocates outside the format_report calls. + +#include +#include +#include +#include +#include +#include + +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +// Monotonic: these are durations, and a wall-clock adjustment mid-run would +// otherwise produce negative ones. +using MetricsClock = std::chrono::steady_clock; +using MetricsTime = MetricsClock::time_point; + +ET_EXPERIMENTAL inline std::int64_t us_between( + MetricsTime from, + MetricsTime to) { + return std::chrono::duration_cast(to - from) + .count(); +} + +// A default-constructed time_point means the event never happened, which is +// distinct from it happening at time zero: the clock's epoch is the process's, +// so no real event lands there. +ET_EXPERIMENTAL inline bool stamped(MetricsTime t) { + return t.time_since_epoch().count() != 0; +} + +// One generation's timeline and counts. The finish reason is not duplicated +// here; GenerationHandle::finish_reason() already carries it, and two copies +// would be free to disagree. +struct ET_EXPERIMENTAL GenerationMetrics { + SessionId sid = 0; + + // generate_async() was called. Taken on the caller's thread, before the + // request is queued. + MetricsTime t_submit{}; + // The first batch this generation appeared in. Everything between here and + // t_submit is time the scheduler did not pick it. + MetricsTime t_first_step{}; + MetricsTime t_first_token{}; + MetricsTime t_end{}; + + std::int64_t n_prompt_tokens = 0; + std::int64_t n_generated_tokens = 0; + std::int32_t n_prefill_steps = 0; + std::int32_t n_decode_steps = 0; + + // Caller-visible inter-token latency, excluding the gap to the first token, + // which is TTFT. The first token in a later callback gets the elapsed gap; + // further tokens in that callback get zero-time samples. Summary rather than + // samples: exact mean at four scalars, where a per-token vector would cost + // 8 KB on a long generation. + std::int64_t itl_count = 0; + std::int64_t itl_sum_us = 0; + std::int64_t itl_min_us = std::numeric_limits::max(); + std::int64_t itl_max_us = 0; + + // Time to first token: what a caller waits before anything appears. + std::int64_t ttft_us() const { + return stamped(t_submit) && stamped(t_first_token) + ? us_between(t_submit, t_first_token) + : 0; + } + + // The queueing share of ttft_us(). Large means the scheduler was busy, not + // that prefill was slow. + std::int64_t queue_wait_us() const { + return stamped(t_submit) && stamped(t_first_step) + ? us_between(t_submit, t_first_step) + : 0; + } + + // The compute share of ttft_us(). Includes other sessions' work in the steps + // this one's prefill was spread across, so read it with n_prefill_steps. + std::int64_t prefill_span_us() const { + return stamped(t_first_step) && stamped(t_first_token) + ? us_between(t_first_step, t_first_token) + : 0; + } + + // Generation proper, after the first token. + std::int64_t decode_span_us() const { + return stamped(t_first_token) && stamped(t_end) + ? us_between(t_first_token, t_end) + : 0; + } + + std::int64_t e2e_us() const { + return stamped(t_submit) && stamped(t_end) ? us_between(t_submit, t_end) + : 0; + } + + double itl_mean_us() const { + return itl_count > 0 ? static_cast(itl_sum_us) / itl_count : 0.0; + } + + // What this one caller saw, which is not the engine's aggregate rate. + double decode_tokens_per_sec() const { + if (itl_count == 0) { + return 0.0; + } + return itl_sum_us > 0 + ? 1e6 * static_cast(itl_count) / itl_sum_us + : std::numeric_limits::infinity(); + } + + // Prompt tokens over the wall time this generation's prefill took. Like the + // decode rate above it is what this caller experienced, so it includes any + // work sharing those steps -- a prompt that waits behind another session's + // chunks reports a lower rate, which is what that caller actually got. + double prefill_tokens_per_sec() const { + const std::int64_t span = prefill_span_us(); + return span > 0 ? 1e6 * static_cast(n_prompt_tokens) / span : 0.0; + } +}; + +// The engine's own view, accumulated on the engine thread and read once it has +// stopped. +struct ET_EXPERIMENTAL EngineMetrics { + std::uint64_t steps = 0; + std::uint64_t steps_failed = 0; + + // Summed over steps, so dividing by `steps` gives the mean. Sequences, not + // tokens: a prefill chunk is one sequence and many tokens. + std::uint64_t decode_seqs_total = 0; + std::uint64_t prefill_seqs_total = 0; + // Generations that could have decoded in this step: alive and past their + // first token. Prefilling generations are excluded because they are not + // waiting on a decode slot, they are doing their own work. Against + // decode_seqs_total this says how much of the eligible work the scheduler + // ran, without naming any scheduler's limits. + std::uint64_t ready_total = 0; + + std::int64_t step_latency_sum_us = 0; + std::int64_t step_latency_max_us = 0; + + // Every token the engine processed. Task::is_decode classifies each one + // exactly, so these are complete. + // + // There is no decode-tokens-per-second here on purpose. A step mixing both + // kinds runs them in one forward pass over one weight read, so no share of + // its latency belongs to either. Dividing these totals by the time of the + // steps that held them would make packing prefill alongside decode -- which + // raises total throughput -- look like a decode regression. What the run + // actually delivered is model_input_tokens_per_sec() and the per-generation + // inter-token latencies; both are exact and neither moves with the mix. + std::uint64_t decode_tokens_total = 0; + std::uint64_t prefill_tokens_total = 0; + + // Steps holding at least one task of each kind. A step can hold both, so + // these overlap by mixed_steps(). + std::uint64_t steps_with_decode = 0; + std::uint64_t steps_with_prefill = 0; + + // Step latency charged once to every decode sequence in the step. A latency + // may be counted for several sequences because each of them really did wait + // it -- unlike time as a cost, waiting is not divided up. Includes mixed + // steps, where a decode stuck behind a prefill chunk waited the whole thing. + std::int64_t decode_wait_sum_us = 0; + + std::uint64_t generations_started = 0; + std::uint64_t generations_completed = 0; + std::uint64_t finished_stop_token = 0; + std::uint64_t finished_token_limit = 0; + std::uint64_t finished_cancelled = 0; + std::uint64_t finished_failed = 0; + + // Over generations that reached a first token, which is not every + // completion: one cancelled or failed during prefill has no TTFT to report. + std::uint64_t ttft_count = 0; + std::int64_t ttft_sum_us = 0; + std::int64_t ttft_min_us = std::numeric_limits::max(); + std::int64_t ttft_max_us = 0; + + std::int64_t total_prompt_tokens = 0; + std::int64_t total_generated_tokens = 0; + + MetricsTime t_first_step{}; + MetricsTime t_last_step{}; + + double wall_us() const { + return stamped(t_first_step) && stamped(t_last_step) + ? static_cast(us_between(t_first_step, t_last_step)) + : 0.0; + } + + // Mean decode sequences per step, and the mean that were eligible. Raw + // counts: normalising against a scheduler's decode limit would tie these to + // one scheduler, and against a workload smaller than that limit it would + // report idle capacity that no prompt existed to fill. + double mean_admitted_decode_seqs() const { + return steps > 0 ? static_cast(decode_seqs_total) / steps : 0.0; + } + + double mean_ready_decode_seqs() const { + return steps > 0 ? static_cast(ready_total) / steps : 0.0; + } + + // Every sequence in the step, both kinds. How full the forward pass was, + // which is the batching question; the decode figures above are the + // scheduling one. + double mean_step_seqs() const { + return steps > 0 + ? static_cast(decode_seqs_total + prefill_seqs_total) / steps + : 0.0; + } + + // Below 1 the scheduler is holding eligible work back rather than running + // out of it. Slightly under 1 is normal: a generation is briefly eligible + // but unqueued between its output being handled and its continuation being + // submitted. + double admitted_ratio() const { + return ready_total > 0 + ? static_cast(decode_seqs_total) / ready_total + : 0.0; + } + + // What the engine processed. Distinct from the generation tier's totals, + // which count what callers were given: a generation's first token comes out + // of a prefill step, and prompt tokens are processed but never emitted. + std::uint64_t model_input_tokens() const { + return decode_tokens_total + prefill_tokens_total; + } + + std::int64_t total_tokens() const { + return total_prompt_tokens + total_generated_tokens; + } + + double mean_step_tokens() const { + return steps > 0 ? static_cast(model_input_tokens()) / steps : 0.0; + } + + // Wall time the engine spent outside execute(): waiting for work, draining + // commands, running callbacks. + double idle_fraction() const { + const double wall = wall_us(); + return wall > 0.0 + ? 1.0 - static_cast(step_latency_sum_us) / wall + : 0.0; + } + + // What callers were given, per second. The comparable figure when swapping + // executors: a speculative one feeds one token per decode step and returns + // several, so processed tokens would stay flat while this rises. + double generated_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 + ? 1e6 * static_cast(total_generated_tokens) / wall + : 0.0; + } + + // What the engine put through the model, per second. Against the rate above + // it shows how much output each processed token bought. + double model_input_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 ? 1e6 * static_cast(model_input_tokens()) / wall + : 0.0; + } + + // Steps that held both kinds at once. Implied by the overlap rather than + // counted: every step holds decode, prefill, or both. + std::uint64_t mixed_steps() const { + const std::uint64_t overlap = steps_with_decode + steps_with_prefill; + return overlap > steps ? overlap - steps : 0; + } + + // Mean wall time a decode sequence waited per step it took part in. The + // engine-side counterpart to per-generation inter-token latency, and a + // cross-check on it. + double mean_decode_wait_us() const { + return decode_seqs_total > 0 + ? static_cast(decode_wait_sum_us) / decode_seqs_total + : 0.0; + } + + double mean_ttft_us() const { + return ttft_count > 0 ? static_cast(ttft_sum_us) / ttft_count : 0.0; + } + + // The sentinel is never shown: with no samples there is no minimum. + std::int64_t min_ttft_us() const { + return ttft_count > 0 ? ttft_min_us : 0; + } +}; + +namespace detail { + +inline std::string fixed(double v, int places) { + std::ostringstream os; + os << std::fixed << std::setprecision(places) << v; + return os.str(); +} + +inline std::string ms(std::int64_t microseconds) { + return fixed(static_cast(microseconds) / 1000.0, 2); +} + +} // namespace detail + +// Returned rather than printed: these headers stay free of ExecuTorch's +// logging, so the caller decides where it goes. +ET_EXPERIMENTAL inline std::string format_report(const GenerationMetrics& m) { + std::ostringstream os; + os << "session " << m.sid << ": " << m.n_prompt_tokens << " prompt + " + << m.n_generated_tokens << " generated tokens\n" + << " ttft " << detail::ms(m.ttft_us()) << " ms = queue " + << detail::ms(m.queue_wait_us()) << " + prefill " + << detail::ms(m.prefill_span_us()) << " (" + << detail::fixed(m.prefill_tokens_per_sec(), 1) << " tok/s, " + << m.n_prefill_steps + << (m.n_prefill_steps == 1 ? " step)\n" : " steps)\n"); + if (m.itl_count > 0) { + // The decode span and its rate are omitted: both follow from the mean + // below, which is decode time over decode tokens by construction. + os << " decode " << detail::ms(m.itl_sum_us / m.itl_count) + << " ms/token, max " << detail::ms(m.itl_max_us) << " -> " + << detail::fixed(m.decode_tokens_per_sec(), 1) << " tok/s over " + << m.n_decode_steps << " steps\n"; + } + os << " total " << detail::ms(m.e2e_us()) << " ms\n"; + return os.str(); +} + +ET_EXPERIMENTAL inline std::string format_report(const EngineMetrics& m) { + std::ostringstream os; + os << "engine\n" + << " elapsed " << detail::fixed(m.wall_us() / 1e6, 2) << " s, " + << m.steps << " steps"; + if (m.steps_failed > 0) { + os << " (" << m.steps_failed << " failed)"; + } + os << ", " << detail::fixed(100.0 * m.idle_fraction(), 1) << "% idle\n" + << " generations " << m.generations_completed << " of " + << m.generations_started << ": " << m.finished_stop_token + << " stop token, " << m.finished_token_limit << " token limit, " + << m.finished_cancelled << " cancelled, " << m.finished_failed + << " failed\n" + << " generated " << m.total_generated_tokens << " tokens -> " + << detail::fixed(m.generated_tokens_per_sec(), 1) << " tok/s\n" + << " model input " << m.model_input_tokens() << " tokens -> " + << detail::fixed(m.model_input_tokens_per_sec(), 1) << " tok/s (" + << m.decode_tokens_total << " decode + " << m.prefill_tokens_total + << " prompt)\n" + << " admitted " << detail::fixed(m.mean_admitted_decode_seqs(), 2) + << " of " << detail::fixed(m.mean_ready_decode_seqs(), 2) + << " ready decode seqs (" + << detail::fixed(100.0 * m.admitted_ratio(), 1) << "%)\n" + << " step size " << detail::fixed(m.mean_step_seqs(), 2) << " seqs, " + << detail::fixed(m.mean_step_tokens(), 2) << " tokens\n" + << " step time mean " + << detail::ms( + m.steps > 0 + ? m.step_latency_sum_us / static_cast(m.steps) + : 0) + << " ms, max " << detail::ms(m.step_latency_max_us) + << " ms, per-seq decode wait " + << detail::fixed(m.mean_decode_wait_us() / 1000.0, 2) << " ms\n" + << " ttft mean " << detail::ms(m.mean_ttft_us()) << " ms, min " + << detail::ms(m.min_ttft_us()) << " ms, max " + << detail::ms(m.ttft_max_us) << " ms over " << m.ttft_count + << " generations\n" + << " step kinds " << m.steps_with_decode << " decode, " + << m.steps_with_prefill << " prefill, " << m.mixed_steps() << " both\n"; + return os.str(); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch