microsoft / microsoft/onnxruntime
Refactor contrib ops under bert/
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
We need a better organization and clean up on LLM contrib op to have a better readability and maintainability.
(1) Identify duplication code and functions and then make shared functions for them if possible.
(2) Regarding how the attention files are organized, some are in the bert folder while others are in the llm folder. I think it would be good to have an attn folder and move all of the remaining bert files into llm. The bert name is a past relic and should be modernized.
---
The main folder of Contrib Op LLM Attention: https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/contrib_ops/cuda/bert
Some important files beginning of tracing:
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/attention_prepare_qkv.cu
Contrib Op specs: https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MultiHeadAttention
https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.GroupQueryAttention
The core of the kernel (flash attention):
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc
---
NOTE: ONNX Attention also consumes contrib op utilities. If there is any path changes, please remember to update it-
The ONNX Attention:
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/llm/attention.cc
The ONNX Attention Spec:
https://github.com/onnx/onnx/blob/main/docs/Operators.md#Attention
Contributor guide
Assessment
This issue has not been assessed yet.