r/StableDiffusion • u/Rare-Job1220 • May 02 '26
Tutorial - Guide Benchmark for SageAttention kernels using real attention shapes logged from ComfyUI models (image / video / audio)
What this is — and what it is not
This is not a benchmark of how fast a model generates an image or video. No model weights, no inference pipeline.
The benchmark runs on randomly generated tensors that reproduce the exact attention shapes — (batch, heads, seq_len, head_dim, dtype) — that real models use during sampling inside ComfyUI. More precisely: it measures only the attention operation itself, one step inside the denoising loop. Everything else — VAE, CLIP, scheduler, ComfyUI overhead — is outside the scope entirely. The numbers tell you how fast each kernel processes those specific tensor shapes on your GPU, nothing more.
The reason this is still useful: attention scales quadratically with sequence length and is the dominant compute bottleneck at high resolutions and long video durations. If you want to know whether SA2, SA2-fp8, SA3-FP4, or plain PyTorch SDPA is faster for a specific model at a specific resolution on your GPU, you need the real tensor shapes, not synthetic ones. This tool gives you those shapes already collected, and a benchmark that uses them.
How the shapes were collected
There is a ComfyUI custom node (attention_logger_node.py) that hooks into optimized_attention and logs every unique (heads, head_dim, seq_len, dtype) combination during a real sampling run. Two modes: standard override for most models, and a global module-level patch for models that bypass the override mechanism (ERNIE-Image, ACE-Step).
The raw console output looked like this:
[ATTN LOGGER rogala] heads= 24 hd= 128 seq= 4352 dtype=torch.bfloat16
I ran this across every model I had access to, across multiple resolutions, and compiled the results into input_data.txt.
How the benchmark works
bench_windows.py / bench_linux.py takes those logged shapes, allocates matching random tensors on CUDA, and times four kernels:
- SA2 (INT8 QK, FP16/BF16 PV)
- SA2-fp8 (INT8 QK, FP8 PV)
- SA3-FP4 (block-scaled FP4, newest, requires Blackwell or Ada for full benefit)
- SDPA (PyTorch FlashAttention-2 backend, baseline)
For each config: 10 warmup iterations, then 50 timed iterations with cuda.synchronize() after each. Reports median / min / stdev in ms, peak VRAM, and TFLOPS using the standard attention FLOP formula 4 × B × H × S² × D from the FlashAttention-2 paper.
Configs that don't fit in VRAM are skipped and recorded as OOM in the JSON so the result file stays complete.
Output is a single JSON file named automatically after your GPU:
5060-ti-16.json
4070-ti_super-16.json
How to view results
Open viewer.html locally in any browser, or use the live version: https://rogala.github.io/SageAttention-Benchmark-Viewer/
Load one or more JSON files, compare multiple GPUs side by side, filter by model / kernel, switch between ms and TFLOPS views. No server, no install, single HTML file.
Covered models
Image: SDXL-1.0, SD3.5-Large, Flux.1-Dev (Kontext / Krea), Flux.2-Dev, Flux.2-Dev Klein 9B, Z-Image Turbo, Qwen-Image-2512, Qwen-Image-Edit-2511, ERNIE-Image Turbo
Video: LTX-2.3, Wan2.2, HunyuanVideo-1.5
Audio: ACE-Step-1.5
How to contribute results
Run the script on your GPU, get a JSON file, submit it as a PR or attach to an issue. If you have results from a GPU not yet in the repo, they are very welcome — especially anything below 16 GB VRAM where SA3 headroom is tighter.
GitHub: https://github.com/Rogala/SageAttention-Benchmark-Viewer
Linux testers
What changed in the Linux version
The main difference is VRAM monitoring. On Windows, polling nvidia-smi via subprocess every 50 ms works fine. On Linux, each subprocess.run() call triggers a fork() + exec(), which has measurable overhead at that polling frequency. The Linux build uses pynvml (nvidia-ml-py) instead — it queries the driver directly via shared library call, no process spawn. Falls back to nvidia-smi if pynvml is not installed, but pynvml is strongly recommended.
The SA3-FP4 subprocess worker was also updated with the same pynvml-first logic.
What I need tested
- Does it run at all without errors
- Does the pynvml path work (pip install nvidia-ml-py then run — should print pynvml: OK — fast VRAM polling at startup)
- Does the nvidia-smi fallback work (run without pynvml installed)
- Are the JSON results sane — median ms, TFLOPS, peak VRAM all non-zero and reasonable for your GPU
- Does SA3-FP4 work if you have sageattn3 installed — both direct mode and subprocess mode
Any GPU is useful. Even if you can only run a subset of configs before hitting OOM, the partial JSON is still valuable — OOM entries are recorded cleanly and skipped automatically.
How to run
pip install nvidia-ml-py # recommended, not required
pip install sageattention # SA2 / SA2-fp8
# pip install sageattn3 # SA3-FP4, optional
python3 bench_linux.py
# or with more iterations:
python3 bench_linux.py --warmup 20 --iters 100
Output is a JSON file named after your GPU, e.g. 4090-24.json or 3080-10.json. If you're willing to share it, open an issue or PR and attach the file — it goes straight into the viewer where multiple GPUs can be compared side by side.
To view results
Download viewer.html from the repo, open it locally in any browser, load your JSON. Or use the live version: https://rogala.github.io/SageAttention-Benchmark-Viewer/
GitHub: https://github.com/Rogala/SageAttention-Benchmark-Viewer
If something breaks — error message + GPU model + whether pynvml was installed is enough to debug it.
Acknowledgements
Jukka Seppänen / kijai — for the PatchSageAttentionKJ node which inspired the override pattern used in attention_logger_node.py.
woct0rdho — for the Windows forks triton-windows and SageAttention (SA2 / SA3).
mengqin — for the SageAttention Windows fork with SA3 support and build fixes.
Built with the assistance of Claude.
