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