# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+# "numpy",
+# "torch",
+# "kernels-benchmark-tools",
+# "matplotlib",
+# ]
+#
+# [tool.uv.sources]
+# kernels-benchmark-tools = { git = "https://github.com/drbh/kernels-benchmark-tools.git", branch = "main" }
+# ///
+import os
+import sys
+from pathlib import Path
+import json
+import torch # noqa: F401 # imported because upstream may expect torch to be importable
+import kernels_benchmark_tools as kbt
+
+# --- Matplotlib setup and helpers ------------------------------------------------
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import csv
+
+
+# Keep text as text (not paths) so CSS can style fonts, size, etc.
+mpl.rcParams["svg.fonttype"] = "none"
+# Make ids deterministic across builds
+mpl.rcParams["svg.hashsalt"] = "latency-benchmark-combined"
+# Avoid auto-closed figures interfering with our tagging
+mpl.rcParams["figure.autolayout"] = True
+# Make background transparent
+mpl.rcParams["figure.facecolor"] = "none"
+mpl.rcParams["axes.facecolor"] = "none"
+mpl.rcParams["savefig.facecolor"] = "none"
+mpl.rcParams["savefig.edgecolor"] = "none"
+
+def _slugify(s: str) -> str:
+ s = (s or "").strip().lower()
+ keep = []
+ for ch in s:
+ if ch.isalnum():
+ keep.append(ch)
+ elif ch in (" ", "-", "_", "/", ".", ":"):
+ keep.append("-")
+ else:
+ keep.append("")
+ out = "".join(keep)
+ while "--" in out:
+ out = out.replace("--", "-")
+ return out.strip("-") or "unnamed"
+
+def _tag_current_figure(default_series_prefix="series"):
+ """Attach SVG ids (gid) to key artists so they can be targeted from CSS."""
+ fig = plt.gcf()
+ if fig is None:
+ return
+
+ # Tag the figure itself
+ fig.set_gid("figure--latency")
+
+ for ax_idx, ax in enumerate(fig.get_axes(), start=1):
+ ax.set_gid(f"axes--{ax_idx}")
+
+ # Axis labels & title
+ if ax.get_title():
+ for t in ax.texts:
+ if t.get_text() == ax.get_title():
+ t.set_gid("title--main")
+ if ax.xaxis and ax.xaxis.get_label():
+ ax.xaxis.label.set_gid("label--x")
+ if ax.yaxis and ax.yaxis.get_label():
+ ax.yaxis.label.set_gid("label--y")
+
+ # Gridlines
+ for i, gl in enumerate(ax.get_xgridlines(), start=1):
+ gl.set_gid(f"grid-x--{i}")
+ for i, gl in enumerate(ax.get_ygridlines(), start=1):
+ gl.set_gid(f"grid-y--{i}")
+
+ # Legend block & entries
+ leg = ax.get_legend()
+ if leg is not None:
+ leg.set_gid("legend")
+ for i, txt in enumerate(leg.get_texts(), start=1):
+ label_slug = _slugify(txt.get_text())
+ txt.set_gid(f"legend-label--{label_slug or i}")
+
+ # Series (lines, patches)
+ # Lines
+ line_seen = {}
+ for ln in getattr(ax, "lines", []):
+ raw_label = ln.get_label() or ""
+ # Matplotlib uses labels beginning with "_" for non-legendable items
+ label = raw_label if not raw_label.startswith("_") else f"{default_series_prefix}"
+ slug = _slugify(label)
+ line_seen[slug] = line_seen.get(slug, 0) + 1
+ suffix = "" if line_seen[slug] == 1 else f"-{line_seen[slug]}"
+ ln.set_gid(f"series--{slug}{suffix}")
+
+ # Patches (bars, areas)
+ patch_seen = {}
+ for pt in getattr(ax, "patches", []):
+ label = getattr(pt, "get_label", lambda: "")() or f"{default_series_prefix}"
+ if isinstance(label, str) and label.startswith("_"):
+ label = default_series_prefix
+ slug = _slugify(label)
+ patch_seen[slug] = patch_seen.get(slug, 0) + 1
+ suffix = "" if patch_seen[slug] == 1 else f"-{patch_seen[slug]}"
+ pt.set_gid(f"series--{slug}{suffix}")
+
+def _postprocess_svg_add_classes(svg_path: Path):
+ """Add convenient CSS classes alongside ids (e.g., class='series grid grid-x')."""
+ try:
+ import xml.etree.ElementTree as ET
+ ET.register_namespace("", "http://www.w3.org/2000/svg")
+ tree = ET.parse(svg_path)
+ root = tree.getroot()
+ for el in root.iter():
+ el_id = el.attrib.get("id", "")
+ if not el_id:
+ continue
+ cls = []
+ if el_id.startswith("figure--"):
+ cls.append("figure")
+ elif el_id.startswith("axes--"):
+ cls.append("axes")
+ elif el_id.startswith("grid-x--"):
+ cls += ["grid", "grid-x"]
+ elif el_id.startswith("grid-y--"):
+ cls += ["grid", "grid-y"]
+ elif el_id.startswith("legend"):
+ cls.append("legend")
+ elif el_id.startswith("label--x"):
+ cls.append("xlabel")
+ elif el_id.startswith("label--y"):
+ cls.append("ylabel")
+ elif el_id.startswith("title--"):
+ cls.append("title")
+ elif el_id.startswith("series--"):
+ cls.append("series")
+ if cls:
+ # Preserve any existing class (unlikely from Matplotlib)
+ existing = el.attrib.get("class", "")
+ el.set("class", (existing + " " + " ".join(cls)).strip())
+ tree.write(svg_path, encoding="utf-8", xml_declaration=True)
+ except Exception as e:
+ print(f"✗ SVG postprocess (classes) skipped: {e}")
+
+# Monkey-patch savefig to force SVG & ensure tagging occurs even if kbt.viz saves internally.
+_orig_savefig = plt.savefig
+def _savefig_svg(fname, *args, **kwargs):
+ # Always save as SVG at a stable path for the artifact system
+ out = Path("latency.svg")
+ kwargs["format"] = "svg"
+ # Ensure everything we care about has ids before export
+ _tag_current_figure()
+ res = _orig_savefig(out, *args, **kwargs)
+ # Add helpful CSS classes on top of ids
+ _postprocess_svg_add_classes(out)
+ print(f"✓ Combined visualization saved as {out}")
+ return res
+
+plt.savefig = _savefig_svg # apply patch
+
+# Capture close calls in case kbt.viz() closes figures before we re-save
+_orig_close = plt.close
+_last_closed = {"fig": None}
+def _capture_close(arg=None):
+ try:
+ if hasattr(arg, "savefig"): # looks like a Figure
+ _last_closed["fig"] = arg
+ else:
+ _last_closed["fig"] = plt.gcf()
+ finally:
+ return _orig_close(arg)
+plt.close = _capture_close
+
+# --- Locate benchmark artifacts --------------------------------------------------
+cache_dirs = {
+ "Flash (PyTorch SDPA)": os.environ.get('UVNOTE_FILE_FLASH_ATTENTION_BENCHMARK'),
+ "MemEff (PyTorch SDPA)": os.environ.get('UVNOTE_FILE_MEM_EFFICIENT_ATTENTION_BENCHMARK'),
+ "Flash Attn 2": os.environ.get('UVNOTE_FILE_FLASH_ATTN2_BENCHMARK'),
+ "xFormers": os.environ.get('UVNOTE_FILE_XFORMERS_BENCHMARK'),
+ "SageAttention": os.environ.get('UVNOTE_FILE_SAGE_ATTENTION_BENCHMARK'),
+ "Compiled (default)": os.environ.get('UVNOTE_FILE_COMPILED_VARIANTS_BENCHMARK_DEFAULT'),
+ "Compiled (max-autotune)": os.environ.get('UVNOTE_FILE_COMPILED_VARIANTS_BENCHMARK_MAX_AUTOTUNE'),
+ "HF Kernels Flash Attn": os.environ.get('UVNOTE_FILE_HF_KERNELS_FLASH_ATTN_BENCHMARK'),
+ "HF Kernels Flash Attn3": os.environ.get('UVNOTE_FILE_HF_KERNELS_FLASH_ATTN3_BENCHMARK'),
+}
+
+print("LOADING BENCHMARK DATA")
+for name, cache_dir in cache_dirs.items():
+ print(f"{name:30s}: {cache_dir}")
+print()
+
+file_mapping = {
+ "Flash (PyTorch SDPA)": "attn.jsonl",
+ "MemEff (PyTorch SDPA)": "attn.jsonl",
+ "Flash Attn 2": "attn.jsonl",
+ "xFormers": "attn.jsonl",
+ "SageAttention": "attn.jsonl",
+ "Compiled (default)": "attn_default.jsonl",
+ "Compiled (max-autotune)": "attn_max_autotune.jsonl",
+ "HF Kernels Flash Attn": "attn.jsonl",
+ "HF Kernels Flash Attn3": "attn.jsonl",
+}
+
+all_paths = []
+for name, cache_dir in cache_dirs.items():
+ if cache_dir:
+ path = Path(cache_dir) / file_mapping[name]
+ if path.exists() and path.stat().st_size > 0:
+ all_paths.append(str(path))
+ print(f"✓ Found {name}: {path}")
+ else:
+ print(f"⊘ Empty/Missing {name}: {path}")
+ else:
+ print(f"✗ No cache dir for {name}")
+print()
+
+if not all_paths:
+ print("ERROR: No benchmark data files found!")
+ # restore patched functions before exiting
+ plt.savefig = _orig_savefig
+ plt.close = _orig_close
+ sys.exit(1)
+
+# --- Summary + Visualization -----------------------------------------------------
+print("COMBINED BENCHMARK SUMMARY\n")
+kbt.summarize(all_paths)
+print("\nGENERATING COMBINED VISUALIZATION\n")
+
+try:
+ # If kbt.viz saves internally, our patched savefig ensures SVG gets written,
+ # and it will carry ids/classes for CSS styling.
+ kbt.viz(all_paths)
+ # Safety net: if kbt.viz didn't save, save now.
+ # if not Path("latency.svg").exists():
+ # _tag_current_figure()
+ # plt.savefig("latency.svg")
+
+ plt.savefig("latency.svg") # ensure saved with tagging
+
+ print("✓ SVG visualization ready: latency.svg!")
+except ImportError as e:
+ print(f"✗ Visualization requires matplotlib: {e}")
+except Exception as e:
+ print(f"✗ Visualization failed: {e}")
+finally:
+ # Clean up patches to avoid side effects in later cells
+ plt.savefig = _orig_savefig
+ plt.close = _orig_close
+
+print()
+print("ANALYSIS COMPLETE")
+print(f"Total implementations analyzed: {len(all_paths)}")
+print(f"\nImplementations included:")
+for name, cache_dir in cache_dirs.items():
+ if cache_dir:
+ path = Path(cache_dir) / file_mapping[name]
+ if path.exists() and path.stat().st_size > 0:
+ print(f" ✓ {name}")
+
+
+
+# Collect all benchmark data and export to CSV
+all_data = {}
+for name, cache_dir in cache_dirs.items():
+ if cache_dir:
+ path = Path(cache_dir) / file_mapping[name]
+ if path.exists() and path.stat().st_size > 0:
+ with open(path, 'r') as f:
+ records = [json.loads(line) for line in f]
+ all_data[name] = records
+
+# Export to CSV
+csv_path = Path("latency.csv")
+with open(csv_path, 'w', newline='') as csvfile:
+ writer = csv.writer(csvfile)
+
+ # Write header
+ header = ["Implementation", "Impl ID", "Workload", "Batch", "Seq Length", "Heads", "Head Dim", "Dtype",
+ "Mean (ms)", "P10 (ms)", "P50 (ms)", "P90 (ms)", "Reps",
+ # "Compile (ms)",
+ "Peak Mem (MB)", "Backend", "Family"]
+ writer.writerow(header)
+
+ # Write data rows
+ for impl_name, records in all_data.items():
+ for record in records:
+ wl = record.get('wl', {})
+ lat = record.get('lat_ms', {})
+ tags = record.get('tags', {})
+
+ row = [
+ impl_name,
+ record.get('impl', ''),
+ wl.get('name', ''),
+ wl.get('batch', ''),
+ wl.get('seq_len', ''),
+ wl.get('heads', ''),
+ wl.get('head_dim', ''),
+ wl.get('dtype', ''),
+ lat.get('mean', ''),
+ lat.get('p10', ''),
+ lat.get('p50', ''),
+ lat.get('p90', ''),
+ lat.get('reps', ''),
+ # record.get('compile_ms', ''),
+ round(record.get('peak_bytes', 0) / 1024 / 1024, 2) if record.get('peak_bytes') else '',
+ tags.get('backend', ''),
+ tags.get('family', ''),
+ ]
+ writer.writerow(row)
+
+print(f"✓ CSV export complete: {csv_path}")
+print(f"Total implementations: {len(all_data)}")
+print(f"Total records: {sum(len(records) for records in all_data.values())}")
+