benchmark_all_prompt_lookup.py

#!/usr/bin/env python3.12
# SPDX-License-Identifier: Apache-2.0
"""
Full regression benchmark: prompt lookup across ALL local models.

For each model, runs 3 representative prompts (high/med/low repeat)
with baseline vs prompt_lookup, and checks:
  1. Correctness: prompt_lookup output matches baseline (for trimmable caches)
  2. Performance: no regression — prompt_lookup should not be slower on average
  3. Graceful fallback: non-trimmable caches auto-disable speculation

Usage:
    python3.12 benchmark_all_prompt_lookup.py [--max-tokens 256]
"""

import argparse
import gc
import sys
import time
from dataclasses import dataclass
from pathlib import Path

MODELS = [
    # (name, path, needs_strict_false)
    (
        "Hermes-3-8B-4bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/Hermes-3-Llama-3.1-8B-4bit",
        False,
    ),
    (
        "gemma-3-12b-4bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/gemma-3-12b-it-qat-4bit",
        False,
    ),
    (
        "Devstral-24B-4bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/Devstral-Small-2-24B-Instruct-2512-4bit",
        False,
    ),
    (
        "Mistral-Small-24B-4bit",
        "/Users/raullenstudio/.lmstudio/models/lmstudio-community/Mistral-Small-3.2-24B-Instruct-2506-MLX-4bit",
        False,
    ),
    (
        "Phi-4-mini-4bit",
        "/Users/raullenstudio/.lmstudio/models/lmstudio-community/Phi-4-mini-reasoning-MLX-4bit",
        False,
    ),
    (
        "gpt-oss-20b",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/gpt-oss-20b-MXFP4-Q8",
        False,
    ),
    (
        "Qwen3.5-4B-4bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/Qwen3.5-4B-MLX-4bit",
        True,
    ),
    (
        "Qwen3.5-9B-4bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/Qwen3.5-9B-4bit",
        True,
    ),
    (
        "Qwen3.5-35B-A3B-8bit",
        "/Users/raullenstudio/.lmstudio/models/mlx-community/Qwen3.5-35B-A3B-8bit",
        True,
    ),
    (
        "Qwen3-Coder-Next-6bit",
        "/Users/raullenstudio/.lmstudio/models/lmstudio-community/Qwen3-Coder-Next-MLX-6bit",
        False,
    ),
]

PROMPTS = [
    (
        "html_list",
        "high",
        "Generate an HTML <ul> with 20 <li> items, each containing an <a href='#'> "
        "with a world capital city name. Output ONLY HTML.",
    ),
    (
        "markdown_table",
        "med",
        "Write a markdown table of 10 programming languages with columns: "
        "Name, Year, Paradigm, Typing, Speed. Output ONLY the table.",
    ),
    (
        "creative_story",
        "low",
        "Write a short creative story (about 200 words) about a robot discovering "
        "music for the first time. Be vivid and original.",
    ),
]


@dataclass
class Result:
    model: str
    prompt_name: str
    category: str
    method: str
    tokens: int = 0
    elapsed: float = 0.0
    tok_per_sec: float = 0.0
    ttft: float = 0.0
    draft_accepted: int = 0
    acceptance_pct: float = 0.0
    cache_type: str = ""
    trimmable: bool = True
    output_match: bool | None = None  # None if not checked
    error: str | None = None


def load_model(path: str, strict_false: bool):
    if strict_false:
        from mlx_lm.utils import load_model as _load_model
        from mlx_lm.utils import load_tokenizer

        model, _ = _load_model(Path(path), strict=False)
        tokenizer = load_tokenizer(Path(path), {})
        return model, tokenizer
    else:
        from mlx_lm import load as mlx_load

        return mlx_load(path)


def apply_chat_template(tokenizer, prompt: str) -> str:
    messages = [{"role": "user", "content": prompt}]
    if hasattr(tokenizer, "apply_chat_template"):
        try:
            return tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        except Exception:
            pass
    return f"User: {prompt}\nAssistant:"


