bench-diffusion-gemma-cli.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import statistics
import subprocess
import sys
import time
from pathlib import Path


PREFILL_RE = re.compile(
    r"prefill \(encoder, no self-cond\):\s+"
    r"(?P<tokens>\d+)\s+tokens\s+in\s+"
    r"(?P<seconds>[0-9.]+)\s+s\s+\((?P<tps>[0-9.]+)\s+tok/s\)"
)

GEN_RE = re.compile(
    r"generation:\s+"
    r"(?P<blocks>\d+)\s+block\(s\),\s+"
    r"(?P<steps>\d+)\s+denoising steps,\s+"
    r"(?P<canvas_tokens>\d+)\s+canvas tokens in\s+"
    r"(?P<seconds>[0-9.]+)\s+s\s+"
    r"\((?P<canvas_tps>[0-9.]+)\s+canvas tok/s,\s+"
    r"(?P<seconds_per_step>[0-9.]+)\s+s/step\);\s+"
    r"answer tokens=(?P<answer_tokens>\d+)"
)


def clean_env() -> dict[str, str]:
    env = dict(os.environ)
    if os.name == "nt":
        path_value = env.get("Path") or env.get("PATH") or os.defpath
        for key in list(env):
            if key.lower() == "path":
                del env[key]
        env["Path"] = path_value
    return env


def quote_cmd(cmd: list[str]) -> str:
    return " ".join(shlex.quote(str(x)) for x in cmd)


def tail(text: str, n_lines: int = 80) -> str:
    lines = text.splitlines()
    return "\n".join(lines[-n_lines:])


def parse_output(text: str) -> dict[str, object]:
    result: dict[str, object] = {}

    prefill_match = None
    for prefill_match in PREFILL_RE.finditer(text):
        pass
    if prefill_match:
        result["prefill"] = {
            "tokens": int(prefill_match.group("tokens")),
            "seconds": float(prefill_match.group("seconds")),
            "tokens_per_second": float(prefill_match.group("tps")),
        }

    gen_match = None
    for gen_match in GEN_RE.finditer(text):
        pass
    if gen_match:
        result["generation"] = {
            "blocks": int(gen_match.group("blocks")),
            "steps": int(gen_match.group("steps")),
            "canvas_tokens": int(gen_match.group("canvas_tokens")),
            "seconds": float(gen_match.group("seconds")),
            "canvas_tokens_per_second": float(gen_match.group("canvas_tps")),
            "seconds_per_step": float(gen_match.group("seconds_per_step")),
            "answer_tokens": int(gen_match.group("answer_tokens")),
        }

    return result


def mean_or_none(values: list[float]) -> float | None:
    return statistics.fmean(values) if values else None


def load_prompts(prompt: str, prompt_file: Path | None) -> list[str]:
    if prompt_file is None:
        return [prompt]
    text = prompt_file.read_text(encoding="utf-8")
    prompts = [line.strip() for line in text.splitlines() if line.strip()]
    if not prompts:
        raise ValueError(f"prompt file has no non-empty prompts: {prompt_file}")
    return prompts


def fmt_num(value: object, digits: int = 3) -> str:
    if value is None:
        return "n/a"
    return f"{float(value):.{digits}f}"


def prompt_preview(prompt: str, limit: int = 90) -> str:
    text = prompt.replace("|", "\\|").replace("\n", " ").strip()
    if len(text) <= limit:
        return text
    return text[: limit - 3] + "..."


