# SPDX-License-Identifier: GNU Affero General Public License v3.0
# Copyright 2023-present the Unsloth team. All rights reserved.

"""Batched steady-state throughput bench through ``FastLanguageModel`` +
``UNSLOTH_FAST_INFERENCE=1``. First ``generate`` call primes CUDA graphs;
subsequent calls report steady state. Compare against April CLI-only
numbers for the same workload."""

import os, sys, time
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

import argparse
import torch


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default = "unsloth/Qwen3-4B-Base")
    p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16")
    p.add_argument("--n_prompts", type = int, default = 8)
    p.add_argument("--max_new_tokens", type = int, default = 64)
    p.add_argument("--max_batch_size", type = int, default = 16)
    p.add_argument("--max_seq_length", type = int, default = 1024)
    p.add_argument("--n_rounds", type = int, default = 3)
    args = p.parse_args()

    os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
    import unsloth
    from unsloth import FastLanguageModel

    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16

    model, tok = FastLanguageModel.from_pretrained(
        model_name = args.model,
        max_seq_length = args.max_seq_length,
        dtype = dtype,
        load_in_4bit = False,
        fast_inference = True,
        max_batch_size = args.max_batch_size,
    )
    print(f"[bench] model={args.model} dtype={args.dtype}")
    prompts = [
        f"In one sentence, a fact about {t} is"
        for t in [
            "the moon",
            "gravity",
            "the ocean",
            "the sun",
            "honey",
            "rain",
            "trees",
            "mountains",
        ][: args.n_prompts]
    ]

    class _SP:
        max_tokens = args.max_new_tokens
        temperature = 0.0

    # Warmup round — captures CUDA graphs.
    print("[bench] warmup (CUDA graph capture)...")
    t0 = time.perf_counter()
    _ = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False)
    print(f"[bench] warmup wall: {time.perf_counter() - t0:.2f}s")

    walls = []
    tok_counts = []
    for r in range(args.n_rounds):
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        outs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False)
        torch.cuda.synchronize()
        dt = time.perf_counter() - t1
        n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
        walls.append(dt)
        tok_counts.append(n_tok)
        print(f"[bench] round {r}: {n_tok} toks in {dt:.2f}s -> {n_tok/dt:.1f} tok/s")

    if walls:
        wall_med = sorted(walls)[len(walls) // 2]
        tok_med = tok_counts[len(walls) // 2]
        print(
            f"[bench] median: {tok_med} toks in {wall_med:.2f}s "
            f"=> {tok_med / wall_med:.1f} tok/s"
        )


if __name__ == "__main__":
    main()