rebased

octoml · Dec 14, 2023 · c0c390f · c0c390f
1 parent 61cdfd4
commit c0c390f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/serve/benchmarks/benchmark_throughput.py b/serve/benchmarks/benchmark_throughput.py
@@ -119,7 +119,7 @@ def run_vllm(
     end = time.perf_counter()
     return end - start
 
-def run_mlc(engine, requests) -> float:
+def run_mlc(engine, requests, num_sequences_to_sample) -> float:
     for i, (prompt, _, output_len) in enumerate(requests):
         engine.add(
             [
@@ -128,7 +128,7 @@ def run_mlc(engine, requests) -> float:
                     messages=[ChatMessage(role="user", content=prompt)],
                     sampling_params=SamplingParams(temperature=SAMPLER_SETTING["temperature"]),
                     stopping_criteria=StoppingCriteria(max_tokens=output_len, stop_sequences=None),
-                    num_sequences=num_sequences,
+                    num_sequences=num_sequences_to_sample,
                     debug_options=DebugOptions(ignore_eos=SAMPLER_SETTING["ignore_eos"], prompt=prompt),
                 )
             ]
@@ -189,7 +189,7 @@ def main(args: argparse.Namespace):
             args.dataset, args.num_prompts, engine.tokenizer._tokenizer
         )
         elapsed_time = run_mlc(
-            engine, requests
+            engine, requests, args.num_sequences_to_sample
         )
     else:
         from transformers import AutoTokenizer