LLM benchmark tool with provider support and accurate tokenization

llm_bench.py

import time
import openai
import argparse
import uuid
import tiktoken  

# ---------------------------------------------------------------------------
# Configurazione provider
# ---------------------------------------------------------------------------
PROVIDERS = {
    "omlx": {
        "base_url": "http://127.0.0.1:8000/v1",
        "api_key": "",
        "model": "gemma-4-26b-a4b-it-4bit"
    },
    "llama.cpp": {
        "base_url": "http://127.0.0.1:8000/v1",
        "api_key": "",
        "model": "gemma-4-26b-a4b-it-4bit"
    },
    "ollama": {
        "base_url": "http://127.0.0.1:11434/v1",
        "api_key": "",
        "model": "gemma4:26b"
    },
    "holodeck": {
        "base_url": "http://192.168.1.25:8001/v1",
        "api_key": "",
        "model": "gemma-4-26b-a4b-it-4bit"
    },
}

# ---------------------------------------------------------------------------
# Parametri benchmark
# ---------------------------------------------------------------------------
CONTEXT_TARGET    = 1024   
GENERATION_TARGET = 1024   
MIN_TTFT_MS       = 30     

# Inizializzazione tokenizer
try:
    encoding = tiktoken.get_encoding("cl100k_base")
except Exception:
    encoding = tiktoken.get_encoding("gpt2")

def get_large_prompt(target_tokens):
    """Genera un blocco di testo per raggiungere circa target_tokens."""
    unique_id = str(uuid.uuid4())
    # Usiamo una stringa di base per costruire il prompt
    base_filler = "The quick brown fox jumps over the lazy dog. "
    # Random for cache invalidation :)
    current_text = f"Random ID: {unique_id}. "
    
    # Aggiungiamo testo finché non raggiungiamo il target di token
    while len(encoding.encode(current_text)) < target_tokens:
        current_text += base_filler
        
    instructions = (
        f"\n\n[Padding Context Above]\n\n"
        "Repeat the following sequence exactly and continuously: alpha bravo charlie delta echo foxtrot. "
        "Do not stop until you reach the token limit. No intro/outro."
    )
    return current_text + instructions

def run_single_test(name, config, is_warmup=False):
    """Esegue una singola iterazione e restituisce le metriche."""
    client = openai.OpenAI(base_url=config["base_url"], api_key=config["api_key"])
    prompt = get_large_prompt(CONTEXT_TARGET)
    
    # Conteggio esatto dei token del prompt tramite tiktoken
    prompt_tokens = len(encoding.encode(prompt))

    start_time = time.time()
    ttft = None
    full_text = ""
    tokens_from_usage = 0

    try:
        stream = client.chat.completions.create(
            model=config["model"],
            messages=[{"role": "user", "content": prompt}],
            max_tokens=GENERATION_TARGET,
            stream=True,
            stream_options={"include_usage": True}
        )

        for chunk in stream:
            if not chunk.choices:
                if hasattr(chunk, 'usage') and chunk.usage:
                    tokens_from_usage = chunk.usage.completion_tokens
                continue

            delta = chunk.choices[0].delta
            content = getattr(delta, 'content', None)
            
            if ttft is None and content:
                ttft = time.time() - start_time

            if content:
                full_text += content

        end_time = time.time()
        total_duration = end_time - start_time

        if is_warmup: return None

        if ttft is None:
            ttft = total_duration 

        # Conteggio Token accurato
        if tokens_from_usage > 0:
            tokens_generated = tokens_from_usage
        else:
            # Fallback accurato usando tiktoken sulla stringa ricevuta
            tokens_generated = len(encoding.encode(full_text))

        # CALCOLO PP (Prompt Processing) - Usa i token reali contati
        pp_tok_s = None
        if ttft >= (MIN_TTFT_MS / 1000):
            pp_tok_s = prompt_tokens / ttft

        # CALCOLO TG (Token Generation)
        gen_duration = total_duration - ttft
        if gen_duration > 0.01:
            tg_tok_s = (tokens_generated - 1) / gen_duration
        else:
            tg_tok_s = tokens_generated / total_duration

        return {
            "pp_tok_s": pp_tok_s,
            "tg_tok_s": tg_tok_s,
            "ttft_ms": ttft * 1000,
            "tokens_generated": tokens_generated,
            "prompt_tokens": prompt_tokens,
            "total_s": total_duration,
        }

    except Exception as e:
        if not is_warmup: print(f"  [Errore {name}]: {e}")
        return None

def run_benchmark_suite(name, config, runs=3):
    print(f"\n--- Benchmark: {name} ---")
    print("  Warmup run (pre-loading VRAM)...", end="", flush=True)
    run_single_test(name, config, is_warmup=True)
    print(" DONE")

    all_results = []
    for i in range(1, runs + 1):
        print(f"  Run {i}/{runs}...", end="", flush=True)
        res = run_single_test(name, config)
        if res:
            all_results.append(res)
            pp_display = f"{res['pp_tok_s']:.1f}" if res['pp_tok_s'] else "N/A"
            print(f" OK | TTFT: {res['ttft_ms']:.0f}ms | PP: {pp_display} | TG: {res['tg_tok_s']:.1f}")
        else:
            print(" FAILED")

    if not all_results: return None

    return {
        "pp_tok_s": sum(r["pp_tok_s"] for r in all_results if r["pp_tok_s"]) / len([r for r in all_results if r["pp_tok_s"]]) if any(r["pp_tok_s"] for r in all_results) else None,
        "tg_tok_s": sum(r["tg_tok_s"] for r in all_results) / len(all_results),
        "ttft_ms": sum(r["ttft_ms"] for r in all_results) / len(all_results),
        "avg_tokens": sum(r["tokens_generated"] for r in all_results) / len(all_results),
        "avg_prompt_tokens": sum(r["prompt_tokens"] for r in all_results) / len(all_results)
    }

def main():
    parser = argparse.ArgumentParser(description="LLM Benchmark Context 1024 (Accurate Tokenization)")
    parser.add_argument("--provider", choices=list(PROVIDERS.keys()))
    parser.add_argument("--runs", type=int, default=3)
    args = parser.parse_args()

    target_providers = {args.provider: PROVIDERS[args.provider]} if args.provider else PROVIDERS
    results = {}

    for name, config in target_providers.items():
        res = run_benchmark_suite(name, config, runs=args.runs)
        if res: results[name] = res

    print("\n" + "="*95)
    print(f"{'PROVIDER':<15} | {'TTFT (ms)':<12} | {'PP (tok/s)':<15} | {'TG (tok/s)':<15} | {'TOK GEN'}")
    print("-" * 95)
    for name, m in results.items():
        pp_val = f"{m['pp_tok_s']:.1f}" if m['pp_tok_s'] else "N/A"
        print(f"{name.upper():<15} | {m['ttft_ms']:<12.1f} | {pp_val:<15} | {m['tg_tok_s']:<15.2f} | {m['avg_tokens']:.0f}")
    print("="*95)

if __name__ == "__main__":
    main()