def summarize(runs: list[dict[str, object]]) -> dict[str, object]:
    wall_ms = [float(r["wall_ms"]) for r in runs]
    gen_seconds = []
    canvas_tps = []
    sec_per_step = []
    prefill_tps = []
    denoise_steps = []
    canvas_tokens = []
    blocks = []
    answer_tokens = []

    for run in runs:
        generation = run.get("generation")
        if isinstance(generation, dict):
            gen_seconds.append(float(generation["seconds"]))
            canvas_tps.append(float(generation["canvas_tokens_per_second"]))
            sec_per_step.append(float(generation["seconds_per_step"]))
            denoise_steps.append(float(generation["steps"]))
            canvas_tokens.append(float(generation["canvas_tokens"]))
            blocks.append(float(generation["blocks"]))
            answer_tokens.append(float(generation["answer_tokens"]))
        prefill = run.get("prefill")
        if isinstance(prefill, dict):
            prefill_tps.append(float(prefill["tokens_per_second"]))

    seconds_per_step_avg = mean_or_none(sec_per_step)
    denoising_steps_total = sum(denoise_steps)
    blocks_total = sum(blocks)
    return {
        "runs": len(runs),
        "wall_ms_avg": mean_or_none(wall_ms),
        "wall_ms_min": min(wall_ms) if wall_ms else None,
        "wall_ms_max": max(wall_ms) if wall_ms else None,
        "generation_seconds_avg": mean_or_none(gen_seconds),
        "canvas_tokens_per_second_avg": mean_or_none(canvas_tps),
        "seconds_per_step_avg": seconds_per_step_avg,
        "ms_per_step_avg": seconds_per_step_avg * 1000.0 if seconds_per_step_avg is not None else None,
        "denoising_steps_total": denoising_steps_total,
        "denoising_steps_avg": denoising_steps_total / blocks_total if blocks_total > 0 else None,
        "canvas_tokens_avg": mean_or_none(canvas_tokens),
        "canvas_tokens_total": sum(canvas_tokens),
        "blocks_total": blocks_total,
        "blocks_avg": mean_or_none(blocks),
        "answer_tokens_avg": mean_or_none(answer_tokens),
        "prefill_tokens_per_second_avg": mean_or_none(prefill_tps),
    }


def prompt_summaries(runs: list[dict[str, object]]) -> list[dict[str, object]]:
    grouped: dict[int, list[dict[str, object]]] = {}
    prompts: dict[int, str] = {}
    for run in runs:
        prompt_index = int(run["prompt_index"])
        grouped.setdefault(prompt_index, []).append(run)
        prompts[prompt_index] = str(run["prompt"])

    rows = []
    for prompt_index in sorted(grouped):
        rows.append({
            "prompt_index": prompt_index,
            "prompt": prompts[prompt_index],
            "summary": summarize(grouped[prompt_index]),
        })
    return rows


def build_command(args: argparse.Namespace, prompt: str) -> list[str]:
    cmd = [
        str(args.binary),
        "-m", str(args.model),
        "-p", prompt,
        "-n", str(args.n_predict),
        "-c", str(args.ctx_size),
        "-ngl", str(args.n_gpu_layers),
        "-t", str(args.threads),
        "--diffusion-steps", str(args.diffusion_steps),
        "-lv", str(args.log_verbosity),
    ]

    if args.diffusion_timing:
        cmd.append("--diffusion-timing")
    if args.no_diffusion_gpu_sampling:
        cmd.append("--no-diffusion-gpu-sampling")
    if args.no_diffusion_device_selfcond:
        cmd.append("--no-diffusion-device-selfcond")
    if args.no_diffusion_device_denoise_loop:
        cmd.append("--no-diffusion-device-denoise-loop")
    if args.diffusion_pin_host_outputs:
        cmd.append("--diffusion-pin-host-outputs")
    if args.diffusion_fused_self_cond_embd:
        cmd.append("--diffusion-fused-self-cond-embd")
    if args.diffusion_fuse_final_softcap:
        cmd.append("--diffusion-fuse-final-softcap")
    if args.diffusion_cuda_direct_self_cond:
        cmd.append("--diffusion-cuda-direct-self-cond")
    if args.diffusion_cuda_final_tokens_on_stop:
        cmd.append("--diffusion-cuda-final-tokens-on-stop")
    if args.diffusion_cuda_fused_top_k_sample:
        cmd.append("--diffusion-cuda-fused-top-k-sample")
    if args.diffusion_cuda_parallel_full_softmax:
        cmd.append("--diffusion-cuda-parallel-full-softmax")
    if args.no_diffusion_cuda_fused_full_softmax:
        cmd.append("--no-diffusion-cuda-fused-full-softmax")
    if args.no_diffusion_cuda_fast_top_k:
        cmd.append("--no-diffusion-cuda-fast-top-k")
    if args.diffusion_cuda_mmq_max_x is not None:
        cmd.extend(["--diffusion-cuda-mmq-max-x", str(args.diffusion_cuda_mmq_max_x)])
    if args.diffusion_self_cond_top_k is not None:
        cmd.extend(["--diffusion-self-cond-top-k", str(args.diffusion_self_cond_top_k)])
    if args.diffusion_input_gpu_groups is not None:
        cmd.extend(["--diffusion-input-gpu-groups", str(args.diffusion_input_gpu_groups)])
    if args.diffusion_default_top_k is not None:
        cmd.extend(["--diffusion-default-top-k", str(args.diffusion_default_top_k)])
    if args.top_k is not None:
        cmd.extend(["--top-k", str(args.top_k)])
    if args.top_k_start is not None:
        cmd.extend(["--top-k-start", str(args.top_k_start)])
    if args.top_k_end is not None:
        cmd.extend(["--top-k-end", str(args.top_k_end)])
    if args.top_k_tail_correction:
        cmd.append("--top-k-tail-correction")

    for extra in args.extra_arg:
        cmd.append(extra)

    return cmd


