Add inference profiling support

Author: Ziminli

infer.py

@@ -1,4 +1,6 @@
 import argparse
+import time
+import torch
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -38,13 +40,30 @@
         default="cpu",
         help='Device to use for inference (e.g., "cuda", "cpu").',
     )
+    parser.add_argument(
+        "--num-warmups",
+        type=int,
+        default=0,
+        help="For profiling. The number of warmup iterations to run before measuring performance.",
+    )
+    parser.add_argument(
+        "--num-iterations",
+        type=int,
+        default=1,
+        help="For profiling. The number of iterations to run for performance measurement.",
+    )
 
     args = parser.parse_args()
 
     model_name_or_path = args.model
     prompts = args.prompts
     max_new_tokens = args.max_new_tokens
     device = args.device
+    num_warmups = args.num_warmups
+    num_iterations = args.num_iterations
+
+    if (num_iterations < 1) or (num_warmups < 0):
+        raise ValueError("num_iterations must be >= 1 and num_warmups must be >= 0")
 
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     model = AutoModelForCausalLM.from_pretrained(model_name_or_path).to(device)
@@ -58,7 +77,22 @@
     replace_module(model, SiLU)
 
     inputs = tokenizer(prompts, padding=True, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+
+    for _ in range(num_warmups):
+        _ = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    if device == "cuda":
+        torch.cuda.synchronize()
+
+    start_time = time.time()
+
+    for _ in range(num_iterations):
+        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    if device == "cuda":
+        torch.cuda.synchronize()
+
+    end_time = time.time()
+    avg_time_ms = (end_time - start_time) * 1000 / num_iterations
     strings = tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
     print(strings)
+    print(f"\nTotal Inference time: {avg_time_ms:.4f} ms")