Add prompt processing tps in oga-bench

Signed-off-by: David Fan <[email protected]>

Author: jiafatom

src/turnkeyml/llm/cache.py

@@ -24,6 +24,7 @@ class Keys:
     MEAN_TOKENS_PER_SECOND = "mean_tokens_per_second"
     STD_DEV_TOKENS_PER_SECOND = "std_dev_tokens_per_second"
     SECONDS_TO_FIRST_TOKEN = "seconds_to_first_token"
+    PROMPT_PROCESSING_TOKENS_PER_SECOND = "prompt_processing_tokens_per_second"
     STD_DEV_SECONDS_TO_FIRST_TOKEN = "std_dev_seconds_to_first_token"
     CHECKPOINT = "checkpoint"
     DTYPE = "dtype"

src/turnkeyml/llm/tools/ort_genai/oga_bench.py

@@ -32,6 +32,7 @@ def __init__(self):
 
         self.status_stats = [
             Keys.SECONDS_TO_FIRST_TOKEN,
+            Keys.PROMPT_PROCESSING_TOKENS_PER_SECOND,
             Keys.MEAN_TOKENS_PER_SECOND,
             Keys.PROMPT_TOKENS,
         ]
@@ -144,9 +145,11 @@ def run(
                 per_iteration_tokens_per_second.append(model.tokens_per_second)
 
         mean_time_to_first_token = statistics.mean(per_iteration_time_to_first_token)
+        prompt_processing_tokens_per_second = input_ids_len / mean_time_to_first_token
         mean_tokens_per_second = statistics.mean(per_iteration_tokens_per_second)
 
         state.save_stat(Keys.SECONDS_TO_FIRST_TOKEN, mean_time_to_first_token)
+        state.save_stat(Keys.PROMPT_PROCESSING_TOKENS_PER_SECOND, prompt_processing_tokens_per_second)
         state.save_stat(Keys.MEAN_TOKENS_PER_SECOND, mean_tokens_per_second)
         state.save_stat(Keys.PROMPT_TOKENS, input_ids_len)