mxflashattn 1.4.0__py3-none-any.whl

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Reproducible MXFlashAttn benchmark tools."""

benchmarks/compare.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+ENVIRONMENT_KEYS = (
+    "device",
+    "device_name",
+    "device_total_memory_bytes",
+    "torch",
+    "runtime",
+    "driver",
+    "flash_attn",
+    "vllm",
+    "warmup",
+    "repeats",
+    "seed",
+)
+
+
+def compare_runs(
+    baseline: dict,
+    candidate: dict,
+    *,
+    allow_mismatched_environment: bool = False,
+) -> list[dict]:
+    baseline_metadata = baseline.get("metadata", {})
+    candidate_metadata = candidate.get("metadata", {})
+    mismatched = [
+        key
+        for key in ENVIRONMENT_KEYS
+        if (key in baseline_metadata or key in candidate_metadata)
+        and baseline_metadata.get(key) != candidate_metadata.get(key)
+    ]
+    if mismatched and not allow_mismatched_environment:
+        raise ValueError(f"benchmark environments differ for: {', '.join(mismatched)}")
+
+    baseline_cases = {case["case_id"]: case for case in baseline.get("cases", [])}
+    candidate_cases = {case["case_id"]: case for case in candidate.get("cases", [])}
+    if not baseline_cases:
+        raise ValueError("baseline run contains no cases")
+
+    results = []
+    for case_id, base in baseline_cases.items():
+        if case_id not in candidate_cases:
+            raise ValueError(f"missing candidate case {case_id}")
+        candidate_case = candidate_cases[case_id]
+        base_latency = float(base["latency_ms"])
+        candidate_latency = float(candidate_case["latency_ms"])
+        if base_latency <= 0 or candidate_latency <= 0:
+            raise ValueError(f"latency must be positive for case {case_id}")
+        results.append(
+            {
+                "case_id": case_id,
+                "baseline_latency_ms": base_latency,
+                "latency_ms": candidate_latency,
+                "speedup": base_latency / candidate_latency,
+                "latency_reduction_pct": (base_latency - candidate_latency) / base_latency * 100,
+                "max_abs_error": candidate_case.get("max_abs_error"),
+                "backend": candidate_case.get("backend"),
+                "fallback_reason": candidate_case.get("fallback_reason"),
+            }
+        )
+    extra = candidate_cases.keys() - baseline_cases.keys()
+    if extra:
+        raise ValueError(f"candidate contains unmatched cases: {', '.join(sorted(extra))}")
+    return results
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Compare baseline and candidate benchmark JSON")
+    parser.add_argument("baseline", type=Path)
+    parser.add_argument("candidate", type=Path)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/comparison.json"))
+    parser.add_argument("--allow-mismatched-environment", action="store_true")
+    args = parser.parse_args()
+    baseline = json.loads(args.baseline.read_text(encoding="utf-8"))
+    candidate = json.loads(args.candidate.read_text(encoding="utf-8"))
+    result = compare_runs(
+        baseline,
+        candidate,
+        allow_mismatched_environment=args.allow_mismatched_environment,
+    )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(
+        json.dumps(
+            {
+                "metadata": {
+                    "baseline": baseline.get("metadata", {}),
+                    "candidate": candidate.get("metadata", {}),
+                    "allow_mismatched_environment": args.allow_mismatched_environment,
+                },
+                "cases": result,
+            },
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+    print(f"Wrote {len(result)} comparisons to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/configs/c500.yaml

@@ -0,0 +1,10 @@
+seed: 20261005
+matrix:
+  batch_size: [1, 2]
+  q_len: [1, 64]
+  kv_len: [128, 512]
+  head_dim: [64, 128]
+  gqa_ratio: [1, 4]
+  layout: [dense, varlen, paged]
+  dtype: [float16, bfloat16]
+  causal: [true, false]

benchmarks/fair_matrix.py

@@ -0,0 +1,191 @@
+"""Run a same-input six-column MXFlashAttn benchmark matrix.
+
+The operator columns are measured in one process on one device.  The vLLM
+columns are imported from separately captured model runs and are never used to
+inflate the operator speedup numbers.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import math
+import os
+import statistics
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from mxflashattn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache
+from mxflashattn import reference
+
+
+def _sync(device: torch.device) -> None:
+    module = getattr(torch, device.type, None)
+    sync = getattr(module, "synchronize", None)
+    if sync:
+        sync(device)
+
+
+def _measure(fn, device: torch.device, repeats: int) -> tuple[float, torch.Tensor]:
+    result = fn()
+    _sync(device)
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        result = fn()
+        _sync(device)
+        samples.append((time.perf_counter() - start) * 1000)
+    return statistics.median(samples), result
+
+
+def _make_cache(k: torch.Tensor, v: torch.Tensor, page_size: int = 16):
+    batch, kv_len, kv_heads, head_dim = k.shape
+    pages = (kv_len + page_size - 1) // page_size
+    cache_k = torch.zeros(batch * pages, page_size, kv_heads, head_dim, device=k.device, dtype=k.dtype)
+    cache_v = torch.zeros_like(cache_k)
+    table = torch.arange(batch * pages, device=k.device, dtype=torch.int32).reshape(batch, pages)
+    for b in range(batch):
+        for p in range(pages):
+            start = p * page_size
+            count = min(page_size, kv_len - start)
+            cache_k[int(table[b, p]), :count] = k[b, start : start + count]
+            cache_v[int(table[b, p]), :count] = v[b, start : start + count]
+    lengths = torch.full((batch,), kv_len, device=k.device, dtype=torch.int32)
+    return cache_k, cache_v, table, lengths
+
+
+def _dense_sdpa(q, k, v, causal):
+    ratio = q.shape[2] // k.shape[2]
+    kh = k.transpose(1, 2).repeat_interleave(ratio, dim=1)
+    vh = v.transpose(1, 2).repeat_interleave(ratio, dim=1)
+    return F.scaled_dot_product_attention(
+        q.transpose(1, 2), kh, vh, is_causal=causal, scale=1.0 / math.sqrt(q.shape[-1])
+    ).transpose(1, 2)
+
+
+def _case(index: int) -> dict:
+    # 36 deterministic combinations cover decode/prefill, both layouts,
+    # FP16/BF16, GQA/MQA, two head sizes, and causal/non-causal attention.
+    return {
+        "case_id": f"case-{index:02d}",
+        "batch_size": 1 + (index % 2),
+        "q_len": (1, 8)[(index // 2) % 2],
+        "kv_len": (64, 128)[(index // 4) % 2],
+        "head_dim": (64, 128)[(index // 8) % 2],
+        "gqa_ratio": (1, 4)[(index // 16) % 2],
+        "dtype": ("float16", "bfloat16")[(index // 32) % 2],
+        "causal": bool(index % 3),
+        "layout": ("dense", "varlen", "paged")[index % 3],
+    }
+
+
+def _run_case(case: dict, device: torch.device, repeats: int, vendor_dense, vendor_varlen, vendor_paged):
+    dtype = getattr(torch, case["dtype"])
+    b, q_len, kv_len, dim = case["batch_size"], case["q_len"], case["kv_len"], case["head_dim"]
+    kv_heads = 2
+    q_heads = kv_heads * case["gqa_ratio"]
+    generator = torch.Generator(device="cpu").manual_seed(20261005 + int(case["case_id"][-2:]))
+    q = torch.randn(b, q_len, q_heads, dim, dtype=dtype, generator=generator).to(device)
+    k = torch.randn(b, kv_len, kv_heads, dim, dtype=dtype, generator=generator).to(device)
+    v = torch.randn_like(k)
+    causal = case["causal"]
+    cache_k, cache_v, table, lengths = _make_cache(k, v)
+
+    def reference_run():
+        return reference.attention_batched(q, k, v, causal=causal)
+
+    def sdpa_run():
+        return _dense_sdpa(q, k, v, causal)
+
+    def mx_run():
+        if case["layout"] == "dense":
+            return flash_attn_func(q, k, v, causal=causal)
+        if case["layout"] == "varlen":
+            cu_q = torch.arange(0, (b + 1) * q_len, q_len, dtype=torch.int32, device=device)
+            cu_k = torch.arange(0, (b + 1) * kv_len, kv_len, dtype=torch.int32, device=device)
+            return flash_attn_varlen_func(q.reshape(-1, q_heads, dim), k.reshape(-1, kv_heads, dim), v.reshape(-1, kv_heads, dim), cu_q, cu_k, max_seqlen_q=q_len, max_seqlen_k=kv_len, causal=causal).reshape_as(q)
+        return flash_attn_with_kvcache(q, cache_k, cache_v, cache_seqlens=lengths, block_table=table, causal=causal)
+
+    def vendor_run():
+        if case["layout"] == "dense":
+            return vendor_dense(q, k, v, causal=causal)
+        if case["layout"] == "varlen":
+            cu_q = torch.arange(0, (b + 1) * q_len, q_len, dtype=torch.int32, device=device)
+            cu_k = torch.arange(0, (b + 1) * kv_len, kv_len, dtype=torch.int32, device=device)
+            return vendor_varlen(q.reshape(-1, q_heads, dim), k.reshape(-1, kv_heads, dim), v.reshape(-1, kv_heads, dim), cu_q, cu_k, max_seqlen_q=q_len, max_seqlen_k=kv_len, causal=causal).reshape_as(q)
+        return vendor_paged(q, cache_k, cache_v, cache_seqlens=lengths, block_table=table, causal=causal)
+
+    funcs = {"reference": reference_run, "torch_sdpa": sdpa_run, "mxflashattn_api": mx_run, "metax_vendor": vendor_run}
+    outputs, timings = {}, {}
+    for name, fn in funcs.items():
+        try:
+            elapsed, output = _measure(fn, device, repeats)
+            timings[name] = elapsed
+            outputs[name] = output.detach().float().cpu()
+        except Exception as exc:
+            timings[name] = None
+            outputs[name] = None
+            case.setdefault("errors", {})[name] = f"{type(exc).__name__}: {exc}"
+    base = outputs["reference"]
+    errors = {}
+    if base is not None:
+        for name, output in outputs.items():
+            if output is not None:
+                delta = (output - base).abs()
+                errors[name] = {"max_abs": float(delta.max()), "max_rel": float((delta / base.abs().clamp_min(1e-6)).max())}
+    case["timings_ms"] = timings
+    case["errors_vs_reference"] = errors
+    case["tokens"] = b * q_len
+    return case
+
+
+def run_matrix(device: torch.device, repeats: int, count: int) -> dict:
+    try:
+        vendor = __import__("flash_attn", fromlist=["flash_attn_func", "flash_attn_varlen_func", "flash_attn_with_kvcache"])
+        vendor_dense = vendor.flash_attn_func
+        vendor_varlen = vendor.flash_attn_varlen_func
+        vendor_paged = vendor.flash_attn_with_kvcache
+    except (ImportError, AttributeError) as exc:
+        vendor_dense = vendor_varlen = vendor_paged = None
+        vendor_error = f"{type(exc).__name__}: {exc}"
+    else:
+        vendor_error = None
+    cases = []
+    for index in range(count):
+        case = _case(index)
+        if vendor_dense is None:
+            case["errors"] = {"metax_vendor": vendor_error}
+            case["timings_ms"] = {name: None for name in ("reference", "torch_sdpa", "mxflashattn_api", "metax_vendor")}
+            cases.append(case)
+        else:
+            cases.append(_run_case(case, device, repeats, vendor_dense, vendor_varlen, vendor_paged))
+    metadata = {"device": str(device), "torch": torch.__version__, "vllm": _version("vllm"), "flash_attn": _version("flash-attn"), "vllm_metax": _version("vllm_metax"), "mxmaca": os.getenv("MXMACA_VERSION", "unknown"), "driver": os.getenv("MXMACA_DRIVER_VERSION", "unknown"), "seed": 20261005, "repeats": repeats, "case_count": len(cases), "vendor_error": vendor_error}
+    return {"schema_version": 1, "metadata": metadata, "cases": cases}
+
+
+def _version(name: str):
+    try:
+        return importlib.metadata.version(name)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--count", type=int, default=36)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/fair-matrix.json"))
+    args = parser.parse_args()
+    result = run_matrix(torch.device(args.device), args.repeats, args.count)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(f"Wrote {len(result['cases'])} fair benchmark cases to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/fair_report.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+from pathlib import Path
+
+
+def _values(cases, method):
+    return [float(c["timings_ms"][method]) for c in cases if c.get("timings_ms", {}).get(method) is not None]
+
+
+def render(matrix: dict, vendor_model: dict | None, candidate_model: dict | None) -> str:
+    cases = matrix["cases"]
+    methods = ("reference", "torch_sdpa", "metax_vendor", "mxflashattn_api")
+    lines = ["# v0.9 Fair Benchmark Report", "", "The operator matrix uses the same generated tensors for every method.", "Reference speedup and vendor speedup are reported separately.", "", f"- Cases: {len(cases)}", f"- Device: {matrix['metadata'].get('device')} / {matrix['metadata'].get('torch')}", f"- MetaX flash-attn: {matrix['metadata'].get('flash_attn')}", ""]
+    for method in methods:
+        values = _values(cases, method)
+        lines.append(f"- Median {method}: {statistics.median(values):.6f} ms" if values else f"- Median {method}: unavailable")
+    ref = _values(cases, "reference")
+    mx = _values(cases, "mxflashattn_api")
+    vendor_times = _values(cases, "metax_vendor")
+    if ref and mx:
+        lines.append(f"- Median reduction vs PyTorch reference: {(1 - statistics.median(mx) / statistics.median(ref)) * 100:.2f}%")
+    if vendor_times and mx:
+        lines.append(f"- Median reduction vs MetaX vendor: {(1 - statistics.median(mx) / statistics.median(vendor_times)) * 100:.2f}%")
+    lines.extend(["", "## Operator Matrix", "", "| Case | Layout | Dtype | Batch | Q | KV | Dim | GQA | Causal | Reference ms | SDPA ms | MetaX vendor ms | MXFlashAttn ms | MX abs error | Vendor abs error |", "|---|---|---:|---:|---:|---:|---:|---:|:---:|---:|---:|---:|---:|---:|---:|---:|"])
+    for c in cases:
+        t = c.get("timings_ms", {})
+        e = c.get("errors_vs_reference", {})
+        lines.append("| {case_id} | {layout} | {dtype} | {batch_size} | {q_len} | {kv_len} | {head_dim} | {gqa_ratio} | {causal} | {reference} | {sdpa} | {vendor} | {mx} | {mxerr} | {verr} |".format(
+            **c, reference=_fmt(t.get("reference")), sdpa=_fmt(t.get("torch_sdpa")), vendor=_fmt(t.get("metax_vendor")), mx=_fmt(t.get("mxflashattn_api")), mxerr=_fmt((e.get("mxflashattn_api") or {}).get("max_abs")), verr=_fmt((e.get("metax_vendor") or {}).get("max_abs"))))
+    lines.extend(["", "## vLLM Model Evidence", "", "These rows are model-level runs and are intentionally not merged into the operator speedup medians.", "", "| Model evidence | vLLM vendor | vLLM MXFlashAttn candidate |", "|---|---:|---:|"])
+    vendor_speed = statistics.median([r.get("tokens_per_second") for r in (vendor_model or {}).get("runs", []) if r.get("tokens_per_second")]) if vendor_model else None
+    candidate_speed = statistics.median([r.get("tokens_per_second") for r in (candidate_model or {}).get("runs", []) if r.get("tokens_per_second")]) if candidate_model else None
+    lines.append(f"| status | {(vendor_model or {}).get('status', 'unavailable')} | {(candidate_model or {}).get('status', 'unavailable')} |")
+    lines.append(f"| median tokens/s | {_fmt(vendor_speed)} | {_fmt(candidate_speed)} |")
+    for label, record in (("vendor", vendor_model), ("candidate", candidate_model)):
+        if not record:
+            lines.append(f"- {label}: unavailable")
+            continue
+        runs = record.get("runs", [])
+        speeds = [r.get("tokens_per_second") for r in runs if r.get("tokens_per_second")]
+        speed_text = f"{statistics.median(speeds):.4f}" if speeds else "unavailable"
+        lines.append(f"- {label}: status={record.get('status')}, runs={len(runs)}, median tokens/s={speed_text}")
+    return "\n".join(lines) + "\n"
+
+
+def _fmt(value):
+    return "" if value is None else f"{float(value):.6g}"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("matrix", type=Path)
+    parser.add_argument("--vendor-model", type=Path)
+    parser.add_argument("--candidate-model", type=Path)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/fair-matrix.md"))
+    parser.add_argument("--summary", type=Path, default=Path("artifacts/fair-matrix-summary.json"))
+    args = parser.parse_args()
+    matrix = json.loads(args.matrix.read_text(encoding="utf-8"))
+    vendor = json.loads(args.vendor_model.read_text(encoding="utf-8")) if args.vendor_model else None
+    candidate = json.loads(args.candidate_model.read_text(encoding="utf-8")) if args.candidate_model else None
+    args.output.write_text(render(matrix, vendor, candidate), encoding="utf-8")
+    summary = {"case_count": len(matrix["cases"]), "methods": ["pytorch_reference", "torch_sdpa", "metax_vendor", "mxflashattn_api", "vllm_vendor_model", "vllm_candidate_model"], "operator_methods": ["reference", "torch_sdpa", "metax_vendor", "mxflashattn_api"], "vllm_vendor_status": vendor.get("status") if vendor else None, "vllm_candidate_status": candidate.get("status") if candidate else None, "claims": {"reference_speedup_separate": True, "vendor_speedup_verified": False}}
+    args.summary.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(f"Wrote {len(matrix['cases'])} cases to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/model_fair_report.py

@@ -0,0 +1,60 @@
+"""Pair same-prompt vLLM vendor and candidate model runs."""
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+from pathlib import Path
+
+
+def build(vendor: dict, candidate: dict) -> dict:
+    if len(vendor.get("runs", [])) != len(candidate.get("runs", [])):
+        raise ValueError("vendor and candidate run counts differ")
+    rows = []
+    for left, right in zip(vendor["runs"], candidate["runs"]):
+        if left.get("prompt") != right.get("prompt"):
+            raise ValueError(f"prompt mismatch at index {left.get('index')}")
+        rows.append({
+            "index": left.get("index"),
+            "prompt": left.get("prompt"),
+            "prompt_tokens": left.get("prompt_tokens"),
+            "vendor_tokens_per_second": left.get("tokens_per_second"),
+            "candidate_tokens_per_second": right.get("tokens_per_second"),
+            "candidate_dispatch_verified": any(e.get("path") == "mxflashattn_decoder" for e in right.get("dispatch_events", [])),
+            "vendor_generated_text": left.get("generated_text"),
+            "candidate_generated_text": right.get("generated_text"),
+        })
+    vendor_speed = [r["vendor_tokens_per_second"] for r in rows if r["vendor_tokens_per_second"]]
+    candidate_speed = [r["candidate_tokens_per_second"] for r in rows if r["candidate_tokens_per_second"]]
+    return {
+        "schema_version": 1,
+        "kind": "vllm_same_prompt_model_matrix",
+        "case_count": len(rows),
+        "same_prompt": True,
+        "vendor_status": vendor.get("status"),
+        "candidate_status": candidate.get("status"),
+        "candidate_dispatch_verified": candidate.get("candidate_dispatch_verified", False),
+        "vendor_median_tokens_per_second": statistics.median(vendor_speed) if vendor_speed else None,
+        "candidate_median_tokens_per_second": statistics.median(candidate_speed) if candidate_speed else None,
+        "rows": rows,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("vendor", type=Path)
+    parser.add_argument("candidate", type=Path)
+    parser.add_argument("--json", type=Path, default=Path("artifacts/vllm-model-matrix-c500.json"))
+    parser.add_argument("--markdown", type=Path, default=Path("artifacts/vllm-model-matrix-c500.md"))
+    args = parser.parse_args()
+    result = build(json.loads(args.vendor.read_text()), json.loads(args.candidate.read_text()))
+    args.json.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
+    lines = ["# vLLM Same-Prompt Model Matrix", "", "This is model-level evidence; it is not merged into operator medians.", "", f"- Cases: {result['case_count']}", f"- Candidate dispatch verified: {result['candidate_dispatch_verified']}", f"- Vendor median tokens/s: {result['vendor_median_tokens_per_second']:.4f}", f"- Candidate median tokens/s: {result['candidate_median_tokens_per_second']:.4f}", "", "| Case | Prompt tokens | Vendor tokens/s | Candidate tokens/s | Candidate dispatch |", "|---:|---:|---:|---:|:---:|"]
+    for row in result["rows"]:
+        lines.append(f"| {row['index']} | {row['prompt_tokens']} | {row['vendor_tokens_per_second']:.4f} | {row['candidate_tokens_per_second']:.4f} | {row['candidate_dispatch_verified']} |")
+    args.markdown.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    print(f"Wrote {result['case_count']} same-prompt model cases")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/plot_summary.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from statistics import median
+
+
+LAYOUTS = ("dense", "varlen", "paged")
+QUERY_LENGTHS = (1, 64)
+COLORS = {"dense": "#0072B2", "varlen": "#009E73", "paged": "#D55E00"}
+
+
+def summarize_matrix(run: dict) -> list[dict]:
+    groups: dict[tuple[int, str], list[float]] = {}
+    for case in run.get("cases", []):
+        if case.get("fallback_reason") or case.get("backend") == "reference":
+            continue
+        baseline = case.get("baseline_latency_ms")
+        latency = case.get("latency_ms")
+        key = (case.get("q_len"), case.get("layout"))
+        if (
+            baseline is None
+            or latency is None
+            or baseline <= 0
+            or key[0] not in QUERY_LENGTHS
+            or key[1] not in LAYOUTS
+        ):
+            continue
+        reduction = (float(baseline) - float(latency)) / float(baseline) * 100
+        groups.setdefault(key, []).append(reduction)
+
+    return [
+        {
+            "q_len": q_len,
+            "layout": layout,
+            "median_reduction_pct": median(values),
+            "target_cases": sum(value >= 20 for value in values),
+            "case_count": len(values),
+        }
+        for q_len in QUERY_LENGTHS
+        for layout in LAYOUTS
+        if (values := groups.get((q_len, layout)))
+    ]
+
+
+def render_plot(run: dict, output: Path) -> None:
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError as error:
+        raise RuntimeError("install mxflashattn[reports] to generate the benchmark chart") from error
+
+    summary = summarize_matrix(run)
+    if not summary:
+        raise ValueError("benchmark run contains no comparable measurements")
+
+    metadata = run.get("metadata", {})
+    figure, axis = plt.subplots(figsize=(12, 6.5))
+    centers = (0.0, 1.7)
+    width = 0.28
+    offsets = {layout: (index - 1) * width for index, layout in enumerate(LAYOUTS)}
+    for layout in LAYOUTS:
+        for q_len, center in zip(QUERY_LENGTHS, centers):
+            item = next(
+                (row for row in summary if row["q_len"] == q_len and row["layout"] == layout),
+                None,
+            )
+            if item is None:
+                continue
+            x = center + offsets[layout]
+            label = layout if q_len == QUERY_LENGTHS[0] else "_nolegend_"
+            bar = axis.bar(x, item["median_reduction_pct"], width, color=COLORS[layout], label=label)
+            axis.annotate(
+                f"{item['median_reduction_pct']:.0f}%\n{item['target_cases']}/{item['case_count']}",
+                (bar[0].get_x() + bar[0].get_width() / 2, bar[0].get_height()),
+                xytext=(0, 5),
+                textcoords="offset points",
+                ha="center",
+                va="bottom",
+                fontsize=9,
+            )
+
+    axis.axhline(20, color="#565D64", linestyle="--", linewidth=1.2, label="20% target")
+    axis.set_xticks(centers, ["Decode (q=1)", "Prefill (q=64)"])
+    axis.set_ylabel("Median latency reduction vs PyTorch reference (%)")
+    axis.set_title("MXFlashAttn C500 Operator Benchmark")
+    axis.grid(axis="y", color="#D9DEE3", linewidth=0.7)
+    axis.set_axisbelow(True)
+    axis.spines[["top", "right"]].set_visible(False)
+    axis.legend(ncols=4, loc="upper right", frameon=False)
+    upper = max(75, max(item["median_reduction_pct"] for item in summary) + 15)
+    axis.set_ylim(min(-10, min(item["median_reduction_pct"] for item in summary) - 15), upper)
+
+    run_details = (
+        f"{metadata.get('device_name', 'unknown')} | MXMACA {metadata.get('runtime', 'unknown')} | "
+        f"PyTorch {metadata.get('torch', 'unknown')} | {len(run.get('cases', []))} cases | "
+        f"seed {metadata.get('seed', 'unknown')}, warmup {metadata.get('warmup', 'unknown')}, "
+        f"repeats {metadata.get('repeats', 'unknown')}"
+    )
+    footer = "Labels show median reduction and cases at or above the 20% target.\n" + run_details
+    figure.text(0.5, 0.035, footer, ha="center", va="bottom", fontsize=9, color="#41484F")
+    figure.tight_layout(rect=(0.02, 0.13, 0.98, 0.98))
+    output.parent.mkdir(parents=True, exist_ok=True)
+    figure.savefig(output, dpi=180, facecolor="white")
+    plt.close(figure)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Plot summarized benchmark latency reductions")
+    parser.add_argument("input", type=Path)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/benchmark-summary.png"))
+    args = parser.parse_args()
+    run = json.loads(args.input.read_text(encoding="utf-8"))
+    render_plot(run, args.output)
+    print(f"Wrote benchmark chart to {args.output}")
+
+
+if __name__ == "__main__":
+    main()