def run_baseline(
    model, tokenizer, prompt_str: str, max_tokens: int
) -> tuple[list[int], float, float]:
    """Returns (token_ids, tok_per_sec, ttft)."""
    import mlx.core as mx
    from mlx_lm.generate import generate_step
    from mlx_lm.models import cache as mlx_cache

    tokens = mx.array(tokenizer.encode(prompt_str), mx.uint32)
    kv = mlx_cache.make_prompt_cache(model)

    t_start = time.perf_counter()
    t_first = None
    ids = []

    for (tid, _), _ in zip(
        generate_step(tokens, model, prompt_cache=kv, max_tokens=max_tokens),
        range(max_tokens),
    ):
        if t_first is None:
            t_first = time.perf_counter()
        t = int(tid)
        ids.append(t)
        if t == tokenizer.eos_token_id:
            break

    t_end = time.perf_counter()
    ttft = (t_first - t_start) if t_first else 0
    decode_time = t_end - (t_first or t_start)
    tps = max(len(ids) - 1, 1) / decode_time if decode_time > 0 else 0
    return ids, tps, ttft


def run_prompt_lookup(
    model, tokenizer, prompt_str: str, max_tokens: int
) -> tuple[list[int], float, float, int, float]:
    """Returns (token_ids, tok_per_sec, ttft, draft_accepted, accept_pct)."""
    import mlx.core as mx

    from vllm_mlx.speculative.prompt_lookup import prompt_lookup_generate_step

    tokens = mx.array(tokenizer.encode(prompt_str), mx.uint32)

    t_start = time.perf_counter()
    t_first = None
    ids = []
    n_draft = 0

    for tid, _, from_draft in prompt_lookup_generate_step(
        tokens,
        model,
        num_draft_tokens=4,
        ngram_size=3,
        max_tokens=max_tokens,
    ):
        if t_first is None:
            t_first = time.perf_counter()
        t = int(tid)
        ids.append(t)
        if from_draft:
            n_draft += 1
        if t == tokenizer.eos_token_id:
            break

    t_end = time.perf_counter()
    ttft = (t_first - t_start) if t_first else 0
    decode_time = t_end - (t_first or t_start)
    tps = max(len(ids) - 1, 1) / decode_time if decode_time > 0 else 0
    accept_pct = n_draft / len(ids) if ids else 0
    return ids, tps, ttft, n_draft, accept_pct


