jzzliu
73ff04d77f
Requests of different lengths share one left-padded cache with explicit per-row RoPE positions and a key mask. Requests join mid-flight and leave on EOS/length, and unused padding columns are trimmed. Greedy output is verified identical to solo generation.
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""Throughput of the continuous-batching engine vs serving requests one at a time.
|
|
|
|
python scripts/bench_batching.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
import torch
|
|
|
|
from microgpt.batching import BatchEngine
|
|
from microgpt.model import GPT, GPTConfig
|
|
|
|
|
|
def main() -> None:
|
|
torch.manual_seed(0)
|
|
cfg = GPTConfig(vocab_size=96, block_size=256, n_layer=6, n_head=6, n_embd=192, n_kv_head=2)
|
|
model = GPT(cfg).eval()
|
|
g = torch.Generator().manual_seed(0)
|
|
reqs = []
|
|
for _ in range(16):
|
|
length = int(torch.randint(8, 48, (1,), generator=g))
|
|
reqs.append((torch.randint(0, 96, (length,), generator=g).tolist(), 64))
|
|
total = sum(n for _, n in reqs)
|
|
print(f"{len(reqs)} requests, prompts 8-47 tokens, 64 new tokens each ({total} total)\n")
|
|
print("| strategy | wall (s) | tokens/s | speedup | mean batch |")
|
|
print("|:--|--:|--:|--:|--:|")
|
|
|
|
model.generate(torch.tensor([reqs[0][0]]), 8, temperature=0) # warm up
|
|
t0 = time.perf_counter()
|
|
for prompt, n in reqs:
|
|
model.generate(torch.tensor([prompt]), n, temperature=0)
|
|
base = time.perf_counter() - t0
|
|
print(f"| one at a time | {base:.1f} | {total / base:.0f} | 1.0x | 1.0 |")
|
|
|
|
for max_batch in (2, 4, 8, 16):
|
|
eng = BatchEngine(model, max_batch=max_batch)
|
|
for prompt, n in reqs:
|
|
eng.submit(prompt, n)
|
|
t0 = time.perf_counter()
|
|
eng.run()
|
|
dt = time.perf_counter() - t0
|
|
print(
|
|
f"| continuous, max_batch={max_batch} | {dt:.1f} | {total / dt:.0f} "
|
|
f"| {base / dt:.1f}x | {eng.mean_batch_size:.1f} |"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|