Flash Attention Benchmarks - Aggregated Results

# /// 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())}")