39 lines
1.6 KiB
Python
39 lines
1.6 KiB
Python
"""Time to first token for a fresh prompt: cold, and warmed up as if it had been typed.
|
|
|
|
usage (prain installed): python3 scripts/warm_effect.py http://127.0.0.1:8000
|
|
"""
|
|
|
|
import sys
|
|
import time
|
|
import uuid
|
|
|
|
from prain.chat import chat_body, first_model, post, stream_text, warm_body
|
|
|
|
PROMPT = ("You are reviewing a design for a home server that runs a mixture-of-experts language "
|
|
"model on a single consumer graphics card with eight gigabytes of memory. The model has "
|
|
"thirty layers with one hundred and twenty-eight experts each, and only a few experts per "
|
|
"layer fit on the card at once. The rest wait in system memory or on an NVMe drive and "
|
|
"are copied over PCI Express when a token needs them. List the three biggest risks of "
|
|
"this design and, for each, one concrete way to measure it.")
|
|
|
|
|
|
def first_token_seconds(base_url, model, text):
|
|
start = time.monotonic()
|
|
body = chat_body(model, [{"role": "user", "content": text}], max_tokens=1)
|
|
with post(base_url, "/v1/chat/completions", body) as response:
|
|
next(stream_text(response), None)
|
|
return time.monotonic() - start
|
|
|
|
|
|
def main(base_url):
|
|
model = first_model(base_url)
|
|
cold = first_token_seconds(base_url, model, f"[{uuid.uuid4().hex[:8]}] {PROMPT}")
|
|
text = f"[{uuid.uuid4().hex[:8]}] {PROMPT}" # a new tag: no prefix-cache hit from above
|
|
post(base_url, "/v1/chat/completions", warm_body(model, [], text[:-40])).read()
|
|
warm = first_token_seconds(base_url, model, text)
|
|
print(f"first token: cold {cold:.2f} s, warmed {warm:.2f} s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(sys.argv[1])
|