File size: 1,001 Bytes
58b76f1 05bebc1 58b76f1 81fff32 58b76f1 81fff32 05bebc1 81fff32 05bebc1 58b76f1 05bebc1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "numpy",
# "torch==2.8.0",
# "kernels-benchmark-tools",
# "matplotlib",
# ]
#
# [tool.uv.sources]
# kernels-benchmark-tools = { path = "../../../../../tools", editable = true }
# ///
from kernels_benchmark_tools.core.visuals import generate_combined_results
# Map display names to uvnote environment variables
cache_env_map = {
"Flash (PyTorch SDPA)": "UVNOTE_FILE_FLASH_ATTENTION_BENCHMARK",
"MemEff (PyTorch SDPA)": "UVNOTE_FILE_MEM_EFFICIENT_ATTENTION_BENCHMARK",
"xFormers": "UVNOTE_FILE_XFORMERS_BENCHMARK",
"HF Kernels Flash Attn": "UVNOTE_FILE_HF_KERNELS_FLASH_ATTN_BENCHMARK",
"HF Kernels Flash Attn3": "UVNOTE_FILE_HF_KERNELS_FLASH_ATTN3_BENCHMARK",
"SageAttention": "UVNOTE_FILE_SAGE_ATTENTION_BENCHMARK",
}
# Generate combined results with visualization
generate_combined_results(
cache_env_map=cache_env_map,
output_filename="attention.jsonl",
svg_filename="latency.svg"
) |