microsoft / microsoft/onnxruntime

Refactor contrib ops under bert/

Open
#27,581 0 comments 0 reactions 1 assignee Claimed by @titaiwangms View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.