def run_once(args: argparse.Namespace, prompt_index: int, prompt: str, run_index: int, keep: bool) -> dict[str, object]:
    cmd = build_command(args, prompt)
    start = time.perf_counter()
    proc = subprocess.run(
        cmd,
        cwd=args.cwd,
        env=clean_env(),
        text=True,
        encoding="utf-8",
        errors="replace",
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=args.timeout,
    )
    wall_ms = (time.perf_counter() - start) * 1000.0

    combined = (proc.stdout or "") + "\n" + (proc.stderr or "")
    parsed = parse_output(combined)
    run = {
        "prompt_index": prompt_index,
        "prompt": prompt,
        "run": run_index,
        "kept": keep,
        "exit_code": proc.returncode,
        "wall_ms": wall_ms,
        "command": cmd,
        **parsed,
    }

    log_base = args.log_dir / f"cli-p{prompt_index:02d}-run-{run_index:02d}"
    log_base.with_suffix(".stdout.log").write_text(proc.stdout or "", encoding="utf-8")
    log_base.with_suffix(".stderr.log").write_text(proc.stderr or "", encoding="utf-8")
    log_base.with_suffix(".json").write_text(json.dumps(run, indent=2), encoding="utf-8")
    run["stdout_log"] = str(log_base.with_suffix(".stdout.log"))
    run["stderr_log"] = str(log_base.with_suffix(".stderr.log"))
    run["json_log"] = str(log_base.with_suffix(".json"))

    if proc.returncode != 0:
        raise RuntimeError(
            f"CLI run {run_index} failed with exit code {proc.returncode}\n"
            f"Command: {quote_cmd(cmd)}\n"
            f"stderr tail:\n{tail(proc.stderr or '')}"
        )
    if keep and "generation" not in run:
        raise RuntimeError(
            f"CLI run {run_index} completed but no generation perf line was found\n"
            f"stderr tail:\n{tail(combined)}"
        )

    return run


