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/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/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 3c84511de3a..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 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) @@ -66,6 +68,16 @@ struct ET_EXPERIMENTAL Stats { aggregate_sampling_timer_start_timestamp = 0; } + inline void on_model_execution_begin() { + 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) { // 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,10 +89,13 @@ 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; 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); @@ -121,10 +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_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 dd8d026631f..ac210106cf8 100644 --- a/extension/llm/runner/test/test_generation_config.cpp +++ b/extension/llm/runner/test/test_generation_config.cpp @@ -11,10 +11,49 @@ using namespace ::testing; using executorch::extension::llm::GenerationConfig; +using executorch::extension::llm::Stats; +using executorch::extension::llm::stats_to_json_string; namespace { class GenerationConfigTest : public Test {}; +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_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\"")); + + 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, 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_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) { // 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);