Batch generation refactoring and various fixes

mlx_lm/benchmark.py

@@ -148,10 +148,13 @@ def batch_bench():
     for i in range(args.num_trials):
         if args.delay > 0:
             time.sleep(args.delay)
+        tic = time.perf_counter()
         response = _bench()
+        toc = time.perf_counter()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]
         results = [f"{k}={v:.3f}" for k, v in results]
+        results.append(f"total_time={toc - tic:.3f}")
         rprint(f"Trial {i+1}:  " + ", ".join(results))
 
     def avg(k):

mlx_lm/examples/batch_generate_response.py

@@ -27,7 +27,7 @@
 
 # Set `verbose=True` to see generation statistics
 result = batch_generate(
-    model, tokenizer, prompts, verbose=False, return_prompt_caches=True
+    model, tokenizer, prompts, verbose=False, return_prompt_caches=True, max_tokens=2048
 )
 print(result.texts[-1])