saymrwulf/onnxruntime: ONNX Runtime: cross-platform, high performance ML inferencing and training accelerator

mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00

ONNX Runtime: cross-platform, high performance ML inferencing and training accelerator

Find a file

Tianlei Wu c5f8389648 [CUDA] Fix MultiHeadAttention thread safe and bias support (#21498 ) ### Description #### Issues Fixed (1) TRT cross attention not thread safe. [Core changes like this](`6fd7aba3d4`) are used to make it thread-safe: * Add an once_flag to CumulatedSequenceLengthCache to make sure it is only initialized once; and change the cache to be read only after initialization. Previously, the content is not read-only so it might be changed by other thread and potentially cause buffer overrun. * The kernel initialization is not guarded (Although the factory of kernel loading has static mutex to guard multiple threading), so the mutable variable might be set by two different threads at the same time. Add an once_flag to avoid that. This requires need some workspace computation change as well. So I did not create a separated pull request. (2) Bias for cross attention That scenario has assumption that only query has bias, but not for key and value. However, such assumption is not verified in runtime and there was no comment of assumption, and there was no test case so the support of scenario was disabled by mistake. Actually, the scenario is used in whisper model (TODO: we shall add tests for whisper to CI pipeline, and also update fusion script to verify such assumptions if needed.) CUDA/CPU kernels supports bias for cross attention as long as bias is zero for key and value. I updated the check to support the scenario and added comments wherever there is such assumption. (3) Fallback support Previously, unfused kernel did not support packed qkv and packed kv formats. That means some case might fail since there is no fallback. I added new AddBiasTranpose cuda kernels for them to support fallback, so that all supported cases will not fail. #### Improvements (4) QKV workspace size. The logic for no_qkv_workspace could be easily out of sync since related code are scattered in different source files. I refactor the code to move all related code to one file (attention_prepare_qkv.cu) and add asserts, so that the logic can be in sync. (5) Remove confusing concept of pass past in kv parameters.pass_past_in_kv is confusing since the k/v in cross attention is not past state. Remove it and use parameters.qkv_format == Q_K_V_BSNH_BNSH_BNSH instead. New code does not use past_key/past_value for cross attention, so the logic is more clear. (6) More coverage and less workspace and less transpose of flash and efficient attention Previously, there is one condition does not run flash or efficient attention: ``` bool past_no_bias = (pass_key_value_as_past \|\| past_key != nullptr \|\| present_key != nullptr) && bias == nullptr; ``` After this change, we can use flash and efficient attention for the case, and also less workspace. For example, cross attention with bias, the original code uses two additional workspaces: ``` transpose: past_key (BxNxSxH) => temp_k_workspace (BxSxNxH), past_value (BxNxSxH_v) => temp_v_workspace (BxSxNxH_v) add bias: query => q, temp_k_workspace => k, temp_v_workspace => v ``` New logic is like ``` if (has bias) Add bias to query, key, value, and store in q, k, v workspace else Use query, key and value directly as q, k and v in kernel ``` We can see that, we do not need allocate temp_k_workspace and temp_v_workspace so use less memory. New code saved two transposes in this case. Flash and efficient attention supports BSNH or BNSH formats for k and v. In old code, k/v are also converted to BSNH format. Some is not necessary. I do some change to convert k/v to BSNH or BNSH case by case. So that there are more cases can be covered by flash or efficient attention to improve performance. (6) Debugging support Previously, there is less debug info. In this change, I add a flag for debug info in the AttentionData. So that we can output debug info during the processing. Also add functions to consolidate the dumping of inputs, QKV processing and outputs; Add an environment variable `ORT_ENABLE_GPU_DUMP` to allow disable dumping from cuda kernel. #### Summary of changes (1) Refactoring the CheckInputs, and pass in operator type. (2) Refactoring the PrepareQKV to support fallback for packed qkv or packed kv inputs. (3) Change a few case of PrepareQKV to allow more case covered by flash and efficient attention. (4) use parameters.qkv_format == Q_K_V_BSNH_BNSH_BNSH to replace parameters.pass_past_in_kv (5) Allow bias input for Q_K_V_BSNH_BNSH_BNSH, and add comments of assumption that key/value has no bias in this case. (6) Fix thread-safe issue in CumulatedSequenceLengthCache handling. (7) Add test cases to cover all supported scenarios. Current support scenarios for MultiHeadAttention for CUDA/CPU: \| Q \| K \| V \| pastK\| pastV \| presentK\| presentV \| Bias \| Op desc \| ---- \| ---- \| ---- \| ------ \| ----- \| --------- \| -------- \| -----\|--------- \| BSNH \| BLNH\| BLNH\| - \| - \| - \| - \| QKV \| not packed \| BLN3H\| - \| - \| - \| - \| - \| - \| QKV \| qkv packed <br> not support in CPU \| BSNH \| BLN2H\| - \| - \| - \| - \| - \| --- \| kv packed <br> not support in CPU \| BSNH \| BNLH\| BNLH\| - \| - \| - \| - \| Q-- \| cross attention <br> bias for Q only \| BSNH \| BLNH \| BLNH \| - \| - \| BNTH \| BNTH \| QKV \| no past <br> only present \| BSNH \| BLNH \| BLNH \| BNPH \| BNPH \| BNTH \| BNTH \| QKV \| past and present <br> (not share buffer) ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> https://github.com/microsoft/onnxruntime/issues/18854		2024-07-31 09:01:05 -07:00
.config	Update tsaoptions.json: update the email alias (#13448 )	2022-10-26 15:56:16 -07:00
.devcontainer	Remove two lines in the Dockerfile for Github Codespace (#12278 )	2022-07-21 20:52:17 -07:00
.gdn	Update win-ci-pipeline.yml: enable xnnpack tests (#16244 )	2023-06-14 19:12:42 -07:00
.github	Update labeling bot (#21548 )	2024-07-29 16:06:03 -07:00
.pipelines	Update DirectML from 1.14.1 to 1.15.0 (#21323 )	2024-07-22 16:59:03 -07:00
.vscode	disable gemm f16 on CPU (#19744 )	2024-03-01 13:44:29 -08:00
cgmanifests	[TensorRT EP] Update TRT OSS Parser to 10.2 (#21552 )	2024-07-29 17:27:38 -07:00
cmake	[TensorRT EP] Update TRT OSS Parser to 10.2 (#21552 )	2024-07-29 17:27:38 -07:00
csharp	Bump Sixlabors.ImageSharp from 2.1.8 to 2.1.9 in /csharp/sample/Microsoft.ML.OnnxRuntime.ResNet50v2Sample (#21444 )	2024-07-26 22:31:16 -07:00
dockerfiles	ORT- OVEP 1.19 PR-follow up (#21546 )	2024-07-29 14:12:36 -07:00
docs	Enable FP16 Clip and Handle Bias in FP16 Depthwise Conv (#21493 )	2024-07-30 03:49:14 -07:00
include/onnxruntime/core	Utilize ext data location to reduce qd matmul memory usage (#21451 )	2024-07-30 15:22:46 -07:00
java	Fix typos according to reviewdog report. (#21335 )	2024-07-22 13:37:32 -07:00
js	[WebNN EP] Support ConvTranspose for TFLite backend (#21291 )	2024-07-30 17:46:08 -07:00
objectivec	Fix Objective-C static analysis warnings. (#20417 )	2024-04-24 11:48:29 -07:00
onnxruntime	[CUDA] Fix MultiHeadAttention thread safe and bias support (#21498 )	2024-07-31 09:01:05 -07:00
orttraining	pick changes from https://github.com/onnx/onnx/pull/6195 to fix heap-buffer-overflow in onnx::convPoolShapeInference (#21507 )	2024-07-27 15:58:36 -07:00
rust	Fix typos according to reviewdog report. (#21335 )	2024-07-22 13:37:32 -07:00
samples	Removed all the deprecated python training code and related tests and utils (#18333 )	2023-11-17 18:19:21 -08:00
tools	CoreML: Add ML Program Split Op (#21456 )	2024-07-30 14:04:47 +10:00
winml	Update ruff and clang-format versions (#21479 )	2024-07-24 11:50:11 -07:00
.clang-format	Prevent GSL_SUPPRESS arguments from being modified by clang-format (#17242 )	2023-08-22 18:26:53 -07:00
.clang-tidy	Create clang-tidy CI (#12653 )	2022-09-30 08:05:38 -07:00
.dockerignore
.gitattributes	Fix typos according to reviewdog report. (#21335 )	2024-07-22 13:37:32 -07:00
.gitignore	Build onnxruntime.dll as arm64x (#18633 )	2023-12-06 16:49:00 -08:00
.gitmodules	[js/web] optimize module export and deployment (#20165 )	2024-05-20 09:51:16 -07:00
.lintrunner.toml	CoreML: Aggregated changes to add all required ops for priority model (#21472 )	2024-07-26 08:29:33 +10:00
build.bat	try to find patch.exe in git default installation folder (#17106 )	2023-08-10 21:48:13 -07:00
build.sh	Upgrade old Python version in packaging pipeline (#16667 )	2023-07-17 08:24:47 -07:00
build_arm64x.bat	remove unnecessary environment variable (#19166 )	2024-01-16 16:24:37 -08:00
CITATION.cff	Fix citation author name issue (#19597 )	2024-02-22 17:03:56 -08:00
CODEOWNERS	Add owners for public facing API files (#15288 )	2023-03-30 17:16:15 -07:00
CONTRIBUTING.md	Fix link to High Level Design (#11786 )	2023-02-28 11:05:54 -08:00
lgtm.yml	Fix lgtm C++ error (#13613 )	2022-11-10 10:06:22 -08:00
LICENSE
NuGet.config
ort.wprp	Fully dynamic ETW controlled logging for ORT and QNN logs (#20537 )	2024-06-06 21:11:14 -07:00
ORT_icon_for_light_bg.png
packages.config	Update DirectML from 1.14.1 to 1.15.0 (#21323 )	2024-07-22 16:59:03 -07:00
pyproject.toml	Ignore ruff rule `N813` (#21477 )	2024-07-24 17:48:22 -07:00
README.md	Update README.md (#18963 )	2024-01-03 17:26:25 -08:00
requirements-dev.txt	ONNX 1.15 integration (#17125 )	2023-09-26 14:44:48 -07:00
requirements-doc.txt
requirements-lintrunner.txt	Update ruff and clang-format versions (#21479 )	2024-07-24 11:50:11 -07:00
requirements-training.txt	ONNX 1.15 integration (#17125 )	2023-09-26 14:44:48 -07:00
requirements.txt	Add compatibility for NumPy 2.0 (#21085 )	2024-06-27 13:50:53 -07:00
SECURITY.md
setup.py	Migraphx ep windows build (#21284 )	2024-07-11 21:21:38 -07:00
ThirdPartyNotices.txt	Fix typos according to reviewdog report. (#21335 )	2024-07-22 13:37:32 -07:00
VERSION_NUMBER	Bump up version in main from 1.18.0 to 1.19.0 (#20489 )	2024-04-29 20:21:41 -07:00

README.md

ONNX Runtime is a cross-platform inference and training machine-learning accelerator.

ONNX Runtime inference can enable faster customer experiences and lower costs, supporting models from deep learning frameworks such as PyTorch and TensorFlow/Keras as well as classical machine learning libraries such as scikit-learn, LightGBM, XGBoost, etc. ONNX Runtime is compatible with different hardware, drivers, and operating systems, and provides optimal performance by leveraging hardware accelerators where applicable alongside graph optimizations and transforms. Learn more →

ONNX Runtime training can accelerate the model training time on multi-node NVIDIA GPUs for transformer models with a one-line addition for existing PyTorch training scripts. Learn more →

Get Started & Resources

General Information: onnxruntime.ai
Usage documentation and tutorials: onnxruntime.ai/docs
YouTube video tutorials: youtube.com/@ONNXRuntime
Upcoming Release Roadmap
Companion sample repositories:
- ONNX Runtime Inferencing: microsoft/onnxruntime-inference-examples
- ONNX Runtime Training: microsoft/onnxruntime-training-examples

Builtin Pipeline Status

System	Inference	Training
Windows
Linux
Mac
Android
iOS
Web
Other

Third-party Pipeline Status

System	Inference	Training
Linux

Data/Telemetry

Windows distributions of this project may collect usage data and send it to Microsoft to help improve our products and services. See the privacy statement for more details.

Contributions and Feedback

We welcome contributions! Please see the contribution guidelines.

For feature requests or bug reports, please file a GitHub Issue.

For general discussion or questions, please use GitHub Discussions.

Code of Conduct

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.

License

This project is licensed under the MIT License.