--- /dev/null
+#!/usr/bin/env python3
+
+import argparse
+import json
+import subprocess
+from time import sleep, time
+from typing import Optional
+
+import datasets
+import logging
+import matplotlib.pyplot as plt
+import numpy as np
+import requests
+from tqdm.contrib.concurrent import thread_map
+
+
+logging.basicConfig(level=logging.INFO, format='%(message)s')
+logger = logging.getLogger("server-bench")
+
+
+def get_prompts(n_prompts: int) -> list[str]:
+ logger.info("Loading MMLU dataset...")
+ ret = datasets.load_dataset("cais/mmlu", "all")["test"]["question"] # type: ignore
+ if n_prompts >= 0:
+ ret = ret[:n_prompts]
+ return ret
+
+
+def get_server(path_server: str, path_model: str, path_log: Optional[str], port: int, n_gpu_layers: int, parallel: int, ctx_size: int) -> dict:
+ logger.info("Starting the llama.cpp server...")
+ address = f"http://localhost:{port}"
+
+ popen_args: list[str] = [
+ path_server,
+ "--flash-attn",
+ "--n-gpu-layers", str(n_gpu_layers),
+ "--parallel", str(parallel),
+ "--ctx-size", str(parallel * ctx_size),
+ "--model", path_model,
+ "--port", str(port),
+ "--swa-full", # FIXME performance bad otherwise
+ # "--attn-streams",
+ ]
+ fout = open("bench.log", "w") if path_log is not None else subprocess.DEVNULL
+ process = subprocess.Popen(popen_args, stdout=fout, stderr=subprocess.STDOUT)
+
+ n_failures: int = 0
+ while True:
+ try:
+ sleep(1.0)
+ exit_code = process.poll()
+ if exit_code is not None:
+ raise RuntimeError(f"llama.cpp server for {path_model} exited unexpectedly with exit code {exit_code}")
+ response = requests.get(f"{address}/health")
+ if response.status_code == 200:
+ break
+ except requests.ConnectionError:
+ n_failures += 1
+ if n_failures >= 10:
+ raise RuntimeError(f"llama.cpp server for {path_model} is not healthy after 10 seconds")
+
+ return {"process": process, "address": address, "fout": fout}
+
+
+def get_prompt_length(data: dict) -> int:
+ session = data["session"]
+ server_address: str = data["server_address"]
+
+ response = session.post(
+ f"{server_address}/apply-template",
+ json={"messages": [{"role": "user", "content": data["prompt"], "stream": True}]}
+ )
+ if response.status_code != 200:
+ raise RuntimeError(f"Server returned status code {response.status_code}: {response.text}")
+ prompt: str = json.loads(response.text)["prompt"]
+ response = session.post(
+ f"{server_address}/tokenize",
+ json={"content": prompt, "add_special": True}
+ )
+ if response.status_code != 200:
+ raise RuntimeError(f"Server returned status code {response.status_code}: {response.text}")
+ tokens: list[str] = json.loads(response.text)["tokens"]
+ return len(tokens)
+
+
+def send_prompt(data: dict) -> tuple[float, list[float]]:
+ session = data["session"]
+ server_address: str = data["server_address"]
+
+ response = session.post(
+ f"{server_address}/apply-template",
+ json={"messages": [{"role": "user", "content": data["prompt"], "stream": True}]}
+ )
+ if response.status_code != 200:
+ raise RuntimeError(f"Server returned status code {response.status_code}: {response.text}")
+ prompt: str = json.loads(response.text)["prompt"]
+
+ json_data: dict = {"prompt": prompt, "seed": data["seed"], "n_predict": data["n_predict"], "stream": True}
+ response = session.post(f"{server_address}/completion", json=json_data, stream=True)
+
+ last_valid_line: str = ""
+ token_arrival_times: list[float] = []
+ for line in response.iter_lines(decode_unicode=True):
+ if not line.startswith("data: "):
+ continue
+ last_valid_line = line
+ token_arrival_times.append(time())
+ token_arrival_times = token_arrival_times[:-1]
+
+ if response.status_code != 200:
+ raise RuntimeError(f"Server returned status code {response.status_code}: {response.text}")
+ timings: dict = json.loads(last_valid_line[6:])["timings"]
+
+ return (timings["prompt_ms"], token_arrival_times)
+
+
+def benchmark(path_server: str, path_model: str, path_log: Optional[str], port: int, n_gpu_layers: int, parallel: int, ctx_size: int, n_prompts: int, n_predict: int):
+ num_workers: int = parallel + 1
+ prompts: list[str] = get_prompts(n_prompts)
+
+ server: Optional[dict] = None
+ session = None
+ try:
+ server = get_server(path_server, path_model, path_log, port, n_gpu_layers, parallel, ctx_size)
+ server_address: str = server["address"]
+
+ adapter = requests.adapters.HTTPAdapter(pool_connections=num_workers, pool_maxsize=num_workers) # type: ignore
+ session = requests.Session()
+ session.mount("http://", adapter)
+ session.mount("https://", adapter)
+
+ data: list[dict] = []
+ for i, p in enumerate(prompts):
+ data.append({"session": session, "server_address": server_address, "prompt": p, "n_predict": n_predict, "seed": i})
+
+ logger.info("Getting the prompt lengths...")
+ prompt_n = [get_prompt_length(d) for d in data]
+
+ logger.info("Starting the benchmark...\n")
+ t0 = time()
+ results: list[tuple[int, list[float]]] = thread_map(send_prompt, data, max_workers=num_workers, chunksize=1)
+ finally:
+ if server is not None:
+ server["process"].terminate()
+ server["process"].wait()
+ if session is not None:
+ session.close()
+
+ prompt_ms = []
+ token_t = []
+ depth_sum: int = 0
+ for pn, (pms, tat) in zip(prompt_n, results):
+ prompt_ms.append(pms)
+ token_t += tat
+ n_tokens: int = len(tat)
+ depth_sum += n_tokens * pn
+ depth_sum += n_tokens * (n_tokens + 1) // 2
+ prompt_n = np.array(prompt_n, dtype=np.int64)
+ prompt_ms = np.array(prompt_ms, dtype=np.float64)
+ token_t = np.array(token_t, dtype=np.float64)
+
+ token_t -= t0
+ token_t_last = np.max(token_t)
+
+ logger.info("")
+ logger.info(f"Benchmark duration: {token_t_last:.2f} s")
+ logger.info(f"Request throughput: {n_prompts / token_t_last:.2f} requests/s = {n_prompts / (token_t_last/60):.2f} requests/min")
+ logger.info(f"Total prompt length: {np.sum(prompt_n)} tokens")
+ logger.info(f"Average prompt length: {np.mean(prompt_n):.2f} tokens")
+ logger.info(f"Average prompt latency: {np.mean(prompt_ms):.2f} ms")
+ logger.info(f"Average prompt speed: {np.sum(prompt_n) / (1e-3 * np.sum(prompt_ms)):.2f} tokens/s")
+ logger.info(f"Total generated tokens: {token_t.shape[0]}")
+ logger.info(f"Average generation depth: {depth_sum / token_t.shape[0]:.2f} tokens")
+ logger.info(f"Average total generation speed: {token_t.shape[0] / token_t_last:.2f} tokens/s")
+ logger.info(f"Average generation speed per slot: {token_t.shape[0] / (parallel * token_t_last):.2f} tokens/s / slot")
+
+ plt.figure()
+ plt.scatter(prompt_n, prompt_ms, s=10.0, marker=".", alpha=0.25)
+ plt.xlim(0, 1.05 * np.max(prompt_n))
+ plt.ylim(0, 1.05 * np.max(prompt_ms))
+ plt.title(path_model)
+ plt.xlabel("Prompt length [tokens]")
+ plt.ylabel("Time to first token [ms]")
+ plt.savefig("prompt_time.png", dpi=240)
+
+ bin_max = np.ceil(token_t_last) + 1
+ plt.figure()
+ plt.hist(token_t, np.arange(0, bin_max))
+ plt.xlim(0, bin_max + 1)
+ plt.title(path_model)
+ plt.xlabel("Time [s]")
+ plt.ylabel("Num. tokens generated per second")
+ plt.savefig("gen_rate.png", dpi=240)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="Tool for benchmarking the throughput of the llama.cpp HTTP server. "
+ "Results are printed to console and visualized as plots (saved to current working directory).")
+ parser.add_argument("--path_server", type=str, default="llama-server", help="Path to the llama.cpp server binary")
+ parser.add_argument("--path_model", type=str, required=True, help="Path to the model to use for the benchmark")
+ parser.add_argument("--path_log", type=str, default=None, help="Path to the model to use for the benchmark")
+ parser.add_argument("--port", type=int, default=18725, help="Port to use for the server during the benchmark")
+ parser.add_argument("--n_gpu_layers", type=int, default=999, help="Number of GPU layers for the server")
+ parser.add_argument("--parallel", type=int, default=16, help="Number of slots for the server")
+ parser.add_argument("--ctx_size", type=int, default=4096, help="Server context size per slot")
+ parser.add_argument("--n_prompts", type=int, default=1000, help="Number of prompts to evaluate")
+ parser.add_argument("--n_predict", type=int, default=2048, help="Max. number of tokens to predict per prompt")
+ args = parser.parse_args()
+ benchmark(**vars(args))
overview / pkg / ggml / sources / llama.cpp / commitdiff