Add Paged Attention Op for CUDA SM80 support #24595
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Dismissed
Show dismissed
Hide dismissed
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Dismissed
Show dismissed
Hide dismissed
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Dismissed
Show dismissed
Hide dismissed
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
onnxruntime/test/python/transformers/test_paged_attention_cuda.py
Outdated
Show resolved
Hide resolved
| } | ||
|
|
||
| ONNX_MS_OPERATOR_SET_SCHEMA( | ||
| PagedAttention, 1, |
There was a problem hiding this comment.
Should we name the input KV caches as past_key and past_value + the output KV caches as present_key and present_value to match the names used in the other attention ops?
There was a problem hiding this comment.
Since the kv cache inputs are always updated in place in this op, and the kv cache outputs are optional, I think it would be innacurate to name them like that
| "cumulative_sequence_length", | ||
| "A tensor with shape (batch_size + 1). It specifies the cumulative sequence lengths between the packed " | ||
| "entries in Q/K/V.", | ||
| "M") | ||
| .Input(6, | ||
| "seqlens", | ||
| "A tensor with shape (batch_size). It specifies the past lengths of cached sequence in the KV cache.", | ||
| "M") | ||
| .Input(7, | ||
| "max_query_len", | ||
| "Scalar tensor equivalent to the maximum length of all sequences in the query tensor (max new seqlen). " | ||
| "This input is allocated on CPU.", | ||
| "M") | ||
| .Input(8, | ||
| "max_seq_len", | ||
| "Scalar tensor equivalent to the maximum length of all sequences in the KV Cache after appending new " | ||
| "tokens (max total seqlen). This input is allocated on CPU.", | ||
| "M") |
There was a problem hiding this comment.
The interface is designed for the convenience of current implementation, that might not be good for other EP (like WebGPU, CPU etc).
A generic design is to have q_seq_lens and kv_seq_lens, then you can deduce these four inputs. For example, for CPU, it is very easy to compute max seq length, so there is no need to have max_query_len and max_seq_len inputs.
I guess that the extra latency caused by such interface change might be tiny.
There was a problem hiding this comment.
We take max query length and total sequence length as inputs because in CUDA we would need to copy the value from device to host within the op for however many layers are in the model. If you think this might be a tiny performance penalty then I'll change it; I agree that the amount of inputs is slightly bloated.
There was a problem hiding this comment.
Since padding has been removed in the input, the input is treated as a long sequence. My understanding is that there is no much use of max query length as long as you know the total length of the long sequence.
For max kv length, it is only used to check the optional input cos/sin cache of rotary. But, there is no guarantee that the max kv length itself is correct. To verify max length is correct, we will have to reduce from length per sequence.
If we do not check max lengths in CheckInputs. For security reason, we can let rotary kernel outputs NaNs when total length is too large, and add a check to GenAI to make sure total length is within the limit of rotary cache length.
There was a problem hiding this comment.
Primarily, max seqlen q and k were included as input because these values are passed to flash attention. However, I believe we can remove these dependencies at a kernel level.
Description
Adds Paged Attention Op which enables of Paged KV Cache. Inputs to this op are unpadded (packed / varlen) so Cumulative Sequence Lengths are a required input.
Motivation and Context
Adding this op to ONNXRuntime is necessary to allow the GenAI team to enable a continuous batching server API.