| """Decode speed of the Ling-mini-2.0 megakernel. |
| |
| python bench.py one chat: 10 launches of 256 greedy tokens after a 32-token prompt, median |
| python bench.py --chats 3 three chats at once, each on its share of the SMs: tokens/s each and in total, |
| and whether each chat wrote the same 128 tokens as alone |
| |
| The GPU is warmed up first (Windows lowers an idle GPU's clocks): launches run until they stop |
| getting faster for two seconds. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import statistics |
| import sys |
| import time |
|
|
| import torch |
|
|
| from ling_kernel import HERE, LingKernel, State |
|
|
| PROMPTS = [ |
| "The history of the printing press begins in the fifteenth century, when Johannes Gutenberg combined " |
| "movable metal type, oil-based ink and a wooden press adapted from those used to make wine.", |
| "def merge_sort(items):\n \"\"\"Return a new list with the items in ascending order.\"\"\"\n" |
| " if len(items) <= 1:\n return list(items)\n middle = len(items) // 2\n", |
| ] |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--chats", type=int, default=1) |
| args = ap.parse_args() |
| sys.stdout.reconfigure(encoding="utf-8") |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(HERE) |
| mk = LingKernel(max_pos=1024, log=lambda s: print(s, flush=True)) |
| prompts = [tok(p)["input_ids"][:32] for p in PROMPTS] |
| firsts = [mk.forced(p)[-1] for p in prompts] |
| best, since, start = float("inf"), time.time(), time.time() |
| while time.time() - start < 30: |
| s = mk.timed_greedy(firsts[0], len(prompts[0]), 64) |
| if s < 0.97 * best: |
| best, since = s, time.time() |
| elif time.time() - since > 2: |
| break |
| if args.chats <= 1: |
| mk.forced(prompts[0]) |
| times = [mk.timed_greedy(firsts[0], len(prompts[0]), 256) for _ in range(10)] |
| per = statistics.median(times) / 256 |
| print(f"one chat: {1 / per:.0f} tokens/s (median of 10 launches of 256 tokens; best {256 / min(times):.0f})") |
| return |
| n = args.chats |
| nb = max(16, mk.sms // n) |
| states = [State(1024) for _ in range(n)] |
| streams = [torch.cuda.Stream() for _ in range(n)] |
| pr = [prompts[i % 2] for i in range(n)] |
| for i in range(n): |
| mk.launch_async(states[i], pr[i], 0, len(pr[i]), True, nb, streams[i]) |
| torch.cuda.synchronize() |
| first = [states[i].out[len(pr[i]) - 1].item() for i in range(n)] |
| alone = [] |
| for i in range(2): |
| mk.forced(prompts[i]) |
| alone.append(mk.greedy(firsts[i], len(prompts[i]), 128)) |
|
|
| def round_(k: int) -> float: |
| torch.cuda.synchronize() |
| t = time.perf_counter() |
| for i in range(n): |
| mk.launch_async(states[i], [first[i]], len(pr[i]), k, False, nb, streams[i]) |
| torch.cuda.synchronize() |
| return time.perf_counter() - t |
|
|
| round_(128) |
| same = [states[i].out[:128].tolist() == alone[i % 2] for i in range(n)] |
| t = statistics.median(round_(256) for _ in range(6)) |
| print(f"{n} chats at once, {nb} SMs each: {256 / t:.0f} tokens/s each, {n * 256 / t:.0f} in total; " |
| f"each wrote the same 128 tokens as alone: {same}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
