You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
We need a better organization and clean up on contrib op MHA and GQA, and ONNX standard Attention: attention CUDA kernels to have a better readability and maintainability.
(1) ONNX Attention needs its own kernel: QkvToContext. It is now using GQA and MHA kernel, but they have different spec, so the implementation of ONNX attention is a bit awkward.
(2) Identify duplication code and functions and then make shared functions for them (MHA/GQA/ONNX Attention) if possible.
(3) The performance prefers flash attention to memory efficient attention and then unfused attention.
(4) ONNX Attention might need 2 kernels, since nonpad_kv_seqlen makes a big difference between 23 and 24 opset version. Also, with inplace kv cache possible in TensorScatter, two modes of Attention should be considered.
We need a better organization and clean up on contrib op MHA and GQA, and ONNX standard Attention: attention CUDA kernels to have a better readability and maintainability.
(1) ONNX Attention needs its own kernel:
QkvToContext. It is now using GQA and MHA kernel, but they have different spec, so the implementation of ONNX attention is a bit awkward.(2) Identify duplication code and functions and then make shared functions for them (MHA/GQA/ONNX Attention) if possible.
(3) The performance prefers flash attention to memory efficient attention and then unfused attention.
(4) ONNX Attention might need 2 kernels, since
nonpad_kv_seqlenmakes a big difference between 23 and 24 opset version. Also, with inplace kv cache possible in TensorScatter, two modes of Attention should be considered.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
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