def main():
    import mlx.core as mx

    parser = argparse.ArgumentParser()
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--models", nargs="*", help="Filter by model name substring")
    args = parser.parse_args()

    models = MODELS
    if args.models:
        models = [
            (n, p, s)
            for n, p, s in MODELS
            if any(f.lower() in n.lower() for f in args.models)
        ]
        if not models:
            print(f"No models matched. Available: {[n for n, _, _ in MODELS]}")
            return

    all_results: list[Result] = []

    for model_name, model_path, strict_false in models:
        print(f"\n{'=' * 80}")
        print(f"  MODEL: {model_name}")
        print(f"{'=' * 80}")

        # Load
        try:
            model, tokenizer = load_model(model_path, strict_false)
        except Exception as e:
            print(f"  SKIP: load failed — {e}")
            for pname, cat, _ in PROMPTS:
                all_results.append(
                    Result(
                        model=model_name,
                        prompt_name=pname,
                        category=cat,
                        method="baseline",
                        error=str(e),
                    )
                )
            continue

        # Check cache
        from mlx_lm.models import cache as mlx_cache

        kv_check = mlx_cache.make_prompt_cache(model)
        cache_type = type(kv_check[0]).__name__
        trimmable = kv_check[0].is_trimmable()
        del kv_check
        print(f"  Cache: {cache_type}, trimmable={trimmable}")

        # Warmup
        warmup = apply_chat_template(tokenizer, "Hi.")
        run_baseline(model, tokenizer, warmup, max_tokens=5)
        gc.collect()
        mx.clear_cache()

        for pname, cat, raw_prompt in PROMPTS:
            prompt_str = apply_chat_template(tokenizer, raw_prompt)

            # Baseline
            print(f"  {pname} [baseline] ...", end="", flush=True)
            gc.collect()
            mx.clear_cache()
            try:
                base_ids, base_tps, base_ttft = run_baseline(
                    model, tokenizer, prompt_str, args.max_tokens
                )
                print(f" {len(base_ids)} tok, {base_tps:.1f} tok/s")
                all_results.append(
                    Result(
                        model=model_name,
                        prompt_name=pname,
                        category=cat,
                        method="baseline",
                        tokens=len(base_ids),
                        tok_per_sec=base_tps,
                        ttft=base_ttft,
                        cache_type=cache_type,
                        trimmable=trimmable,
                    )
                )
            except Exception as e:
                print(f" ERROR: {e}")
                all_results.append(
                    Result(
                        model=model_name,
                        prompt_name=pname,
                        category=cat,
                        method="baseline",
                        error=str(e),
                    )
                )
                continue

            # Prompt Lookup
            print(f"  {pname} [lookup]   ...", end="", flush=True)
            gc.collect()
            mx.clear_cache()
            try:
                pl_ids, pl_tps, pl_ttft, pl_draft, pl_accept = run_prompt_lookup(
                    model, tokenizer, prompt_str, args.max_tokens
                )
                # Check output match (only meaningful for trimmable caches)
                match = None
                if trimmable:
                    match = base_ids == pl_ids
                print(
                    f" {len(pl_ids)} tok, {pl_tps:.1f} tok/s, accept={pl_accept:.0%}"
                    f"{'' if match is None else (' ✅' if match else ' ❌ MISMATCH')}"
                )
                all_results.append(
                    Result(
                        model=model_name,
                        prompt_name=pname,
                        category=cat,
                        method="prompt_lookup",
                        tokens=len(pl_ids),
                        tok_per_sec=pl_tps,
                        ttft=pl_ttft,
                        draft_accepted=pl_draft,
                        acceptance_pct=pl_accept,
                        cache_type=cache_type,
                        trimmable=trimmable,
                        output_match=match,
                    )
                )
            except Exception as e:
                print(f" ERROR: {e}")
                all_results.append(
                    Result(
                        model=model_name,
                        prompt_name=pname,
                        category=cat,
                        method="prompt_lookup",
                        error=str(e),
                    )
                )

        del model, tokenizer
        gc.collect()
        mx.clear_cache()

    # ===================== Summary Table =====================
    print(f"\n\n{'=' * 120}")
    print(f"  REGRESSION BENCHMARK SUMMARY (max_tokens={args.max_tokens})")
    print(f"{'=' * 120}")
    print(
        f"  {'Model':<28} {'Cache':<18} {'Prompt':<16} {'Base tok/s':>10} {'PL tok/s':>10} "
        f"{'Speedup':>8} {'Accept%':>8} {'Match':>6}"
    )
    print(f"  {'-' * 116}")

    # Group by model
    model_names_seen = []
    for r in all_results:
        if r.model not in model_names_seen:
            model_names_seen.append(r.model)

    any_regression = False
    any_mismatch = False

    for mname in model_names_seen:
        model_results = [r for r in all_results if r.model == mname]
        base_results = {
            r.prompt_name: r
            for r in model_results
            if r.method == "baseline" and not r.error
        }
        pl_results = {
            r.prompt_name: r
            for r in model_results
            if r.method == "prompt_lookup" and not r.error
        }

        speedups = []
        for pname, cat, _ in PROMPTS:
            br = base_results.get(pname)
            pr = pl_results.get(pname)
            if not br or not pr:
                cache_info = model_results[0].cache_type if model_results else "?"
                trim_info = model_results[0].trimmable if model_results else "?"
                err = (
                    br.error
                    if br and br.error
                    else (pr.error if pr and pr.error else "missing")
                )
                print(
                    f"  {mname:<28} {cache_info:<18} {pname:<16} {'ERR':>10} {'ERR':>10} "
                    f"{'--':>8} {'--':>8} {'--':>6}"
                )
                continue

            sp = pr.tok_per_sec / br.tok_per_sec if br.tok_per_sec > 0 else 0
            speedups.append(sp)
            cache_str = f"{br.cache_type}({'✓' if br.trimmable else '✗'})"
            match_str = (
                "N/A"
                if pr.output_match is None
                else ("✅" if pr.output_match else "❌")
            )

            if pr.output_match is False:
                any_mismatch = True
            if sp < 0.90 and br.trimmable:  # > 10% regression on trimmable models
                any_regression = True

            print(
                f"  {mname:<28} {cache_str:<18} {pname:<16} {br.tok_per_sec:>10.1f} "
                f"{pr.tok_per_sec:>10.1f} {sp:>7.2f}x {pr.acceptance_pct:>7.0%} {match_str:>6}"
            )

        if speedups:
            avg = sum(speedups) / len(speedups)
            print(
                f"  {'':<28} {'':<18} {'>>> AVERAGE':<16} {'':>10} {'':>10} {avg:>7.2f}x"
            )
        print()

    # ===================== Verdict =====================
    print(f"{'=' * 120}")

    # Output mismatch is EXPECTED for speculative decoding — batch verification
    # produces slightly different floating point results than sequential generation.
    # This is inherent to all speculative decoding (including mlx-lm's own
    # speculative_generate_step). Both outputs are valid model outputs.
    if any_mismatch:
        print(
            "  ℹ️  INFO: Output mismatch detected (expected — batch verify vs sequential FP differences)"
        )

    if any_regression:
        print("  ⚠️  WARNING: >10% regression detected on some trimmable-cache models")
    else:
        print("  ✅ PASS: No significant performance regressions")

    # Check non-trimmable models fallback correctly
    any_fallback_fail = False
    non_trim = [
        r
        for r in all_results
        if not r.trimmable and r.method == "prompt_lookup" and not r.error
    ]
    if non_trim:
        fallback_ok = all(r.draft_accepted == 0 for r in non_trim)
        if fallback_ok:
            print(
                f"  ✅ Non-trimmable models ({len(non_trim)} runs): all correctly fell back to standard generation"
            )
        else:
            print("  ❌ FAIL: Non-trimmable models attempted speculation!")
            any_fallback_fail = True

    # Check non-trimmable fallback has no performance regression (now uses mlx-lm generate_step)
    non_trim_base = {
        r.prompt_name: r
        for r in all_results
        if not r.trimmable and r.method == "baseline" and not r.error
    }
    non_trim_pl = {
        r.prompt_name: r
        for r in all_results
        if not r.trimmable and r.method == "prompt_lookup" and not r.error
    }
    if non_trim_base and non_trim_pl:
        fallback_speedups = []
        for pname in non_trim_base:
            if pname in non_trim_pl and non_trim_base[pname].tok_per_sec > 0:
                fallback_speedups.append(
                    non_trim_pl[pname].tok_per_sec / non_trim_base[pname].tok_per_sec
                )
        if fallback_speedups:
            avg_fallback = sum(fallback_speedups) / len(fallback_speedups)
            if avg_fallback < 0.95:
                print(
                    f"  ⚠️  WARNING: Non-trimmable fallback avg {avg_fallback:.2f}x vs baseline (>5% overhead)"
                )
            else:
                print(
                    f"  ✅ Non-trimmable fallback: {avg_fallback:.2f}x vs baseline (no overhead)"
                )

    print(f"{'=' * 120}")

    return 1 if any_fallback_fail else 0


if __name__ == "__main__":
    sys.exit(main())