# Copyright 2025 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import time from typing import Dict class Stats: """Python equivalent of the ExecuTorch C++ Stats structure for measuring LLM execution latency. This class provides methods to track various timestamps during model execution, including model loading, inference, token generation, and sampling times. """ def __init__(self): # Scaling factor for timestamps - in ms self.SCALING_FACTOR_UNITS_PER_SECOND = 1000 # Time stamps for different execution stages self.model_load_start_ms = 0 self.model_load_end_ms = 0 self.inference_start_ms = 0 self.token_encode_end_ms = 0 self.model_execution_start_ms = 0 self.model_execution_end_ms = 0 self.prompt_eval_end_ms = 0 self.first_token_ms = 0 self.inference_end_ms = 0 # Sampling time tracking self.aggregate_sampling_time_ms = 0 self._aggregate_sampling_timer_start_timestamp = 0 # Token counts self.num_prompt_tokens = 0 self.num_generated_tokens = 0 def on_model_load_start(self): """Mark the start of model loading.""" self.model_load_start_ms = self._time_in_ms() def on_model_load_end(self): """Mark the end of model loading.""" self.model_load_end_ms = self._time_in_ms() def on_inference_start(self): """Mark the start of inference (includes tokenizer encode time).""" self.inference_start_ms = self._time_in_ms() def on_token_encode_end(self): """Mark the end of token encoding.""" self.token_encode_end_ms = self._time_in_ms() def on_model_execution_start(self): """Mark the start of the model's forward execution.""" self.model_execution_start_ms = self._time_in_ms() def on_model_execution_end(self): """Mark the end of the model's forward execution.""" self.model_execution_end_ms = self._time_in_ms() def on_prompt_eval_end(self): """Mark the end of prompt evaluation.""" self.prompt_eval_end_ms = self._time_in_ms() def on_first_token(self): """Mark when the first generated token is emitted.""" self.first_token_ms = self._time_in_ms() def on_inference_end(self): """Mark the end of inference/generation.""" self.inference_end_ms = self._time_in_ms() def on_sampling_begin(self): """Mark the start of sampling.""" self._aggregate_sampling_timer_start_timestamp = self._time_in_ms() def on_sampling_end(self): """Mark the end of sampling and update the aggregate sampling time.""" self.aggregate_sampling_time_ms += self._time_in_ms() - self._aggregate_sampling_timer_start_timestamp self._aggregate_sampling_timer_start_timestamp = 0 def set_num_prompt_tokens(self, count: int): """Set the number of prompt tokens.""" self.num_prompt_tokens = count def set_num_generated_tokens(self, count: int): """Set the number of generated tokens.""" self.num_generated_tokens = count def reset(self, all_stats: bool = False): """Reset stats, optionally including model load times.""" if all_stats: self.model_load_start_ms = 0 self.model_load_end_ms = 0 self.inference_start_ms = 0 self.token_encode_end_ms = 0 self.model_execution_start_ms = 0 self.model_execution_end_ms = 0 self.prompt_eval_end_ms = 0 self.first_token_ms = 0 self.inference_end_ms = 0 self.aggregate_sampling_time_ms = 0 self.num_prompt_tokens = 0 self.num_generated_tokens = 0 self._aggregate_sampling_timer_start_timestamp = 0 def to_json(self) -> Dict: """Convert the stats to a JSON-serializable dictionary.""" return { "prompt_tokens": self.num_prompt_tokens, "generated_tokens": self.num_generated_tokens, "model_load_start_ms": self.model_load_start_ms, "model_load_end_ms": self.model_load_end_ms, "inference_start_ms": self.inference_start_ms, "token_encode_end_ms": self.token_encode_end_ms, "model_execution_start_ms": self.model_execution_start_ms, "model_execution_end_ms": self.model_execution_end_ms, "inference_end_ms": self.inference_end_ms, "prompt_eval_end_ms": self.prompt_eval_end_ms, "first_token_ms": self.first_token_ms, "aggregate_sampling_time_ms": self.aggregate_sampling_time_ms, "SCALING_FACTOR_UNITS_PER_SECOND": self.SCALING_FACTOR_UNITS_PER_SECOND, } def to_json_string(self) -> str: """Convert the stats to a JSON string.""" return json.dumps(self.to_json()) def print_report(self): """Print a report of the stats, similar to the C++ implementation.""" print( "\n⚠️ DISCLAIMER: Python-based perf measurements are approximate and may not " "match absolute speeds on Android/iOS apps. They are intended for relative " "comparisons—-e.g. SDPA vs. custom SDPA, FP16 vs. FP32—-so you can gauge " "performance improvements from each optimization step. For end-to-end, " "platform-accurate benchmarks, please use the official ExecuTorch apps:\n" " • iOS: https://github.com/pytorch/executorch/tree/main/extension/benchmark/apple/Benchmark\n" " • Android: https://github.com/pytorch/executorch/tree/main/extension/benchmark/android/benchmark\n" ) print(f"PyTorchObserver {self.to_json_string()}") print(f"\tPrompt Tokens: {self.num_prompt_tokens} Generated Tokens: {self.num_generated_tokens}") model_load_time = (self.model_load_end_ms - self.model_load_start_ms) / self.SCALING_FACTOR_UNITS_PER_SECOND print(f"\tModel Load Time:\t\t{model_load_time:.6f} (seconds)") inference_time_ms = self.inference_end_ms - self.inference_start_ms inference_time = inference_time_ms / self.SCALING_FACTOR_UNITS_PER_SECOND if inference_time_ms > 0 and self.num_generated_tokens > 0: inference_rate = (self.num_generated_tokens / inference_time_ms) * self.SCALING_FACTOR_UNITS_PER_SECOND else: inference_rate = 0 print( f"\tTotal inference time:\t\t{inference_time:.6f} (seconds)\t\t Rate: \t{inference_rate:.6f} (tokens/second)" ) prompt_eval_time = (self.prompt_eval_end_ms - self.inference_start_ms) / self.SCALING_FACTOR_UNITS_PER_SECOND if (self.prompt_eval_end_ms - self.inference_start_ms) > 0 and self.num_prompt_tokens > 0: prompt_eval_rate = ( self.num_prompt_tokens / (self.prompt_eval_end_ms - self.inference_start_ms) ) * self.SCALING_FACTOR_UNITS_PER_SECOND else: prompt_eval_rate = 0 print( f"\t\tPrompt evaluation:\t{prompt_eval_time:.6f} (seconds)\t\t Rate: \t{prompt_eval_rate:.6f} (tokens/second)" ) eval_time = (self.inference_end_ms - self.prompt_eval_end_ms) / self.SCALING_FACTOR_UNITS_PER_SECOND if (self.inference_end_ms - self.prompt_eval_end_ms) > 0 and self.num_generated_tokens > 0: eval_rate = ( self.num_generated_tokens / (self.inference_end_ms - self.prompt_eval_end_ms) ) * self.SCALING_FACTOR_UNITS_PER_SECOND else: eval_rate = 0 print( f"\t\tGenerated {self.num_generated_tokens} tokens:\t{eval_time:.6f} (seconds)\t\t Rate: \t{eval_rate:.6f} (tokens/second)" ) time_to_first_token = (self.first_token_ms - self.inference_start_ms) / self.SCALING_FACTOR_UNITS_PER_SECOND print(f"\tTime to first generated token:\t{time_to_first_token:.6f} (seconds)") sampling_time = self.aggregate_sampling_time_ms / self.SCALING_FACTOR_UNITS_PER_SECOND print( f"\tSampling time over {self.num_prompt_tokens + self.num_generated_tokens} tokens:\t{sampling_time:.6f} (seconds)" ) def _time_in_ms(self) -> int: """Get the current time in milliseconds.""" return int(time.time() * 1000)