main.py

"""
Agentic PyTorch inference optimization benchmark.

Flow:
  1. An ExperimentAgent proposes a config (LLM-driven if HF_TOKEN is set,
     otherwise a deterministic heuristic).
  2. The benchmark "tool" (benchmark.run_benchmark) runs it and measures
     latency, throughput, peak memory, and correctness.
  3. The agent records the result and repeats until the trial budget is
     spent or it runs out of candidates.
  4. The agent selects the best config under the stated constraint.
  5. A standalone kernel-level comparison (eager vs torch.compile vs a
     hand-written Triton softmax) is run and reported separately.

Usage:
    python main.py --budget 6 --model gpt2
    python main.py --budget 4 --model distilgpt2 --constraint "..."
"""
from __future__ import annotations

import argparse
import asyncio

import torch

from agent import ExperimentAgent
from benchmark import DEFAULT_MODEL_NAME, run_benchmark
from config import BenchmarkResult
from kernels import benchmark_softmax_variants

DEFAULT_CONSTRAINT = (
    "Minimize mean latency while keeping peak GPU memory under 8000 MB "
    "and correctness (top-1 token agreement with the fp32 eager reference) passing."
)


def format_results_table(results: list[BenchmarkResult]) -> str:
    headers = ["config", "latency(ms)", "±std", "throughput(tok/s)", "peak mem(MB)", "compile(s)", "top1 agree", "status"]
    rows = []
    for r in results:
        if not r.ok:
            rows.append([r.config.key(), "-", "-", "-", "-", "-", "-", f"FAILED: {r.error}"])
            continue
        rows.append(
            [
                r.config.key(),
                f"{r.mean_latency_ms:.2f}",
                f"{r.std_latency_ms:.2f}",
                f"{r.throughput_tokens_per_s:.1f}",
                f"{r.peak_memory_mb:.1f}",
                f"{r.compile_overhead_s:.2f}" if r.compile_overhead_s is not None else "-",
                f"{r.top1_agreement:.3f}" if r.top1_agreement is not None else "-",
                "ok" if r.correctness_passed else "CHECK",
            ]
        )
    widths = [max(len(h), *(len(row[i]) for row in rows)) if rows else len(h) for i, h in enumerate(headers)]
    lines = ["  ".join(h.ljust(w) for h, w in zip(headers, widths))]
    lines.append("  ".join("-" * w for w in widths))
    for row in rows:
        lines.append("  ".join(c.ljust(w) for c, w in zip(row, widths)))
    return "\n".join(lines)


def format_softmax_table(rows: list[dict]) -> str:
    if rows and rows[0].get("error") and rows[0].get("variant") == "all":
        return f"  ({rows[0]['error']})"
    headers = ["variant", "latency(ms)", "±std", "max_abs_error"]
    lines = ["  ".join(h.ljust(14) for h in headers)]
    lines.append("  ".join("-" * 14 for _ in headers))
    for r in rows:
        if "error" in r:
            lines.append(f"{r['variant'].ljust(14)}  FAILED: {r['error']}")
            continue
        lines.append(
            "  ".join(
                [
                    r["variant"].ljust(14),
                    f"{r['mean_latency_ms']:.4f}".ljust(14),
                    f"{r['std_latency_ms']:.4f}".ljust(14),
                    f"{r['max_abs_error']:.2e}",
                ]
            )
        )
    return "\n".join(lines)


async def main():
    parser = argparse.ArgumentParser(description="Agentic PyTorch inference optimization benchmark")
    parser.add_argument("--model", default=DEFAULT_MODEL_NAME, help="HF causal LM name, e.g. gpt2 / distilgpt2")
    parser.add_argument("--budget", type=int, default=6, help="number of trials the agent may run")
    parser.add_argument("--constraint", default=DEFAULT_CONSTRAINT)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iters", type=int, default=20)
    parser.add_argument("--no-llm", action="store_true", help="force the deterministic heuristic agent")
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        print(
            "WARNING: no CUDA GPU detected. Running on CPU for structural testing only -- "
            "latency / throughput / memory / Triton results will NOT be representative. "
            "Run this on an NVIDIA GPU machine for the real numbers.\n"
        )

    agent = ExperimentAgent(constraint=args.constraint, trial_budget=args.budget, use_llm=not args.no_llm)
    print(f"Agent mode: {'LLM-driven' if agent.use_llm else 'heuristic (no HF_TOKEN found)'}")
    print(f"Model: {args.model}  |  device: {device}  |  budget: {args.budget} trials\n")

    for trial in range(1, args.budget + 1):
        config = await agent.propose_next()
        if config is None:
            print("No more candidate configs to try.")
            break

        print(f"[trial {trial}/{args.budget}] {config.key()}")
        result = run_benchmark(
            config, model_name=args.model, device=device, n_warmup=args.warmup, n_iters=args.iters
        )
        agent.record(result)

        if not result.ok:
            print(f"  -> FAILED: {result.error}\n")
            continue
        print(
            f"  -> latency={result.mean_latency_ms:.2f}ms (±{result.std_latency_ms:.2f})  "
            f"throughput={result.throughput_tokens_per_s:.1f} tok/s  "
            f"peak_mem={result.peak_memory_mb:.1f}MB  "
            f"top1_agree={result.top1_agreement:.3f}\n"
        )

    print("=== All trials ===")
    print(format_results_table(agent.results))

    best = agent.select_best()
    print("\n=== Agent's pick ===")
    print(f"Constraint: {args.constraint}")
    if best:
        print(
            f"Best: {best.config.key()} -> {best.mean_latency_ms:.2f}ms, "
            f"{best.throughput_tokens_per_s:.1f} tok/s, {best.peak_memory_mb:.1f}MB, "
            f"top1_agree={best.top1_agreement:.3f}"
        )
    else:
        print("No successful trial to recommend.")

    print("\n=== Kernel-level check: softmax (eager vs torch.compile vs hand-written Triton) ===")
    softmax_rows = benchmark_softmax_variants(device=device)
    print(format_softmax_table(softmax_rows))


if __name__ == "__main__":
    asyncio.run(main())