Skip to content

Refactor all of Attention CUDA kernels to have separate kernels and shared functionality #27516

Description

@titaiwangms

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.

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

Metadata

Metadata

Labels

ep:CUDAissues related to the CUDA execution provider

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions