#!/usr/bin/env python3 """Plot comparable latency progress from the kernel ablation benchmark ledgers. Only complete, all-PASS canonical official-suite runs are comparable: * GDN prefill: 100 workloads * FP8 MoE: 19 workloads * DSA: 23 workloads The time plots include clean and dirty-worktree development runs. The token plots select the lowest-mean clean run for each model/workflow when available. """ from __future__ import annotations import argparse import bisect import csv import json import math import statistics from collections import Counter, defaultdict from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path import matplotlib.pyplot as plt from matplotlib.lines import Line2D ROOT = Path(__file__).resolve().parent KERNELS = { "gdn-prefill": { "title": "GDN Prefill", "expected": 100, "token_key": "total_seq_len", "token_label": "Total sequence tokens", "xscale": "log", }, "fp8-moe": { "title": "FP8 MoE", "expected": 19, "token_key": "seq_len", "token_label": "Sequence tokens", "xscale": "log", }, "dsa": { "title": "DeepSeek Sparse Attention", "expected": 23, "token_key": "num_tokens", "token_label": "Query tokens", "xscale": "linear", }, } MODELS = { "claude-opus-4.8": ("Opus 4.8", "#1f77b4"), "fable-5": ("Fable 5", "#ff7f0e"), "gpt-5.5": ("GPT-5.5", "#2ca02c"), "gpt-5.6-sol": ("GPT-5.6-sol", "#d62728"), } WORKFLOWS = { "goal": ("goal", "-"), "goal-arar": ("goal-arar", "-"), "kda-humanize": ("kda-humanize", "--"), } @dataclass(frozen=True) class Series: model_dir: str workflow_dir: str kernel: str path: Path @property def model(self) -> str: return MODELS[self.model_dir][0] @property def workflow(self) -> str: return WORKFLOWS[self.workflow_dir][0] @property def label(self) -> str: return f"{self.model} | {self.workflow}" @property def color(self) -> str: return MODELS[self.model_dir][1] @property def linestyle(self) -> str: return WORKFLOWS[self.workflow_dir][1] @dataclass class Run: series: Series timestamp: datetime timestamp_text: str commit: str dirty: bool rows: list[dict[str, str]] mean_kernel_ms: float mean_baseline_ms: float arithmetic_mean_speedup: float def parse_bool(value: str) -> bool: return value.strip().lower() == "true" def discover_series() -> list[Series]: found: list[Series] = [] for model_dir in MODELS: for workflow_dir in WORKFLOWS: for kernel in KERNELS: path = ROOT / model_dir / workflow_dir / kernel / "benchmark.csv" if path.exists(): found.append(Series(model_dir, workflow_dir, kernel, path)) return found def read_valid_runs(series: Series) -> list[Run]: expected = KERNELS[series.kernel]["expected"] groups: dict[tuple[str, str, str, str, str], list[dict[str, str]]] = defaultdict(list) with series.path.open(newline="") as handle: for row in csv.DictReader(handle): if row.get("suite") != "official": continue if int(row.get("workload_count", "0")) != expected: continue key = ( row["timestamp_utc"], row["git_commit"], row["git_dirty"], row["op"], row["workload_count"], ) groups[key].append(row) runs: list[Run] = [] for key, rows in groups.items(): if len(rows) != expected: continue if len({row["workload_id"] for row in rows}) != expected: continue if not all(parse_bool(row["passed"]) for row in rows): continue try: kernel_ms = [float(row["kernel_ms"]) for row in rows] baseline_ms = [float(row["baseline_ms"]) for row in rows] speedups = [float(row["speedup"]) for row in rows] except (KeyError, ValueError): continue if not all(math.isfinite(x) and x > 0 for x in kernel_ms + baseline_ms): continue timestamp_text, commit, dirty_text, _, _ = key runs.append( Run( series=series, timestamp=datetime.fromisoformat(timestamp_text.replace("Z", "+00:00")), timestamp_text=timestamp_text, commit=commit, dirty=parse_bool(dirty_text), rows=rows, mean_kernel_ms=statistics.fmean(kernel_ms), mean_baseline_ms=statistics.fmean(baseline_ms), arithmetic_mean_speedup=statistics.fmean(speedups), ) ) return sorted(runs, key=lambda run: run.timestamp) def cumulative_timeline(events: list[tuple[datetime, int]]) -> list[tuple[datetime, int]]: events.sort(key=lambda item: item[0]) total = 0 timeline = [] for timestamp, count in events: if count <= 0: continue total += count timeline.append((timestamp, total)) return timeline def parse_timestamp(value: str | None) -> datetime | None: if not value: return None try: return datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError: return None def claude_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]: encoded = str(series.path.parent.resolve()).replace("/", "-").replace(".", "-") project_dir = Path.home() / ".claude" / "projects" / encoded events: list[tuple[datetime, int]] = [] seen_responses: set[str] = set() for path in project_dir.rglob("*.jsonl") if project_dir.exists() else []: with path.open(errors="replace") as handle: for line in handle: try: record = json.loads(line) except json.JSONDecodeError: continue message = record.get("message") or {} usage = message.get("usage") if record.get("type") != "assistant" or not isinstance(usage, dict): continue response_id = record.get("requestId") or message.get("id") or record.get("uuid") if not response_id or response_id in seen_responses: continue timestamp = parse_timestamp(record.get("timestamp")) if timestamp is None: continue seen_responses.add(response_id) count = sum( int(usage.get(field, 0) or 0) for field in ( "input_tokens", "cache_creation_input_tokens", "cache_read_input_tokens", "output_tokens", ) ) events.append((timestamp, count)) return cumulative_timeline(events), "Claude session JSONL (main + subagents)" def codex_session_files(cwd: Path) -> list[Path]: result = [] session_root = Path.home() / ".codex" / "sessions" for path in session_root.rglob("*.jsonl") if session_root.exists() else []: try: with path.open() as handle: first = json.loads(handle.readline()) if first.get("type") != "session_meta": continue recorded_cwd = Path(first["payload"]["cwd"]).resolve() except (OSError, KeyError, json.JSONDecodeError): continue if recorded_cwd == cwd.resolve(): result.append(path) return result def codex_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]: events: list[tuple[datetime, int]] = [] files = codex_session_files(series.path.parent) for path in files: previous_total = 0 with path.open(errors="replace") as handle: for line in handle: try: record = json.loads(line) except json.JSONDecodeError: continue payload = record.get("payload") or {} if record.get("type") != "event_msg" or payload.get("type") != "token_count": continue total = ( ((payload.get("info") or {}).get("total_token_usage") or {}).get("total_tokens") ) timestamp = parse_timestamp(record.get("timestamp")) if total is None or timestamp is None: continue total = int(total) delta = total - previous_total if total >= previous_total else total previous_total = total if delta > 0: events.append((timestamp, delta)) return cumulative_timeline(events), f"Codex token_count events ({len(files)} sessions)" def omh_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]: artifact_dir = series.path.parent / "workflow-output" / "omh-runtime" / "artifacts" events: list[tuple[datetime, int]] = [] seen_responses: set[str] = set() files = list(artifact_dir.rglob("*.jsonl")) if artifact_dir.exists() else [] for path in files: with path.open(errors="replace") as handle: for line in handle: try: record = json.loads(line) except json.JSONDecodeError: continue message = record.get("message") or {} usage = message.get("usage") if ( record.get("type") != "message" or message.get("role") != "assistant" or not isinstance(usage, dict) ): continue response_id = message.get("responseId") or record.get("id") if not response_id or response_id in seen_responses: continue timestamp = parse_timestamp(record.get("timestamp") or message.get("timestamp")) if timestamp is None: continue seen_responses.add(response_id) count = usage.get("totalTokens") if count is None: count = sum( int(usage.get(field, 0) or 0) for field in ("input", "output", "cacheRead", "cacheWrite") ) events.append((timestamp, int(count or 0))) return cumulative_timeline(events), f"OMH artifact usage ({len(files)} transcripts)" def token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]: if series.workflow_dir == "kda-humanize": return omh_token_timeline(series) if series.model_dir in {"claude-opus-4.8", "fable-5"}: return claude_token_timeline(series) return codex_token_timeline(series) def tokens_at(timeline: list[tuple[datetime, int]], timestamp: datetime) -> int | None: times = [item[0] for item in timeline] index = bisect.bisect_right(times, timestamp) - 1 return timeline[index][1] if index >= 0 else None def set_plot_style() -> None: plt.rcParams.update( { "figure.dpi": 140, "savefig.dpi": 180, "font.size": 10, "axes.titlesize": 14, "axes.labelsize": 11, "axes.grid": True, "grid.alpha": 0.24, "grid.linestyle": ":", "legend.fontsize": 8.5, } ) def plot_time(kernel: str, runs_by_series: dict[Series, list[Run]], out_dir: Path) -> None: config = KERNELS[kernel] fig, ax = plt.subplots(figsize=(11.8, 6.8)) for series in sorted(runs_by_series, key=lambda item: item.label): runs = runs_by_series[series] if not runs: continue start = runs[0].timestamp hours = [(run.timestamp - start).total_seconds() / 3600 for run in runs] values = [run.mean_kernel_ms for run in runs] best = [] current = math.inf for value in values: current = min(current, value) best.append(current) if len(hours) == 1: ax.scatter(hours, best, color=series.color, s=38, label=series.label, zorder=3) else: ax.plot( hours, best, color=series.color, linestyle=series.linestyle, linewidth=2.2, label=series.label, ) ax.set_yscale("log") ax.set_xlabel("Hours since first complete all-PASS official run (per experiment)") ax.set_ylabel("Arithmetic mean kernel latency across official suite (ms, log scale)") ax.set_title(f"{config['title']}: latency progress over experiment time") ax.text( 0.01, 0.01, "Continuous lines connect observed best-so-far full-suite checkpoints; isolated dots denote one-result experiments", transform=ax.transAxes, fontsize=8.5, color="#555555", ) ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), frameon=False) fig.tight_layout() save_figure(fig, out_dir / f"{kernel}_latency_vs_time") def choose_best_run(runs: list[Run]) -> Run: clean = [run for run in runs if not run.dirty] return min(clean or runs, key=lambda run: run.mean_kernel_ms) def token_curve( runs: list[Run], timeline: list[tuple[datetime, int]] ) -> list[tuple[int, float, float, Run]]: mapped = [] for run in runs: token_count = tokens_at(timeline, run.timestamp) if token_count is not None and token_count > 0: mapped.append((token_count, run.mean_kernel_ms, run)) mapped.sort(key=lambda item: (item[0], item[2].timestamp)) # Multiple benchmark runs can occur before another model response updates # the session counter. Keep the lowest latency at each cumulative count. collapsed: dict[int, tuple[float, Run]] = {} for token_count, latency, run in mapped: previous = collapsed.get(token_count) if previous is None or latency < previous[0]: collapsed[token_count] = (latency, run) result = [] best = math.inf for token_count, (latency, run) in sorted(collapsed.items()): best = min(best, latency) result.append((token_count, latency, best, run)) return result def plot_tokens( kernel: str, runs_by_series: dict[Series, list[Run]], timelines: dict[Series, tuple[list[tuple[datetime, int]], str]], out_dir: Path, ) -> None: config = KERNELS[kernel] fig, ax = plt.subplots(figsize=(11.8, 6.8)) for series in sorted(runs_by_series, key=lambda item: item.label): runs = runs_by_series[series] if not runs: continue points = token_curve(runs, timelines[series][0]) if not points: continue xs = [point[0] / 1_000_000 for point in points] ys = [point[2] for point in points] if len(xs) == 1: ax.scatter(xs, ys, color=series.color, s=38, label=series.label, zorder=3) else: ax.plot( xs, ys, color=series.color, linestyle=series.linestyle, linewidth=2.2, label=series.label, ) ax.set_xscale("log") ax.set_yscale("log") ax.set_xlabel("Cumulative agent-session tokens processed (millions, log scale)") ax.set_ylabel("Best-so-far mean official-suite kernel latency (ms, log scale)") ax.set_title(f"{config['title']}: latency versus cumulative agent tokens") ax.text( 0.01, 0.01, "Lines connect observed monotonic best-so-far checkpoints; usage includes cached input/output and workflow subagents/reviewers", transform=ax.transAxes, fontsize=8.5, color="#555555", ) ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), frameon=False) fig.tight_layout() save_figure(fig, out_dir / f"{kernel}_latency_vs_tokens") def save_figure(fig: plt.Figure, stem: Path) -> None: fig.savefig(stem.with_suffix(".png"), bbox_inches="tight") fig.savefig(stem.with_suffix(".svg"), bbox_inches="tight") plt.close(fig) def write_summary( all_series: list[Series], runs_by_kernel: dict[str, dict[Series, list[Run]]], timelines: dict[Series, tuple[list[tuple[datetime, int]], str]], out_dir: Path, ) -> None: fields = [ "kernel", "model", "workflow", "benchmark_csv", "valid_full_runs", "token_mapped_runs", "token_source", "tokens_at_first_mapped_run", "tokens_at_last_mapped_run", "first_valid_utc", "last_valid_utc", "observed_hours", "best_commit", "best_git_dirty", "best_mean_kernel_ms", "best_mean_baseline_ms", "best_arithmetic_mean_speedup", ] series_lookup = {(s.kernel, s.model_dir, s.workflow_dir): s for s in all_series} rows = [] for kernel in KERNELS: for model_dir in MODELS: workflows = ("goal", "kda-humanize") if model_dir in {"claude-opus-4.8", "fable-5"} else ("goal-arar", "kda-humanize") for workflow_dir in workflows: series = series_lookup.get((kernel, model_dir, workflow_dir)) runs = runs_by_kernel[kernel].get(series, []) if series else [] if runs: best = choose_best_run(runs) curve = token_curve(runs, timelines[series][0]) observed_hours = (runs[-1].timestamp - runs[0].timestamp).total_seconds() / 3600 rows.append( { "kernel": kernel, "model": MODELS[model_dir][0], "workflow": WORKFLOWS[workflow_dir][0], "benchmark_csv": str(series.path.relative_to(ROOT)), "valid_full_runs": len(runs), "token_mapped_runs": len(curve), "token_source": timelines[series][1], "tokens_at_first_mapped_run": curve[0][0] if curve else "", "tokens_at_last_mapped_run": curve[-1][0] if curve else "", "first_valid_utc": runs[0].timestamp_text, "last_valid_utc": runs[-1].timestamp_text, "observed_hours": f"{observed_hours:.4f}", "best_commit": best.commit, "best_git_dirty": best.dirty, "best_mean_kernel_ms": f"{best.mean_kernel_ms:.9g}", "best_mean_baseline_ms": f"{best.mean_baseline_ms:.9g}", "best_arithmetic_mean_speedup": f"{best.arithmetic_mean_speedup:.9g}", } ) else: rows.append( { "kernel": kernel, "model": MODELS[model_dir][0], "workflow": WORKFLOWS[workflow_dir][0], "benchmark_csv": str(series.path.relative_to(ROOT)) if series else "missing", "valid_full_runs": 0, "token_mapped_runs": 0, "token_source": timelines[series][1] if series else "missing benchmark.csv", "tokens_at_first_mapped_run": "", "tokens_at_last_mapped_run": "", "first_valid_utc": "", "last_valid_utc": "", "observed_hours": "", "best_commit": "", "best_git_dirty": "", "best_mean_kernel_ms": "", "best_mean_baseline_ms": "", "best_arithmetic_mean_speedup": "", } ) with (out_dir / "coverage_summary.csv").open("w", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fields) writer.writeheader() writer.writerows(rows) token_fields = [ "kernel", "model", "workflow", "benchmark_timestamp_utc", "git_commit", "git_dirty", "cumulative_agent_tokens", "mean_kernel_ms", "best_so_far_mean_kernel_ms", "token_source", ] token_rows = [] for kernel, by_series in runs_by_kernel.items(): for series, runs in by_series.items(): source = timelines[series][1] for token_count, latency, best, run in token_curve(runs, timelines[series][0]): token_rows.append( { "kernel": kernel, "model": series.model, "workflow": series.workflow, "benchmark_timestamp_utc": run.timestamp_text, "git_commit": run.commit, "git_dirty": run.dirty, "cumulative_agent_tokens": token_count, "mean_kernel_ms": f"{latency:.9g}", "best_so_far_mean_kernel_ms": f"{best:.9g}", "token_source": source, } ) token_rows.sort(key=lambda row: (row["kernel"], row["model"], row["workflow"], int(row["cumulative_agent_tokens"]))) with (out_dir / "token_curve_points.csv").open("w", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=token_fields) writer.writeheader() writer.writerows(token_rows) missing = [row for row in rows if row["valid_full_runs"] == 0] with (out_dir / "README.md").open("w") as handle: handle.write("# Kernel ablation latency plots\n\n") generated_at = datetime.now(timezone.utc).isoformat(timespec="seconds") handle.write(f"Snapshot generated at `{generated_at}` from the workspace `benchmark.csv` ledgers.\n\n") handle.write("## Comparison rules\n\n") handle.write("- Only canonical `official` runs with the full expected workload count and all rows passing are plotted.\n") handle.write("- Expected suite sizes: GDN prefill 100, FP8 MoE 19, DSA 23.\n") handle.write("- Time starts at each experiment's first valid full-suite result; curves are not extended to 12 hours.\n") handle.write("- Time charts connect observed best-so-far mean kernel latency checkpoints with continuous lines.\n") handle.write("- Token charts use cumulative agent-session usage from local session logs, not workload sequence length.\n") handle.write("- Token curves connect best-so-far checkpoints and are therefore monotonically non-increasing.\n") handle.write("- Isolated dots are retained only for experiments with exactly one comparable result.\n") handle.write("- Claude goal runs include main/subagent usage; Codex goal runs include all exact-cwd sessions; KDA includes OMH builder/reviewer/judge artifacts.\n") handle.write("- Usage is provider-native total processed tokens, including cached input and output.\n") handle.write("- Latency axes use milliseconds and logarithmic scaling.\n\n") handle.write("## Missing comparable results\n\n") if missing: for row in missing: handle.write(f"- {row['kernel']}: {row['model']} | {row['workflow']} (no complete all-PASS official run)\n") else: handle.write("None.\n") handle.write("\nSee `coverage_summary.csv` for coverage and `token_curve_points.csv` for every plotted token/latency point.\n") handle.write("\nRegenerate with `uv run --with matplotlib python plot_experiment_results.py`.\n") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--output-dir", type=Path, default=ROOT / "ablation-plots", help="Directory for PNG, SVG, and coverage metadata", ) args = parser.parse_args() out_dir = args.output_dir.resolve() out_dir.mkdir(parents=True, exist_ok=True) set_plot_style() all_series = discover_series() runs_by_kernel: dict[str, dict[Series, list[Run]]] = {kernel: {} for kernel in KERNELS} for series in all_series: runs_by_kernel[series.kernel][series] = read_valid_runs(series) timelines = {series: token_timeline(series) for series in all_series} for kernel in KERNELS: plot_time(kernel, runs_by_kernel[kernel], out_dir) plot_tokens(kernel, runs_by_kernel[kernel], timelines, out_dir) write_summary(all_series, runs_by_kernel, timelines, out_dir) print(f"Wrote plots to {out_dir}") for kernel in KERNELS: count = sum(bool(runs) for runs in runs_by_kernel[kernel].values()) points = sum(len(runs) for runs in runs_by_kernel[kernel].values()) print(f" {kernel}: {count} series, {points} complete all-PASS runs") if __name__ == "__main__": main()