From 3cb8bf2a05bd6f2e27c04590d5b975c7c50bc341 Mon Sep 17 00:00:00 2001 From: yenhao <46972327+yenhao-huang@users.noreply.github.com> Date: Sun, 12 Jul 2026 04:38:48 +0000 Subject: [PATCH 1/3] Record LLM model execution latency (#11771) --- extension/llm/runner/llm_runner_helper.cpp | 9 ++++-- extension/llm/runner/stats.h | 15 ++++++++++ .../runner/test/test_generation_config.cpp | 29 +++++++++++++++++++ extension/llm/runner/text_decoder_runner.cpp | 18 ++++++++++-- extension/llm/runner/text_decoder_runner.h | 6 +++- extension/llm/runner/text_llm_runner.cpp | 2 ++ 6 files changed, 74 insertions(+), 5 deletions(-) diff --git a/extension/llm/runner/llm_runner_helper.cpp b/extension/llm/runner/llm_runner_helper.cpp index 0744c09e641..90f8170d576 100644 --- a/extension/llm/runner/llm_runner_helper.cpp +++ b/extension/llm/runner/llm_runner_helper.cpp @@ -271,10 +271,16 @@ std::unique_ptr create_text_llm_runner( float init_temp = temperature == -1.0f ? 0.0f : temperature; auto sampler = std::make_unique(vocab_size, init_temp); + auto stats = std::make_unique(); + // Create text_decoder_runner ET_LOG(Info, "Using method: %s", method_name.c_str()); auto text_decoder_runner = std::make_unique( - module.get(), io_manager.get(), method_name, std::move(sampler)); + module.get(), + io_manager.get(), + method_name, + std::move(sampler), + stats.get()); // Create text_prefiller auto text_prefiller = std::make_unique( @@ -284,7 +290,6 @@ std::unique_ptr create_text_llm_runner( metadata.at(kMaxSeqLen)); // Create text_token_generator with stats - auto stats = std::make_unique(); auto text_token_generator = std::make_unique( tokenizer.get(), text_decoder_runner.get(), diff --git a/extension/llm/runner/stats.h b/extension/llm/runner/stats.h index 3c84511de3a..c6e21e5a2d2 100644 --- a/extension/llm/runner/stats.h +++ b/extension/llm/runner/stats.h @@ -66,6 +66,16 @@ struct ET_EXPERIMENTAL Stats { aggregate_sampling_timer_start_timestamp = 0; } + inline void on_model_execution_begin() { + if (model_execution_start_ms == 0) { + model_execution_start_ms = time_in_ms(); + } + } + + inline void on_model_execution_end() { + model_execution_end_ms = time_in_ms(); + } + void reset(bool all_stats = false) { // Not resetting model_load_start_ms and model_load_end_ms because reset is // typically called after warmup and before running the actual run. @@ -77,6 +87,8 @@ struct ET_EXPERIMENTAL Stats { model_load_end_ms = 0; } inference_start_ms = 0; + model_execution_start_ms = 0; + model_execution_end_ms = 0; prompt_eval_end_ms = 0; first_token_ms = 0; inference_end_ms = 0; @@ -121,6 +133,9 @@ inline std::string stats_to_json_string(const Stats& stats) { << "\"model_load_end_ms\":" << stats.model_load_end_ms << "," << "\"inference_start_ms\":" << stats.inference_start_ms << "," << "\"inference_end_ms\":" << stats.inference_end_ms << "," + << "\"model_execution_start_ms\":" + << stats.model_execution_start_ms << "," + << "\"model_execution_end_ms\":" << stats.model_execution_end_ms << "," << "\"prompt_eval_end_ms\":" << stats.prompt_eval_end_ms << "," << "\"first_token_ms\":" << stats.first_token_ms << "," << "\"aggregate_sampling_time_ms\":" << stats.aggregate_sampling_time_ms diff --git a/extension/llm/runner/test/test_generation_config.cpp b/extension/llm/runner/test/test_generation_config.cpp index dd8d026631f..f98a26af904 100644 --- a/extension/llm/runner/test/test_generation_config.cpp +++ b/extension/llm/runner/test/test_generation_config.cpp @@ -11,10 +11,39 @@ using namespace ::testing; using executorch::extension::llm::GenerationConfig; +using executorch::extension::llm::stats_to_json_string; +using executorch::extension::llm::Stats; namespace { class GenerationConfigTest : public Test {}; +TEST(StatsTest, SerializesAndResetsModelExecutionTimestamps) { + Stats stats; + stats.model_execution_start_ms = 123; + stats.model_execution_end_ms = 456; + + const std::string json = stats_to_json_string(stats); + EXPECT_NE( + json.find("\"model_execution_start_ms\":123"), std::string::npos); + EXPECT_NE(json.find("\"model_execution_end_ms\":456"), std::string::npos); + EXPECT_LT( + json.find("\"model_execution_start_ms\""), + json.find("\"model_execution_end_ms\"")); + + stats.reset(); + EXPECT_EQ(stats.model_execution_start_ms, 0); + EXPECT_EQ(stats.model_execution_end_ms, 0); +} + +TEST(StatsTest, RecordsOrderedModelExecutionTimestamps) { + Stats stats; + stats.on_model_execution_begin(); + stats.on_model_execution_end(); + + EXPECT_GT(stats.model_execution_start_ms, 0); + EXPECT_GE(stats.model_execution_end_ms, stats.model_execution_start_ms); +} + TEST_F(GenerationConfigTest, TestResolveMaxNewTokensBothDefault) { // Test when both seq_len and max_new_tokens are -1 (default) GenerationConfig config; diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index 8d234a4f306..bb64d79dd2e 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -26,11 +26,13 @@ TextDecoderRunner::TextDecoderRunner( Module* module, IOManager* io_manager, std::string method_name, - std::unique_ptr sampler) + std::unique_ptr sampler, + Stats* stats) : module_(module), io_manager_(io_manager), method_name_(std::move(method_name)), - sampler_(std::move(sampler)) {} + sampler_(std::move(sampler)), + stats_(stats) {} // This function is functional, meaning it shouldn't modify any state of the // input. It should be safe to call multiple times with the same inputs. The @@ -66,7 +68,13 @@ ::executorch::runtime::Result TextDecoderRunner::step( io_manager_->prepare_decode(tokens, start_pos_tensor, method_name_); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); + if (stats_ != nullptr) { + stats_->on_model_execution_begin(); + } auto outputs_res = module_->execute(method_name_, inputs); + if (stats_ != nullptr) { + stats_->on_model_execution_end(); + } ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); auto update_err = @@ -86,7 +94,13 @@ ::executorch::runtime::Result TextDecoderRunner::step( (void)start_pos; // unused std::vector inputs{tokens}; + if (stats_ != nullptr) { + stats_->on_model_execution_begin(); + } auto outputs_res = module_->execute(method_name_, inputs); + if (stats_ != nullptr) { + stats_->on_model_execution_end(); + } ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); ET_CHECK_MSG( outputs_res.get().size() == 1, diff --git a/extension/llm/runner/text_decoder_runner.h b/extension/llm/runner/text_decoder_runner.h index 486a9f5490f..4dbce70a830 100644 --- a/extension/llm/runner/text_decoder_runner.h +++ b/extension/llm/runner/text_decoder_runner.h @@ -18,13 +18,16 @@ namespace executorch { namespace extension { namespace llm { +struct Stats; + class ET_EXPERIMENTAL TextDecoderRunner { public: explicit TextDecoderRunner( Module* module, IOManager* io_manager, std::string method_name = "forward", - std::unique_ptr sampler = nullptr); + std::unique_ptr sampler = nullptr, + Stats* stats = nullptr); virtual ~TextDecoderRunner() = default; @@ -102,6 +105,7 @@ class ET_EXPERIMENTAL TextDecoderRunner { IOManager* io_manager_; std::string method_name_; std::unique_ptr sampler_; + Stats* stats_; }; } // namespace llm diff --git a/extension/llm/runner/text_llm_runner.cpp b/extension/llm/runner/text_llm_runner.cpp index cf7ab50b9c8..801af11f825 100644 --- a/extension/llm/runner/text_llm_runner.cpp +++ b/extension/llm/runner/text_llm_runner.cpp @@ -108,6 +108,8 @@ Error TextLLMRunner::generate( // return a response token. stats_->inference_start_ms = time_in_ms(); + stats_->model_execution_start_ms = 0; + stats_->model_execution_end_ms = 0; // Get max_seq_len for single prefill chunk limit int64_t max_seq_len = metadata_.at(kMaxSeqLen); From 6f8928d1ed29664f3f368a093c4490e9d423a1b0 Mon Sep 17 00:00:00 2001 From: yenhao <46972327+yenhao-huang@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:22:21 +0000 Subject: [PATCH 2/3] Track aggregate LLM model execution time Authored with Claude. --- extension/llm/runner/_llm_runner.pyi | 7 +++++-- extension/llm/runner/pybindings.cpp | 3 +++ extension/llm/runner/stats.h | 19 ++++++++++-------- .../runner/test/test_generation_config.cpp | 20 ++++++++++++++----- 4 files changed, 34 insertions(+), 15 deletions(-) diff --git a/extension/llm/runner/_llm_runner.pyi b/extension/llm/runner/_llm_runner.pyi index 271cf1e1540..1c81dc9a1bc 100644 --- a/extension/llm/runner/_llm_runner.pyi +++ b/extension/llm/runner/_llm_runner.pyi @@ -83,10 +83,10 @@ class Stats: """End time of tokenizer encoding in milliseconds.""" model_execution_start_ms: int - """Start time of model execution in milliseconds.""" + """Start time of the most recent model execution window in milliseconds.""" model_execution_end_ms: int - """End time of model execution in milliseconds.""" + """End time of the most recent model execution window in milliseconds.""" prompt_eval_end_ms: int """End time of prompt evaluation in milliseconds.""" @@ -100,6 +100,9 @@ class Stats: aggregate_sampling_time_ms: int """Total time spent in sampling across all tokens.""" + aggregate_model_execution_time_ms: int + """Total time spent in model execution across all forward calls.""" + num_prompt_tokens: int """Number of tokens in the input prompt.""" diff --git a/extension/llm/runner/pybindings.cpp b/extension/llm/runner/pybindings.cpp index 3188b5390c4..87ea66b1b08 100644 --- a/extension/llm/runner/pybindings.cpp +++ b/extension/llm/runner/pybindings.cpp @@ -325,6 +325,9 @@ PYBIND11_MODULE(_llm_runner, m) { .def_readonly("inference_end_ms", &Stats::inference_end_ms) .def_readonly( "aggregate_sampling_time_ms", &Stats::aggregate_sampling_time_ms) + .def_readonly( + "aggregate_model_execution_time_ms", + &Stats::aggregate_model_execution_time_ms) .def_readonly("num_prompt_tokens", &Stats::num_prompt_tokens) .def_readonly("num_generated_tokens", &Stats::num_generated_tokens) .def("on_sampling_begin", &Stats::on_sampling_begin) diff --git a/extension/llm/runner/stats.h b/extension/llm/runner/stats.h index c6e21e5a2d2..2e9ce96ea76 100644 --- a/extension/llm/runner/stats.h +++ b/extension/llm/runner/stats.h @@ -32,9 +32,9 @@ struct ET_EXPERIMENTAL Stats { long inference_start_ms = 0; // End of the tokenizer encode time. long token_encode_end_ms = 0; - // Start of the model execution (forward function) time. + /// Start timestamp of the most recent model execution (forward) window. long model_execution_start_ms = 0; - // End of the model execution (forward function) time. + /// End timestamp of the most recent model execution (forward) window. long model_execution_end_ms = 0; // prompt_eval_end_ms: Prompt array allocation and tokenization. Ends right // before the inference loop starts @@ -45,6 +45,8 @@ struct ET_EXPERIMENTAL Stats { long inference_end_ms = 0; // Keep a running total of the time spent in sampling. long aggregate_sampling_time_ms = 0; + /// Running total of time spent in model execution (forward). + long aggregate_model_execution_time_ms = 0; // Token count from prompt int64_t num_prompt_tokens = 0; // Token count from generated (total - prompt) @@ -67,13 +69,13 @@ struct ET_EXPERIMENTAL Stats { } inline void on_model_execution_begin() { - if (model_execution_start_ms == 0) { - model_execution_start_ms = time_in_ms(); - } + model_execution_start_ms = time_in_ms(); } inline void on_model_execution_end() { model_execution_end_ms = time_in_ms(); + aggregate_model_execution_time_ms += + model_execution_end_ms - model_execution_start_ms; } void reset(bool all_stats = false) { @@ -93,6 +95,7 @@ struct ET_EXPERIMENTAL Stats { first_token_ms = 0; inference_end_ms = 0; aggregate_sampling_time_ms = 0; + aggregate_model_execution_time_ms = 0; num_prompt_tokens = 0; num_generated_tokens = 0; gpu_total_bytes = static_cast(-1); @@ -133,13 +136,13 @@ inline std::string stats_to_json_string(const Stats& stats) { << "\"model_load_end_ms\":" << stats.model_load_end_ms << "," << "\"inference_start_ms\":" << stats.inference_start_ms << "," << "\"inference_end_ms\":" << stats.inference_end_ms << "," - << "\"model_execution_start_ms\":" - << stats.model_execution_start_ms << "," + << "\"model_execution_start_ms\":" << stats.model_execution_start_ms << "," << "\"model_execution_end_ms\":" << stats.model_execution_end_ms << "," << "\"prompt_eval_end_ms\":" << stats.prompt_eval_end_ms << "," << "\"first_token_ms\":" << stats.first_token_ms << "," << "\"aggregate_sampling_time_ms\":" << stats.aggregate_sampling_time_ms - << ","; + << ",\"aggregate_model_execution_time_ms\":" + << stats.aggregate_model_execution_time_ms << ","; // Only include GPU fields in the JSON if gpu_total_bytes is valid (not // equal to sentinel -1) if (stats.gpu_total_bytes != static_cast(-1)) { diff --git a/extension/llm/runner/test/test_generation_config.cpp b/extension/llm/runner/test/test_generation_config.cpp index f98a26af904..ac210106cf8 100644 --- a/extension/llm/runner/test/test_generation_config.cpp +++ b/extension/llm/runner/test/test_generation_config.cpp @@ -11,8 +11,8 @@ using namespace ::testing; using executorch::extension::llm::GenerationConfig; -using executorch::extension::llm::stats_to_json_string; using executorch::extension::llm::Stats; +using executorch::extension::llm::stats_to_json_string; namespace { class GenerationConfigTest : public Test {}; @@ -21,11 +21,14 @@ TEST(StatsTest, SerializesAndResetsModelExecutionTimestamps) { Stats stats; stats.model_execution_start_ms = 123; stats.model_execution_end_ms = 456; + stats.aggregate_model_execution_time_ms = 333; const std::string json = stats_to_json_string(stats); - EXPECT_NE( - json.find("\"model_execution_start_ms\":123"), std::string::npos); + EXPECT_NE(json.find("\"model_execution_start_ms\":123"), std::string::npos); EXPECT_NE(json.find("\"model_execution_end_ms\":456"), std::string::npos); + EXPECT_NE( + json.find("\"aggregate_model_execution_time_ms\":333"), + std::string::npos); EXPECT_LT( json.find("\"model_execution_start_ms\""), json.find("\"model_execution_end_ms\"")); @@ -33,15 +36,22 @@ TEST(StatsTest, SerializesAndResetsModelExecutionTimestamps) { stats.reset(); EXPECT_EQ(stats.model_execution_start_ms, 0); EXPECT_EQ(stats.model_execution_end_ms, 0); + EXPECT_EQ(stats.aggregate_model_execution_time_ms, 0); } -TEST(StatsTest, RecordsOrderedModelExecutionTimestamps) { +TEST(StatsTest, RecordsLatestModelExecutionAndAggregateTime) { Stats stats; + stats.model_execution_start_ms = 1; + stats.aggregate_model_execution_time_ms = 7; + stats.on_model_execution_begin(); + EXPECT_GT(stats.model_execution_start_ms, 1); stats.on_model_execution_end(); - EXPECT_GT(stats.model_execution_start_ms, 0); EXPECT_GE(stats.model_execution_end_ms, stats.model_execution_start_ms); + EXPECT_EQ( + stats.aggregate_model_execution_time_ms, + 7 + stats.model_execution_end_ms - stats.model_execution_start_ms); } TEST_F(GenerationConfigTest, TestResolveMaxNewTokensBothDefault) { From 418c2ecdfd795830e998554441a625daa3b64a8d Mon Sep 17 00:00:00 2001 From: yenhao <46972327+yenhao-huang@users.noreply.github.com> Date: Tue, 21 Jul 2026 23:42:24 +0000 Subject: [PATCH 3/3] Fix LLM Stats comment formatting --- extension/llm/runner/stats.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/extension/llm/runner/stats.h b/extension/llm/runner/stats.h index 2e9ce96ea76..9d32aaa40a3 100644 --- a/extension/llm/runner/stats.h +++ b/extension/llm/runner/stats.h @@ -32,9 +32,9 @@ struct ET_EXPERIMENTAL Stats { long inference_start_ms = 0; // End of the tokenizer encode time. long token_encode_end_ms = 0; - /// Start timestamp of the most recent model execution (forward) window. + // Start timestamp of the most recent model execution (forward) window. long model_execution_start_ms = 0; - /// End timestamp of the most recent model execution (forward) window. + // End timestamp of the most recent model execution (forward) window. long model_execution_end_ms = 0; // prompt_eval_end_ms: Prompt array allocation and tokenization. Ends right // before the inference loop starts @@ -45,7 +45,7 @@ struct ET_EXPERIMENTAL Stats { long inference_end_ms = 0; // Keep a running total of the time spent in sampling. long aggregate_sampling_time_ms = 0; - /// Running total of time spent in model execution (forward). + // Running total of time spent in model execution (forward). long aggregate_model_execution_time_ms = 0; // Token count from prompt int64_t num_prompt_tokens = 0;