bench_context.py

#!/usr/bin/env python3
"""Measure uncached chat prefill and decode throughput across context lengths."""

from __future__ import annotations

import argparse
import json
import secrets
import time
import urllib.request


def post_json(url: str, payload: dict, timeout: int = 900) -> dict:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, separators=(",", ":")).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.load(response)


def messages_for(repetitions: int) -> list[dict]:
    content = (
        "This is an inert throughput benchmark. Read the filler, then emit "
        "the word ok repeatedly.\n"
        + " x" * repetitions
    )
    return [{"role": "user", "content": content}]


def count_tokens(base_url: str, repetitions: int) -> int:
    result = post_json(
        f"{base_url}/tokenize",
        {
            "model": "GLM-5.2",
            "messages": messages_for(repetitions),
            "chat_template_kwargs": {"enable_thinking": False},
        },
    )
    return int(result["count"])


def repetitions_for_target(base_url: str, target: int) -> tuple[int, int]:
    empty_count = count_tokens(base_url, 0)
    repetitions = max(0, target - empty_count)
    actual = count_tokens(base_url, repetitions)
    for _ in range(8):
        if actual == target:
            break
        repetitions = max(0, repetitions + target - actual)
        actual = count_tokens(base_url, repetitions)
    if actual != target:
        raise RuntimeError(
            f"could not construct {target} tokens: repetitions={repetitions}, actual={actual}"
        )
    return repetitions, actual


def stream_benchmark(
    base_url: str, target: int, completion_budget: int
) -> dict:
    repetitions, tokenized_count = repetitions_for_target(base_url, target)
    payload = {
        "model": "GLM-5.2",
        "messages": messages_for(repetitions),
        "chat_template_kwargs": {"enable_thinking": False},
        "max_completion_tokens": completion_budget,
        "min_tokens": completion_budget,
        "ignore_eos": True,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
        "return_token_ids": True,
        "cache_salt": secrets.token_urlsafe(32),
    }
    request = urllib.request.Request(
        f"{base_url}/v1/chat/completions",
        data=json.dumps(payload, separators=(",", ":")).encode(),
        headers={"Content-Type": "application/json"},
    )

    started = time.perf_counter()
    first_token_at = None
    last_token_at = None
    usage = None
    streamed_token_ids = 0
    with urllib.request.urlopen(request, timeout=1800) as response:
        for raw_line in response:
            if not raw_line.startswith(b"data: "):
                continue
            data = raw_line[6:].strip()
            if data == b"[DONE]":
                break
            chunk = json.loads(data)
            if chunk.get("usage"):
                usage = chunk["usage"]
            choices = chunk.get("choices") or []
            if not choices:
                continue
            choice = choices[0]
            delta = choice.get("delta") or {}
            token_ids = choice.get("token_ids") or delta.get("token_ids") or []
            has_token = bool(
                token_ids
                or delta.get("content")
                or delta.get("reasoning")
                or delta.get("reasoning_content")
            )
            if has_token:
                now = time.perf_counter()
                if first_token_at is None:
                    first_token_at = now
                last_token_at = now
                streamed_token_ids += len(token_ids) if token_ids else 1

    finished = time.perf_counter()
    if first_token_at is None or last_token_at is None or usage is None:
        raise RuntimeError(
            f"incomplete stream: first={first_token_at}, last={last_token_at}, usage={usage}"
        )

    prompt_tokens = int(usage["prompt_tokens"])
    completion_tokens = int(usage["completion_tokens"])
    ttft = first_token_at - started
    total = finished - started
    decode_span = max(last_token_at - first_token_at, 1e-9)
    decode_rate = (
        (completion_tokens - 1) / decode_span if completion_tokens > 1 else 0.0
    )
    return {
        "target_prompt_tokens": target,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_context_tokens": prompt_tokens + completion_tokens,
        "ttft_s": round(ttft, 4),
        "prefill_tokens_per_s": round(prompt_tokens / ttft, 2),
        "decode_tokens_per_s": round(decode_rate, 2),
        "end_to_end_s": round(total, 4),
        "end_to_end_tokens_per_s": round(
            (prompt_tokens + completion_tokens) / total, 2
        ),
        "streamed_token_events": streamed_token_ids,
        "tokenized_count_before_request": tokenized_count,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default="http://127.0.0.1:8000")
    parser.add_argument(
        "--lengths",
        default="256,1024,4096,16384,65536,131072,262144,500000",
    )
    parser.add_argument("--completion-budget", type=int, default=64)
    parser.add_argument("--output")
    args = parser.parse_args()

    lengths = [int(value) for value in args.lengths.split(",") if value]
    results = []
    for target in lengths:
        print(f"benchmarking prompt_tokens={target}", flush=True)
        result = stream_benchmark(args.base_url, target, args.completion_budget)
        results.append(result)
        print(json.dumps(result, sort_keys=True), flush=True)

    document = {
        "model": "GLM-5.2",
        "vision": True,
        "mtp_speculative_tokens": 3,
        "max_model_len": 524288,
        "results": results,
    }
    rendered = json.dumps(document, indent=2, sort_keys=True)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as handle:
            handle.write(rendered + "\n")
    print(rendered)


if __name__ == "__main__":
    main()