Support more breakdown of latency metrics/stats for Llama (#6072)

Summary:

Support more breakdown of latency metrics/stats for Llama
- This is needed when we debugging the Frame-LLM project across teams

Reviewed By: cccclai

Differential Revision: D64139460

Author: Di Xu (SWE)

extension/llm/runner/stats.h

@@ -29,7 +29,14 @@ struct Stats {
   long model_load_end_ms;
   // inference_start_ms: Immediately after the model is loaded (or we check
   // for model load), measure the inference time.
+  // NOTE: It's actually the tokenizer encode + model execution time.
   long inference_start_ms;
+  // End of the tokenizer encode time.
+  long token_encode_end_ms;
+  // Start of the model execution (forward function) time.
+  long model_execution_start_ms;
+  // End of the model execution (forward function) time.
+  long model_execution_end_ms;
   // prompt_eval_end_ms: Prompt array allocation and tokenization. Ends right
   // before the inference loop starts
   long prompt_eval_end_ms;
@@ -139,6 +146,23 @@ inline void print_report(const Stats& stats) {
       ((double)(stats.first_token_ms - stats.inference_start_ms) /
        stats.SCALING_FACTOR_UNITS_PER_SECOND));
 
+  if (stats.token_encode_end_ms != 0) {
+    ET_LOG(
+        Info,
+        "\tTokenizer encoding time:\t%f (seconds)",
+        ((double)(stats.token_encode_end_ms - stats.inference_start_ms) /
+         stats.SCALING_FACTOR_UNITS_PER_SECOND));
+  }
+
+  if (stats.model_execution_end_ms - stats.model_execution_start_ms > 0) {
+    ET_LOG(
+        Info,
+        "\tForward/inference only time:\t%f (seconds)",
+        ((double)(stats.model_execution_end_ms -
+            stats.model_execution_start_ms) /
+            stats.SCALING_FACTOR_UNITS_PER_SECOND));
+  }
+
   ET_LOG(
       Info,
       "\tSampling time over %" PRIu64 " tokens:\t%f (seconds)",