def print_summary(summary: dict[str, object], runs: list[dict[str, object]]) -> None:
    print("\nCLI diffusion benchmark")
    print("-----------------------")
    print(f"runs kept: {summary['runs']}")
    print(f"wall ms avg/min/max: {summary['wall_ms_avg']:.2f} / {summary['wall_ms_min']:.2f} / {summary['wall_ms_max']:.2f}")

    if summary["generation_seconds_avg"] is not None:
        print(f"generation seconds avg: {summary['generation_seconds_avg']:.3f}")
        print(f"canvas tok/s avg: {summary['canvas_tokens_per_second_avg']:.3f}")
        print(f"ms/step avg: {summary['ms_per_step_avg']:.3f}")
        print(f"mean denoising steps/block: {summary['denoising_steps_avg']:.3f}")
    if summary["prefill_tokens_per_second_avg"] is not None:
        print(f"prefill tok/s avg: {summary['prefill_tokens_per_second_avg']:.3f}")

    last = runs[-1]
    generation = last.get("generation")
    if isinstance(generation, dict):
        print(
            "last run: "
            f"{generation['blocks']} block(s), "
            f"{generation['steps']} steps, "
            f"{generation['canvas_tokens']} canvas tokens, "
            f"{generation['answer_tokens']} answer tokens"
        )
    print(f"logs: {last['stdout_log']}, {last['stderr_log']}")


def write_report(args: argparse.Namespace, result: dict[str, object]) -> None:
    summary = result["summary"]
    prompt_rows = result["prompt_summaries"]
    runs = result["all_runs"]

    lines = [
        "# Diffusion Gemma CLI Benchmark Report",
        "",
        f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}",
        f"Binary: `{args.binary}`",
        f"Model: `{args.model}`",
        f"Prompts: {len(result['prompts'])}",
        f"Warmup per prompt: {args.warmup}",
        f"Runs per prompt: {args.repeat}",
        f"Max output tokens: {args.n_predict}",
        f"Configured diffusion steps: {args.diffusion_steps}",
        "",
        "## Overall",
        "",
        "| Metric | Value |",
        "| --- | ---: |",
        f"| Kept runs | {summary['runs']} |",
        f"| Canvas token throughput avg, tok/s | {fmt_num(summary['canvas_tokens_per_second_avg'])} |",
        f"| Per-step time avg, ms/step | {fmt_num(summary['ms_per_step_avg'])} |",
        f"| Total denoising steps | {fmt_num(summary['denoising_steps_total'], 0)} |",
        f"| Mean denoising steps/block | {fmt_num(summary['denoising_steps_avg'])} |",
        f"| Canvas tokens avg | {fmt_num(summary['canvas_tokens_avg'])} |",
        f"| Canvas tokens total | {fmt_num(summary['canvas_tokens_total'], 0)} |",
        f"| Blocks total | {fmt_num(summary['blocks_total'], 0)} |",
        f"| Blocks avg | {fmt_num(summary['blocks_avg'])} |",
        f"| Generation seconds avg | {fmt_num(summary['generation_seconds_avg'])} |",
        "",
        "## Prompt Summary",
        "",
        "| Prompt | Runs | Canvas tok/s avg | ms/step avg | Mean denoising steps/block | Canvas tokens avg | Blocks avg |",
        "| --- | ---: | ---: | ---: | ---: | ---: | ---: |",
    ]

    for row in prompt_rows:
        prompt_summary = row["summary"]
        lines.append(
            f"| {row['prompt_index']}. {prompt_preview(row['prompt'])} "
            f"| {prompt_summary['runs']} "
            f"| {fmt_num(prompt_summary['canvas_tokens_per_second_avg'])} "
            f"| {fmt_num(prompt_summary['ms_per_step_avg'])} "
            f"| {fmt_num(prompt_summary['denoising_steps_avg'])} "
            f"| {fmt_num(prompt_summary['canvas_tokens_avg'])} "
            f"| {fmt_num(prompt_summary['blocks_avg'])} |"
        )

    lines.extend([
        "",
        "## All Runs",
        "",
        "| Prompt | Run | Kept | Wall ms | Canvas tok/s | ms/step | Denoising steps | Canvas tokens | Blocks | Answer tokens |",
        "| --- | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
    ])

    for run in runs:
        generation = run.get("generation")
        generation = generation if isinstance(generation, dict) else {}
        sec_per_step = generation.get("seconds_per_step")
        lines.append(
            f"| {run['prompt_index']}. {prompt_preview(str(run['prompt']), 60)} "
            f"| {run['run']} "
            f"| {'yes' if run.get('kept') else 'warmup'} "
            f"| {fmt_num(run.get('wall_ms'))} "
            f"| {fmt_num(generation.get('canvas_tokens_per_second'))} "
            f"| {fmt_num(float(sec_per_step) * 1000.0 if sec_per_step is not None else None)} "
            f"| {fmt_num(generation.get('steps'), 0)} "
            f"| {fmt_num(generation.get('canvas_tokens'), 0)} "
            f"| {fmt_num(generation.get('blocks'), 0)} "
            f"| {fmt_num(generation.get('answer_tokens'), 0)} |"
        )

    args.report_out.parent.mkdir(parents=True, exist_ok=True)
    args.report_out.write_text("\n".join(lines) + "\n", encoding="utf-8")


