Record LLM model execution latency

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."""
 

extension/llm/runner/llm_runner_helper.cpp

@@ -271,10 +271,16 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
   float init_temp = temperature == -1.0f ? 0.0f : temperature;
   auto sampler = std::make_unique<Sampler>(vocab_size, init_temp);
 
+  auto stats = std::make_unique<Stats>();
+
   // Create text_decoder_runner
   ET_LOG(Info, "Using method: %s", method_name.c_str());
   auto text_decoder_runner = std::make_unique<TextDecoderRunner>(
-      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<TextPrefiller>(
@@ -284,7 +290,6 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
       metadata.at(kMaxSeqLen));
 
   // Create text_token_generator with stats
-  auto stats = std::make_unique<Stats>();
   auto text_token_generator = std::make_unique<TextTokenGenerator>(
       tokenizer.get(),
       text_decoder_runner.get(),

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)

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<uint64_t>(-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<uint64_t>(-1)) {

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;

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> sampler)
+    std::unique_ptr<Sampler> 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<executorch::aten::Tensor> 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<executorch::aten::Tensor> TextDecoderRunner::step(
     (void)start_pos; // unused
 
     std::vector<runtime::EValue> 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,

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> sampler = nullptr);
+      std::unique_ptr<Sampler> 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> sampler_;
+  Stats* stats_;
 };
 
 } // namespace llm

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);