88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
|
"""pp512 / tg128 for Spark-X2.5 on one GPU, for vLLM or SGLang.
|
||
|
|
|
||
|
|
Same protocol as the CPU grid and llama-bench: a 512-token prefill measured on
|
||
|
|
its own, then 128 generated tokens with the one-token prefill subtracted.
|
||
|
|
Prefix caching is disabled on both engines so a replayed prompt is actually
|
||
|
|
recomputed.
|
||
|
|
"""
|
||
|
|
import argparse, json, os, time
|
||
|
|
|
||
|
|
|
||
|
|
def bench_vllm(a):
|
||
|
|
import vllm_omni # noqa: F401 registers Spark2_5ForCausalLM
|
||
|
|
from vllm import LLM, SamplingParams
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=a.max_model_len,
|
||
|
|
max_num_seqs=1, enable_prefix_caching=False,
|
||
|
|
gpu_memory_utilization=0.85, trust_remote_code=False,
|
||
|
|
disable_log_stats=True)
|
||
|
|
init = time.perf_counter() - t0
|
||
|
|
sp1 = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True)
|
||
|
|
spg = SamplingParams(temperature=0.0, max_tokens=a.tg, ignore_eos=True)
|
||
|
|
pp_prompt = {"prompt_token_ids": list(range(10, 10 + a.pp))}
|
||
|
|
tg_prompt = {"prompt_token_ids": [10]}
|
||
|
|
|
||
|
|
def run(p, sp):
|
||
|
|
s = time.perf_counter()
|
||
|
|
llm.generate([p], sp, use_tqdm=False)
|
||
|
|
return time.perf_counter() - s
|
||
|
|
|
||
|
|
run(pp_prompt, sp1)
|
||
|
|
pp = min(run(pp_prompt, sp1) for _ in range(a.reps))
|
||
|
|
tg = min(run(tg_prompt, spg) for _ in range(a.reps))
|
||
|
|
tg1 = min(run(tg_prompt, sp1) for _ in range(a.reps))
|
||
|
|
return init, pp, tg, tg1
|
||
|
|
|
||
|
|
|
||
|
|
def bench_sglang(a):
|
||
|
|
import sglang as sgl
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
llm = sgl.Engine(model_path=a.model, dtype=a.dtype,
|
||
|
|
context_length=a.max_model_len,
|
||
|
|
disable_radix_cache=True, mem_fraction_static=0.85,
|
||
|
|
disable_cuda_graph=False, log_level="error")
|
||
|
|
init = time.perf_counter() - t0
|
||
|
|
pp_ids = list(range(10, 10 + a.pp))
|
||
|
|
|
||
|
|
def run(ids, n):
|
||
|
|
s = time.perf_counter()
|
||
|
|
llm.generate(input_ids=[ids],
|
||
|
|
sampling_params={"temperature": 0.0, "max_new_tokens": n,
|
||
|
|
"ignore_eos": True})
|
||
|
|
return time.perf_counter() - s
|
||
|
|
|
||
|
|
run(pp_ids, 1)
|
||
|
|
pp = min(run(pp_ids, 1) for _ in range(a.reps))
|
||
|
|
tg = min(run([10], a.tg) for _ in range(a.reps))
|
||
|
|
tg1 = min(run([10], 1) for _ in range(a.reps))
|
||
|
|
llm.shutdown()
|
||
|
|
return init, pp, tg, tg1
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--engine", choices=["vllm", "sglang"], required=True)
|
||
|
|
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||
|
|
ap.add_argument("--dtype", default="bfloat16")
|
||
|
|
ap.add_argument("--pp", type=int, default=512)
|
||
|
|
ap.add_argument("--tg", type=int, default=128)
|
||
|
|
ap.add_argument("--reps", type=int, default=3)
|
||
|
|
ap.add_argument("--max-model-len", type=int, default=2048)
|
||
|
|
ap.add_argument("--out", default="gpu_bench.jsonl")
|
||
|
|
a = ap.parse_args()
|
||
|
|
|
||
|
|
init, pp, tg, tg1 = (bench_vllm if a.engine == "vllm" else bench_sglang)(a)
|
||
|
|
res = {
|
||
|
|
"engine": a.engine, "dtype": a.dtype, "gpu": os.environ.get("GPU_NAME", ""),
|
||
|
|
"init_s": round(init, 2),
|
||
|
|
"pp_tokens": a.pp, "pp_s": round(pp, 4), "pp_tps": round(a.pp / pp, 2),
|
||
|
|
"tg_tokens": a.tg, "tg_s": round(tg - tg1, 4),
|
||
|
|
"tg_tps": round((a.tg - 1) / max(tg - tg1, 1e-9), 2),
|
||
|
|
}
|
||
|
|
print(json.dumps(res, indent=2))
|
||
|
|
with open(a.out, "a") as f:
|
||
|
|
f.write(json.dumps(res) + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|