46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
|
"""Aggregate decode throughput across concurrent sequences (vLLM, CPU).
|
||
|
|
|
||
|
|
Mirrors `llama-batched-bench`: B sequences each with a short prompt generate
|
||
|
|
N tokens; the figure is total generated tokens per second.
|
||
|
|
"""
|
||
|
|
import argparse, json, os, time
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||
|
|
ap.add_argument("--dtype", default="bfloat16")
|
||
|
|
ap.add_argument("--batches", default="1,2,4,8,16")
|
||
|
|
ap.add_argument("--gen", type=int, default=128)
|
||
|
|
ap.add_argument("--prompt", type=int, default=128)
|
||
|
|
ap.add_argument("--tag", default="vllm")
|
||
|
|
ap.add_argument("--out", default="batch_bench.jsonl")
|
||
|
|
a = ap.parse_args()
|
||
|
|
|
||
|
|
import vllm_omni # noqa: F401
|
||
|
|
from vllm import LLM, SamplingParams
|
||
|
|
batches = [int(b) for b in a.batches.split(",")]
|
||
|
|
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=1024,
|
||
|
|
max_num_seqs=max(batches), enable_prefix_caching=False,
|
||
|
|
trust_remote_code=False, disable_log_stats=True)
|
||
|
|
sp = SamplingParams(temperature=0.0, max_tokens=a.gen, ignore_eos=True)
|
||
|
|
res = {}
|
||
|
|
for b in batches:
|
||
|
|
prompts = [{"prompt_token_ids": list(range(10 + i * 7, 10 + i * 7 + a.prompt))}
|
||
|
|
for i in range(b)]
|
||
|
|
llm.generate(prompts[:1], SamplingParams(temperature=0.0, max_tokens=1,
|
||
|
|
ignore_eos=True), use_tqdm=False)
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
llm.generate(prompts, sp, use_tqdm=False)
|
||
|
|
dt = time.perf_counter() - t0
|
||
|
|
res[b] = round(b * a.gen / dt, 2)
|
||
|
|
print(f"batch {b:>3}: {res[b]:>8.2f} gen tok/s total", flush=True)
|
||
|
|
row = {"tag": a.tag, "threads": os.environ.get("OMP_NUM_THREADS", ""),
|
||
|
|
"gen": a.gen, "prompt": a.prompt, "tok_s_by_batch": res}
|
||
|
|
with open(a.out, "a") as f:
|
||
|
|
f.write(json.dumps(row) + "\n")
|
||
|
|
print(json.dumps(row))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|