aciddelgado 406cd324e0 [CUDA] GroupQueryAttention operator using FlashAttention (#17674) ... ### Description Added Group Query Attention op, supporting integer multiple number of heads for Q / KV. As of now, this op can only use FlashAttention kernel, meaning it only supports sm>=80 on Linux. Results from onnxruntime/test/python/transformers/benchmark_gqa.py show an on-average ~37% speed-up over Decoder Masked Multi-Head Attention, with even greater improvements for long past sequence lengths. ``` op batch s_kv heads h_dim ms TFLOPS gqa 16 2048 8 32 0.34 0.10 dmmha 16 2048 8 32 0.39 0.09 --------- gqa 16 2048 8 64 0.45 0.15 dmmha 16 2048 8 64 0.61 0.11 --------- gqa 16 2048 8 128 0.54 0.25 dmmha 16 2048 8 128 0.83 0.16 --------- gqa 16 2048 16 32 0.45 0.15 dmmha 16 2048 16 32 0.69 0.10 --------- gqa 16 2048 16 64 0.69 0.19 dmmha 16 2048 16 64 0.83 0.16 --------- gqa 16 2048 16 128 0.71 0.38 dmmha 16 2048 16 128 1.28 0.21 --------- gqa 16 2048 32 32 0.58 0.23 dmmha 16 2048 32 32 0.77 0.17 --------- gqa 16 2048 32 64 0.58 0.46 dmmha 16 2048 32 64 1.25 0.21 --------- gqa 16 2048 32 128 0.76 0.71 dmmha 16 2048 32 128 2.15 0.25 --------- gqa 16 2048 64 32 0.68 0.39 dmmha 16 2048 64 32 1.23 0.22 --------- gqa 16 2048 64 64 0.77 0.70 dmmha 16 2048 64 64 2.11 0.25 --------- gqa 16 2048 64 128 1.10 0.97 dmmha 16 2048 64 128 4.06 0.26 --------- gqa 16 2048 128 32 1.00 0.54 dmmha 16 2048 128 32 2.09 0.26 --------- gqa 16 2048 128 64 1.10 0.97 dmmha 16 2048 128 64 4.08 0.26 ``` ### Motivation and Context As of now, this op is targeted for use on LLama models, as it supports kv-caching and different number of heads for Q and KV (Grouped Query Attention). We plan to add support for more platforms, input formats, etc. in the future. --------- Co-authored-by: Tianlei Wu <tlwu@microsoft.com> Co-authored-by: tlwu@microsoft.com <tlwu@a100.crj0ad2y1kku1j4yxl4sj10o4e.gx.internal.cloudapp.net>		2023-10-09 12:43:12 -07:00
..
github	[ROCm] Update ROCm and MIGraphX CI to ROCm5.7 (#17834)	2023-10-09 10:29:11 +08:00
__init__.py
amd_hipify.py
build.py	[CUDA] GroupQueryAttention operator using FlashAttention (#17674)	2023-10-09 12:43:12 -07:00
clean_docker_image_cache.py
compile_triton.py	[Better Engineering] Bump ruff to 0.0.278 and fix new lint errors (#16789)	2023-07-21 12:53:41 -07:00
coverage.py
gen_def.py	[TensorRT EP] Refactor OrtTensorRTProviderOptions initialization and make it easy to add new field (#17617)	2023-10-06 14:12:20 -07:00
get_docker_image.py	[Better Engineering] Bump ruff to 0.0.278 and fix new lint errors (#16789)	2023-07-21 12:53:41 -07:00
logger.py
op_registration_utils.py	[CI] Removes type2 in process_registration and fix Windows GPU Reduced Ops CI Pipeline (#16530)	2023-07-07 18:21:06 +02:00
op_registration_validator.py	[CI] Removes type2 in process_registration and fix Windows GPU Reduced Ops CI Pipeline (#16530)	2023-07-07 18:21:06 +02:00
patch_manylinux.py	[Better Engineering] Bump ruff to 0.0.278 and fix new lint errors (#16789)	2023-07-21 12:53:41 -07:00
policheck_exclusions.xml
reduce_op_kernels.py	Re-organize the transpose optimization and layout transformation files. (#16246)	2023-07-07 08:24:47 +10:00
replace_urls_in_deps.py
requirements.txt	Flash Attention v2 MHA (#17227)	2023-08-31 13:52:21 -07:00
set-trigger-rules.py	Pr trggiers generated by code (#17247)	2023-08-30 05:57:03 +08:00
update_tsaoptions.py
upload_python_package_to_azure_storage.py	[Linter] Bump ruff and remove pylint (#17797)	2023-10-05 21:07:33 -07:00