Function main

benchmarks/benchmark_mtp.py:92–125

(args)

def main(args):
    base_url = f"http://{args.host}:{args.port}"

    input_requests = prepare_input_requests(args.num_prompts, args.dataset_name, args.dataset_path)

    if len(args.max_concurrency) != len(args.s_itl_base_model):
        raise ValueError("--max_concurrency should be same length as --s_itl_base_model")

    for max_concurrency, s_itl in zip(args.max_concurrency, args.s_itl_base_model):
        # Warmup
        print("Starting warmup...")
        with open(os.devnull, "w") as f:
            with contextlib.redirect_stdout(f):
                send_one_batch(
                    base_url,
                    max_concurrency,
                    input_requests[0:max_concurrency],
                    True,
                )

        # Benchmark
        record = send_one_batch(base_url, max_concurrency, input_requests, False)

        metric_header = "Speed up"
        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
        for draft_token_step in args.draft_token_steps:
            speedup = calculate_speedup(
                args.acceptance_rate,
                draft_token_step,
                s_itl,
                record["mean_s_itl_ms"],
            )
            print("{:<40} {:<10.2f}".format(f"Speed up on {draft_token_step} steps draft", speedup))
        print("=" * 50)


if __name__ == "__main__":