char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ruixiang Wangserver-bench : add speed-bench for speculative decoding benchmarking (#23869)689a9a470

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
16.1 KiB432 linesraw
1#!/usr/bin/env python3
2from __future__ import annotations
3
4import argparse
5import concurrent.futures
6import json
7import statistics
8import sys
9import time
10from dataclasses import asdict, dataclass
11from typing import Any
12from urllib.parse import urlparse
13
14import requests
15from datasets import get_dataset_config_names, load_dataset
16from tqdm import tqdm
17
18
19DATASET_REPO = "nvidia/SPEED-Bench"
20
21@dataclass
22class Sample:
23    id: str
24    category: str
25    turns: list[str]
26
27
28@dataclass
29class RequestResult:
30    id: str
31    category: str
32    ok: bool
33    turns: int
34    latency_s: float
35    prompt_tokens: int
36    completion_tokens: int
37    total_tokens: int
38    finish_reason: str | None
39    draft_n: int
40    draft_n_accepted: int
41    prompt_ms: float | None
42    predicted_ms: float | None
43    prompt_per_second: float | None
44    predicted_per_second: float | None
45    error: str | None
46
47
48def normalize_base_url(url: str) -> str:
49    url = url.strip().rstrip("/")
50    if not url:
51        raise ValueError("--url cannot be empty")
52    if "://" not in url:
53        url = "http://" + url
54    parsed = urlparse(url)
55    if not parsed.scheme or not parsed.netloc:
56        raise ValueError(f"invalid --url: {url}")
57    if not parsed.path.rstrip("/").endswith("/v1"):
58        url = url + "/v1"
59    return url.rstrip("/")
60
61
62def parse_extra_inputs(value: str) -> dict[str, Any]:
63    extra = json.loads(value)
64    if not isinstance(extra, dict):
65        raise ValueError("--extra-inputs must be a JSON object")
66    return extra
67
68
69def extract_turns(row: dict[str, Any]) -> list[str]:
70    turns = row.get("turns")
71    if isinstance(turns, list) and turns:
72        clean_turns = [str(turn).strip() for turn in turns if turn and str(turn).strip()]
73        if clean_turns:
74            return clean_turns
75    raise ValueError("missing or empty turns")
76
77
78def load_samples(args: argparse.Namespace) -> list[Sample]:
79    bench_names = get_dataset_config_names(DATASET_REPO)
80    if args.bench not in bench_names:
81        raise ValueError(
82            f"unknown --bench {args.bench!r}; available benches: {', '.join(bench_names)}"
83        )
84
85    dataset = load_dataset(DATASET_REPO, name=args.bench, split="test")
86    categories = list(dict.fromkeys(str(category) for category in dataset["category"]))
87    requested_categories = None
88    if args.category != "all":
89        requested_list = [category.strip() for category in args.category.split(",") if category.strip()]
90        if not requested_list:
91            raise ValueError(
92                f"--category must be 'all' or a comma-separated list; available categories: {', '.join(categories)}"
93            )
94        requested_categories = set(requested_list)
95        unknown_categories = [category for category in requested_list if category not in categories]
96        if unknown_categories:
97            unknown = ", ".join(unknown_categories)
98            raise ValueError(
99                f"unknown --category {unknown!r} for bench {args.bench!r}; "
100                f"available categories: all, {', '.join(categories)}"
101            )
102
103    samples: list[Sample] = []
104    samples_per_category: dict[str, int] = {}
105    skipped = 0
106    for index, row_raw in enumerate(dataset):
107        row = dict(row_raw)
108        category_raw = row.get("category")
109        if not isinstance(category_raw, str) or not category_raw.strip():
110            skipped += 1
111            continue
112        category = category_raw.strip()
113        if requested_categories is not None and category not in requested_categories:
114            continue
115        if args.limit is not None and samples_per_category.get(category, 0) >= args.limit:
116            continue
117
118        try:
119            turns = extract_turns(row)
120        except ValueError:
121            skipped += 1
122            continue
123        question_id = row.get("question_id")
124        if not isinstance(question_id, str) or not question_id.strip():
125            skipped += 1
126            continue
127        sample_id = question_id.strip()
128        samples.append(Sample(id=sample_id, category=category, turns=turns))
129        samples_per_category[category] = samples_per_category.get(category, 0) + 1
130
131    if not samples:
132        raise RuntimeError(f"no samples selected from bench={args.bench} category={args.category}")
133
134    if skipped:
135        print(f"speed_bench: skipped {skipped} rows without usable turns")
136    return samples
137
138
139def parse_completion_response(data: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], str | None, str]:
140    usage = data.get("usage") or {}
141    timings = data.get("timings") or {}
142    finish_reason = None
143    content = ""
144    choices = data.get("choices")
145    if isinstance(choices, list) and choices and isinstance(choices[0], dict):
146        choice = choices[0]
147        finish_reason = choice.get("finish_reason")
148        message = choice.get("message")
149        if isinstance(message, dict) and isinstance(message.get("content"), str):
150            content = message["content"]
151        elif isinstance(choice.get("text"), str):
152            content = choice["text"]
153    return usage, timings, finish_reason, content
154
155
156def run_request(
157    endpoint: str,
158    model: str | None,
159    messages: list[dict[str, str]],
160    osl: int,
161    extra_inputs: dict[str, Any],
162    timeout: float,
163) -> tuple[dict[str, Any], float]:
164    payload: dict[str, Any] = {
165        "messages": messages,
166        "max_tokens": osl,
167        "stream": False,
168    }
169    if model:
170        payload["model"] = model
171    payload.update(extra_inputs)
172    payload["max_tokens"] = osl
173
174    start = time.perf_counter()
175    response = requests.post(endpoint, json=payload, timeout=timeout)
176    latency_s = time.perf_counter() - start
177    if response.status_code != 200:
178        body = response.text[:500].replace("\n", "\\n")
179        raise RuntimeError(f"HTTP {response.status_code}: {body}")
180    return response.json(), latency_s
181
182
183def run_one(
184    sample: Sample,
185    endpoint: str,
186    model: str | None,
187    osl: int,
188    extra_inputs: dict[str, Any],
189    timeout: float,
190) -> RequestResult:
191    selected_turns = sample.turns
192    messages: list[dict[str, str]] = []
193    total_latency_s = 0.0
194    prompt_tokens = 0
195    completion_tokens = 0
196    total_tokens = 0
197    draft_n = 0
198    draft_n_accepted = 0
199    prompt_ms = 0.0
200    predicted_ms = 0.0
201    prompt_per_second = None
202    predicted_per_second = None
203    finish_reason: str | None = None
204    try:
205        for turn in selected_turns:
206            messages.append({"role": "user", "content": turn})
207            data, latency_s = run_request(endpoint, model, messages, osl, extra_inputs, timeout)
208            total_latency_s += latency_s
209            usage, timings, finish_reason, assistant_text = parse_completion_response(data)
210
211            turn_prompt_tokens = int(usage.get("prompt_tokens") or timings.get("prompt_n") or 0)
212            turn_completion_tokens_count = int(usage.get("completion_tokens") or timings.get("predicted_n") or 0)
213            turn_total_tokens_count = int(usage.get("total_tokens") or (turn_prompt_tokens + turn_completion_tokens_count))
214            prompt_tokens += turn_prompt_tokens
215            completion_tokens += turn_completion_tokens_count
216            total_tokens += turn_total_tokens_count
217            draft_n += int(timings.get("draft_n") or 0)
218            draft_n_accepted += int(timings.get("draft_n_accepted") or 0)
219            prompt_ms += float(timings.get("prompt_ms") or 0)
220            predicted_ms += float(timings.get("predicted_ms") or 0)
221            if len(selected_turns) == 1 and isinstance(timings.get("prompt_per_second"), (int, float)):
222                prompt_per_second = float(timings["prompt_per_second"])
223            if len(selected_turns) == 1 and isinstance(timings.get("predicted_per_second"), (int, float)):
224                predicted_per_second = float(timings["predicted_per_second"])
225
226            messages.append({"role": "assistant", "content": assistant_text})
227
228        if total_tokens == 0:
229            total_tokens = prompt_tokens + completion_tokens
230        if len(selected_turns) > 1:
231            prompt_per_second = (prompt_tokens / (prompt_ms / 1000)) if prompt_ms > 0 else None
232            predicted_per_second = (completion_tokens / (predicted_ms / 1000)) if predicted_ms > 0 else None
233
234        return RequestResult(
235            id=sample.id,
236            category=sample.category,
237            ok=True,
238            turns=len(selected_turns),
239            latency_s=total_latency_s,
240            prompt_tokens=prompt_tokens,
241            completion_tokens=completion_tokens,
242            total_tokens=total_tokens,
243            finish_reason=finish_reason,
244            draft_n=draft_n,
245            draft_n_accepted=draft_n_accepted,
246            prompt_ms=prompt_ms if prompt_ms > 0 else None,
247            predicted_ms=predicted_ms if predicted_ms > 0 else None,
248            prompt_per_second=prompt_per_second,
249            predicted_per_second=predicted_per_second,
250            error=None,
251        )
252    except Exception as exc:
253        return RequestResult(
254            id=sample.id,
255            category=sample.category,
256            ok=False,
257            turns=len(selected_turns),
258            latency_s=total_latency_s,
259            prompt_tokens=0,
260            completion_tokens=0,
261            total_tokens=0,
262            finish_reason=None,
263            draft_n=0,
264            draft_n_accepted=0,
265            prompt_ms=None,
266            predicted_ms=None,
267            prompt_per_second=None,
268            predicted_per_second=None,
269            error=str(exc),
270        )
271
272
273def summarize_group(category: str, results: list[RequestResult]) -> dict[str, Any]:
274    ok_results = [result for result in results if result.ok]
275    latencies = [result.latency_s for result in ok_results]
276    server_prompt_speeds = [
277        result.prompt_per_second
278        for result in ok_results
279        if result.prompt_per_second is not None
280    ]
281    server_completion_speeds = [
282        result.predicted_per_second
283        for result in ok_results
284        if result.predicted_per_second is not None
285    ]
286    turns = sum(result.turns for result in ok_results)
287    draft_n = sum(result.draft_n for result in ok_results)
288    accepted = sum(result.draft_n_accepted for result in ok_results)
289
290    return {
291        "category": category,
292        "requests": len(ok_results),
293        "turns": turns,
294        "failed": len(results) - len(ok_results),
295        "avg_prompt_t_s": statistics.mean(server_prompt_speeds) if server_prompt_speeds else None,
296        "avg_pred_t_s": statistics.mean(server_completion_speeds) if server_completion_speeds else None,
297        "avg_latency": statistics.mean(latencies) if latencies else None,
298        "draft_n": draft_n,
299        "accepted": accepted,
300        "accept_rate": (accepted / draft_n) if draft_n > 0 else None,
301    }
302
303
304def fmt_value(value: Any, kind: str = "") -> str:
305    if value is None:
306        return "n/a"
307    if kind == "int":
308        return str(int(value))
309    if kind == "rate":
310        return f"{float(value):.4f}"
311    if kind == "seconds":
312        return f"{float(value):.3f}s"
313    if kind == "speed":
314        return f"{float(value):.2f}"
315    if kind == "speedup":
316        return f"{float(value):.2f}x"
317    return str(value)
318
319
320def print_table(rows: list[dict[str, Any]]) -> None:
321    columns = [
322        ("category", "category", ""),
323        ("samples", "requests", "int"),
324        ("avg_prompt_t/s", "avg_prompt_t_s", "speed"),
325        ("avg_pred_t/s", "avg_pred_t_s", "speed"),
326        ("avg_latency", "avg_latency", "seconds"),
327        ("accept_rate", "accept_rate", "rate"),
328    ]
329    print_rows(rows, columns)
330
331
332def print_rows(rows: list[dict[str, Any]], columns: list[tuple[str, str, str]]) -> None:
333    rendered_rows = []
334    for row in rows:
335        rendered_rows.append([fmt_value(row.get(key), kind) for _, key, kind in columns])
336
337    widths = [len(header) for header, _, _ in columns]
338    for rendered in rendered_rows:
339        for i, cell in enumerate(rendered):
340            widths[i] = max(widths[i], len(cell))
341
342    header = "  ".join(header.ljust(widths[i]) for i, (header, _, _) in enumerate(columns))
343    print(header)
344    print("  ".join("-" * width for width in widths))
345    for rendered in rendered_rows:
346        print("  ".join(cell.ljust(widths[i]) for i, cell in enumerate(rendered)))
347
348
349def save_output(path: str, args: argparse.Namespace, samples: list[Sample], results: list[RequestResult], summary: list[dict[str, Any]]) -> None:
350    payload = {
351        "config": {
352            "url": args.url,
353            "model": args.model,
354            "bench": args.bench,
355            "category": args.category,
356            "osl": args.osl,
357            "concurrency": args.concurrency,
358            "extra_inputs": args.extra_inputs,
359        },
360        "selected_samples": len(samples),
361        "completed_samples": sum(1 for result in results if result.ok),
362        "failed_samples": sum(1 for result in results if not result.ok),
363        "summary": summary,
364        "results": [asdict(result) for result in results],
365    }
366    with open(path, "w", encoding="utf-8") as f:
367        json.dump(payload, f, indent=2, sort_keys=True)
368
369
370def main(argv: list[str] | None = None) -> int:
371    parser = argparse.ArgumentParser(description="Run SPEED-Bench against an OpenAI-compatible llama-server.")
372    parser.add_argument("--url", default="localhost:8080", help="Server URL, for example localhost:8080 or http://localhost:8080/v1")
373    parser.add_argument("--model", default=None, help="Optional model name to send in OpenAI requests")
374    parser.add_argument("--bench", default="qualitative", help="SPEED-Bench config to run, for example qualitative or throughput_1k")
375    parser.add_argument("--category", default="all", help="Category to run within the selected bench; use all for no category filter")
376    parser.add_argument("--osl", type=int, default=4096, help="Output sequence length, mapped to max_tokens")
377    parser.add_argument("--extra-inputs", default='{"temperature":0}', help="Extra request fields as a JSON object")
378    parser.add_argument("--concurrency", type=int, default=1, help="Concurrent client requests; usually match llama-server --np")
379    parser.add_argument("--limit", type=int, default=None, help="Optional sample limit per category for smoke tests")
380    parser.add_argument("--timeout", type=float, default=600, help="Per-request timeout in seconds")
381    parser.add_argument("--output", default=None, help="Optional path to save raw results JSON")
382    args = parser.parse_args(argv)
383    try:
384        base_url = normalize_base_url(args.url)
385        endpoint = base_url + "/chat/completions"
386        extra_inputs = parse_extra_inputs(args.extra_inputs)
387        args.extra_inputs = extra_inputs
388        samples = load_samples(args)
389    except Exception as exc:
390        print(f"speed_bench: setup failed: {exc}", file=sys.stderr)
391        return 2
392
393    print(f"speed_bench: loaded {len(samples)} samples from bench={args.bench} category={args.category}")
394
395    results: list[RequestResult] = []
396    started = time.perf_counter()
397    with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as executor:
398        futures = [
399            executor.submit(run_one, sample, endpoint, args.model, args.osl, extra_inputs, args.timeout)
400            for sample in samples
401        ]
402        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="speed_bench", unit="sample"):
403            result = future.result()
404            results.append(result)
405
406    elapsed = time.perf_counter() - started
407    categories = list(dict.fromkeys(sample.category for sample in samples))
408    summary = [
409        summarize_group(category, [result for result in results if result.category == category])
410        for category in categories
411    ]
412    summary.append(summarize_group("overall", results))
413    print()
414    print(f"Summary (elapsed={elapsed:.2f}s)")
415    print_table(summary)
416
417    if args.output:
418        save_output(args.output, args, samples, results, summary)
419        print(f"\nspeed_bench: wrote {args.output}")
420
421    failed = sum(1 for result in results if not result.ok)
422    if failed:
423        print(f"\nspeed_bench: {failed} samples failed", file=sys.stderr)
424        first_error = next((result.error for result in results if result.error), None)
425        if first_error:
426            print(f"first error: {first_error}", file=sys.stderr)
427        return 1
428    return 0
429
430
431if __name__ == "__main__":
432    raise SystemExit(main())