112 lines
5.2 KiB
Python
112 lines
5.2 KiB
Python
"""Generate a markdown summary for attention path benchmark JSON output."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
def ms(value: float) -> float:
|
|
return value * 1000.0
|
|
|
|
|
|
def ok_results(data: dict) -> list[dict]:
|
|
return [item for item in data["results"] if item.get("status") == "ok"]
|
|
|
|
|
|
def timing_ms(result: dict, name: str, field: str = "mean_s") -> float:
|
|
return ms(result["timings"][name][field])
|
|
|
|
|
|
def label(result: dict) -> str:
|
|
layout = result.get("layout_mode", "")
|
|
return f"{result['backend']}:{layout}" if layout else result["backend"]
|
|
|
|
|
|
def timing_names(result: dict) -> tuple[str, ...]:
|
|
layout = result.get("layout_mode")
|
|
if layout == "sol_fused":
|
|
return ("qkv_to_bshd", "attention_kernel", "output_reshape")
|
|
if layout == "sol_native":
|
|
return ("q_bshd_contiguous", "k_bshd_contiguous", "v_bshd_contiguous", "attention_kernel", "output_reshape")
|
|
return ("q_transpose_contiguous", "k_transpose_contiguous", "v_transpose_contiguous", "attention_kernel", "output_reshape")
|
|
|
|
|
|
def table(results: list[dict]) -> list[str]:
|
|
lines = ["| Segment | Backend | Kernel mean ms | Total mean ms | Layout mean ms | Kernel p50 ms | Total p50 ms |", "| --- | --- | ---: | ---: | ---: | ---: | ---: |"]
|
|
for item in sorted(results, key=lambda row: (row["segment"], row["layout_attention_total_mean_s"])):
|
|
total = ms(item["layout_attention_total_mean_s"])
|
|
kernel = timing_ms(item, "attention_kernel")
|
|
layout = total - kernel
|
|
p50_total = sum(timing_ms(item, name, "p50_s") for name in timing_names(item))
|
|
lines.append(f"| {item['segment']} | {label(item)} | {kernel:.3f} | {total:.3f} | {layout:.3f} | {timing_ms(item, 'attention_kernel', 'p50_s'):.3f} | {p50_total:.3f} |")
|
|
return lines
|
|
|
|
|
|
def fastest_by_segment(results: list[dict]) -> list[str]:
|
|
segments = sorted({item["segment"] for item in results})
|
|
lines = []
|
|
for segment in segments:
|
|
segment_results = [item for item in results if item["segment"] == segment]
|
|
best = min(segment_results, key=lambda item: item["layout_attention_total_mean_s"])
|
|
best_kernel = min(segment_results, key=lambda item: item["timings"]["attention_kernel"]["mean_s"])
|
|
lines.append(f"- `{segment}` fastest total: `{label(best)}` at {ms(best['layout_attention_total_mean_s']):.3f} ms; fastest kernel: `{label(best_kernel)}` at {timing_ms(best_kernel, 'attention_kernel'):.3f} ms.")
|
|
return lines
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--whole", type=Path, required=True)
|
|
parser.add_argument("--segments", type=Path, required=True)
|
|
parser.add_argument("--extra", type=Path, action="append", default=[])
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
whole = json.loads(args.whole.read_text(encoding="utf-8"))
|
|
segments = json.loads(args.segments.read_text(encoding="utf-8"))
|
|
extras = [json.loads(path.read_text(encoding="utf-8")) for path in args.extra]
|
|
whole_ok = ok_results(whole)
|
|
segment_ok = ok_results(segments)
|
|
extra_ok = [item for data in extras for item in ok_results(data)]
|
|
extra_lines = table(extra_ok) if extra_ok else ["No extra runs."]
|
|
failures = [item for item in segments["results"] + whole["results"] + [item for data in extras for item in data["results"]] if item.get("status") != "ok"]
|
|
|
|
lines = [
|
|
"# H3 Attention Path Profile",
|
|
"",
|
|
f"Whole-sequence input shape: `{whole['metadata']['h_msa_shape']}`; segments: `{whole['segments']}`.",
|
|
f"Measurements: `{whole['warmup']}` warmups and `{whole['iterations']}` iterations per backend.",
|
|
"",
|
|
"## Key Findings",
|
|
"",
|
|
*fastest_by_segment(whole_ok + segment_ok + extra_ok),
|
|
"- Sol-native BSHD removes the generic HND round trip and keeps exact Sol output parity.",
|
|
"- The current fused QKV split/layout prototype is opt-in because it is not faster than three PyTorch contiguous copies yet.",
|
|
"- Full-sequence Sol-Attn stays under the initial 50-55 ms target for attention plus layout on this captured H3 shape.",
|
|
"- Q/K/V layout remains the next memory-bandwidth target; the Sage-style output reshape is not worth attacking for Sol.",
|
|
"- `sdpa` is fastest for tiny text/secondary segments but loses heavily on the video-dominated path, so it is only interesting for segment-specialized dispatch.",
|
|
"",
|
|
"## Whole Sequence",
|
|
"",
|
|
*table(whole_ok),
|
|
"",
|
|
"## Extra Whole-Sequence Layout Runs",
|
|
"",
|
|
*extra_lines,
|
|
"",
|
|
"## Segment Sweep",
|
|
"",
|
|
*table(segment_ok),
|
|
]
|
|
if failures:
|
|
lines.extend(["", "## Expected Failures", ""])
|
|
lines.extend(f"- `{item['segment']}` / `{label(item)}`: `{item['error']}`" for item in failures)
|
|
lines.extend(["", "## Source Files", "", f"- Whole JSON: `{args.whole}`", f"- Segment JSON: `{args.segments}`"])
|
|
lines.extend(f"- Extra JSON: `{path}`" for path in args.extra)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|