def write_outputs_jsonl(args: argparse.Namespace, result: dict[str, object]) -> None:
    args.outputs_jsonl.parent.mkdir(parents=True, exist_ok=True)
    with args.outputs_jsonl.open("w", encoding="utf-8") as f:
        for run in result["all_runs"]:
            stdout_log = Path(str(run["stdout_log"]))
            text = stdout_log.read_text(encoding="utf-8", errors="replace") if stdout_log.exists() else ""
            record = {
                "prompt_index": run["prompt_index"],
                "prompt": run["prompt"],
                "run": run["run"],
                "kept": run["kept"],
                "text": text,
                "prefill": run.get("prefill"),
                "generation": run.get("generation"),
                "wall_ms": run.get("wall_ms"),
                "stdout_log": run.get("stdout_log"),
                "stderr_log": run.get("stderr_log"),
                "json_log": run.get("json_log"),
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Benchmark llama-diffusion-gemma-cli and summarize performance.")
    parser.add_argument("--binary", required=True, type=Path, help="Path to llama-diffusion-gemma-cli.")
    parser.add_argument("--model", required=True, type=Path, help="Path to the GGUF model.")
    parser.add_argument("--prompt", default="Hello")
    parser.add_argument("--prompt-file", type=Path, default=None, help="Path to a text file with one prompt per non-empty line.")
    parser.add_argument("--n-predict", type=int, default=256)
    parser.add_argument("--ctx-size", type=int, default=8096)
    parser.add_argument("--n-gpu-layers", type=int, default=999)
    parser.add_argument("--threads", type=int, default=8)
    parser.add_argument("--diffusion-steps", type=int, default=48)
    parser.add_argument("--log-verbosity", type=int, default=3)
    parser.add_argument("--repeat", type=int, default=1)
    parser.add_argument("--warmup", type=int, default=0)
    parser.add_argument("--timeout", type=float, default=600.0)
    parser.add_argument("--output-dir", type=Path, default=None, help="Output root directory. Defaults to the current directory.")
    parser.add_argument("--log-dir", type=Path, default=None)
    parser.add_argument("--json-out", type=Path, default=None)
    parser.add_argument("--report-out", type=Path, default=None)
    parser.add_argument("--outputs-jsonl", type=Path, default=None)
    parser.add_argument("--cwd", type=Path, default=None)

    parser.add_argument("--diffusion-timing", action="store_true")
    parser.add_argument("--no-diffusion-gpu-sampling", action="store_true")
    parser.add_argument("--no-diffusion-device-selfcond", action="store_true")
    parser.add_argument("--no-diffusion-device-denoise-loop", action="store_true")
    parser.add_argument("--diffusion-pin-host-outputs", action="store_true")
    parser.add_argument("--diffusion-fused-self-cond-embd", action="store_true")
    parser.add_argument("--diffusion-fuse-final-softcap", action="store_true")
    parser.add_argument("--diffusion-cuda-direct-self-cond", action="store_true")
    parser.add_argument("--diffusion-cuda-final-tokens-on-stop", action="store_true")
    parser.add_argument("--diffusion-cuda-fused-top-k-sample", action="store_true")
    parser.add_argument("--diffusion-cuda-parallel-full-softmax", action="store_true")
    parser.add_argument("--no-diffusion-cuda-fused-full-softmax", action="store_true")
    parser.add_argument("--no-diffusion-cuda-fast-top-k", action="store_true")
    parser.add_argument("--diffusion-cuda-mmq-max-x", type=int, default=64)
    parser.add_argument("--diffusion-self-cond-top-k", type=int, default=None)
    parser.add_argument("--diffusion-input-gpu-groups", type=int, default=None)
    parser.add_argument("--diffusion-default-top-k", type=int, default=None)
    parser.add_argument("--top-k", type=int, default=None)
    parser.add_argument("--top-k-start", type=int, default=None)
    parser.add_argument("--top-k-end", type=int, default=None)
    parser.add_argument("--top-k-tail-correction", action="store_true")
    parser.add_argument("--extra-arg", action="append", default=[], help="Extra CLI arg to append; repeat as needed.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if args.repeat < 1:
        print("error: --repeat must be at least 1", file=sys.stderr)
        return 2
    if args.warmup < 0:
        print("error: --warmup must be non-negative", file=sys.stderr)
        return 2

    args.binary = args.binary.expanduser()
    args.model = args.model.expanduser()
    args.prompt_file = args.prompt_file.expanduser() if args.prompt_file else None
    args.cwd = args.cwd.expanduser() if args.cwd else None
    args.output_dir = args.output_dir.expanduser() if args.output_dir else Path.cwd()
    args.log_dir = args.log_dir.expanduser() if args.log_dir else args.output_dir / "diffusion-cli"
    args.log_dir.mkdir(parents=True, exist_ok=True)
    args.report_out = args.report_out.expanduser() if args.report_out else args.log_dir / "report.md"
    args.outputs_jsonl = args.outputs_jsonl.expanduser() if args.outputs_jsonl else args.log_dir / "outputs.jsonl"

    if not args.binary.exists():
        print(f"error: binary not found: {args.binary}", file=sys.stderr)
        return 2
    if not args.model.exists():
        print(f"error: model not found: {args.model}", file=sys.stderr)
        return 2
    if args.prompt_file is not None and not args.prompt_file.exists():
        print(f"error: prompt file not found: {args.prompt_file}", file=sys.stderr)
        return 2
    try:
        prompts = load_prompts(args.prompt, args.prompt_file)
    except ValueError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    print(f"binary: {args.binary}")
    print(f"model: {args.model}")
    print(f"prompts: {len(prompts)}")
    print(f"output dir: {args.output_dir}")
    print(f"log dir: {args.log_dir}")

    kept_runs: list[dict[str, object]] = []
    all_runs: list[dict[str, object]] = []
    total = args.warmup + args.repeat
    for prompt_index, prompt in enumerate(prompts, start=1):
        print(f"prompt {prompt_index}/{len(prompts)}: {prompt_preview(prompt, 120)}")
        for i in range(total):
            keep = i >= args.warmup
            label = "run" if keep else "warmup"
            print(f"{label} {i + 1}/{total}...")
            run = run_once(args, prompt_index, prompt, i + 1, keep)
            all_runs.append(run)
            if keep:
                kept_runs.append(run)

    summary = summarize(kept_runs)
    result = {
        "kind": "diffusion-gemma-cli",
        "prompts": prompts,
        "summary": summary,
        "prompt_summaries": prompt_summaries(kept_runs),
        "runs": kept_runs,
        "all_runs": all_runs,
        "report": str(args.report_out),
        "outputs_jsonl": str(args.outputs_jsonl),
    }
    write_report(args, result)
    write_outputs_jsonl(args, result)
    print(f"report: {args.report_out}")
    print(f"outputs: {args.outputs_jsonl}")
    if args.json_out:
        args.json_out.parent.mkdir(parents=True, exist_ok=True)
        args.json_out.write_text(json.dumps(result, indent=2), encoding="utf-8")
        print(f"json: {args.json_out}")

    print_summary(summary, kept_runs)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())