Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update benchmark_mha.py to compare with PyTorch SDPA #21449

Merged
merged 6 commits into from
Jul 27, 2024
Merged

Update benchmark_mha.py to compare with PyTorch SDPA #21449

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6209dcf
update benchmark_mha to compare with PyTorch SDPA
tianleiwu Jul 23, 2024
e4caf29
add comment
tianleiwu Jul 23, 2024
f3f95f2
remove unused variable
tianleiwu Jul 23, 2024
ecff18b
Add Windows batch file
tianleiwu Jul 23, 2024
85f9a13
minor change
tianleiwu Jul 23, 2024
8fc1dc3
Merge branch 'main' into tlwu/benchmark_torch_mha
tianleiwu Jul 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions onnxruntime/test/python/transformers/benchmark_mha.cmd
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
echo "Benchmark Scaled Dot Product Attention (SDPA) performance on GPU:"

set CUDA_VISIBLE_DEVICES=0
python benchmark_mha.py --use_gpu
python benchmark_mha.py --use_gpu --use_cuda_graph
python benchmark_mha.py --use_gpu --torch

type benchmark_mha_gpu_*.csv > mha_gpu_benchmark_results.csv

echo "Benchmark performance on CPU with number of threads:"
set MKL_DYNAMIC=FALSE
set OMP_NUM_THREADS=1
tianleiwu marked this conversation as resolved.
Show resolved Hide resolved
python benchmark_mha.py --torch

set OMP_NUM_THREADS=2
python benchmark_mha.py --torch

set OMP_NUM_THREADS=4
python benchmark_mha.py --torch

set OMP_NUM_THREADS=8
python benchmark_mha.py --torch

set MKL_DYNAMIC=
set OMP_NUM_THREADS=

set ORT_DISABLE_FLASH_ATTENTION=0
python benchmark_mha.py --intra_op_num_threads 1
python benchmark_mha.py --intra_op_num_threads 2
python benchmark_mha.py --intra_op_num_threads 4
python benchmark_mha.py --intra_op_num_threads 8

echo "Benchmark performance on CPU with default threads settings:"
python benchmark_mha.py

python benchmark_mha.py --torch

python benchmark_mha.py --causal
python benchmark_mha.py --torch --causal

python benchmark_mha.py --causal --has_past

set ORT_DISABLE_FLASH_ATTENTION=1
python benchmark_mha.py
set ORT_DISABLE_FLASH_ATTENTION=

type benchmark_mha_cpu_*.csv > mha_cpu_benchmark_results.csv
Loading
Loading