CUDA Cross-attention kernel

Is your feature request related to a problem? Please describe.

I'm able to use the onnxruntime.transformers codebase to optimize Tranformer-based model using self-attention, however it's not possible to use the self-attention kernel for cross-attention.

System information

• ONNX Runtime version (you are using): 1.10.0

Describe the solution you'd like

I would like to know if the implementation of a CUDA kernel for cross-attention is something you've considered adding to ONNXRuntime - or simply a modification of the current self-attention kernel to take in one input for queries and one input for keys and values.

Describe alternatives you've considered

For generative models I think the self-attention kernel can be used after a first pass, as we can simply reuse past keys and values. However that is not the case more generally, when you only perform one inference on a given pair of input.

Start in the onnxruntime.transformers codebase by tracing the existing self-attention kernel and how it handles queries, keys, values, and cached state. Done would mean defining and validating a CUDA cross-attention path that supports separate query and key/value inputs for one-shot inference.