Attention#
Domain:
ai.onnxSince version: 24
Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed.
This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V.
For self attention, kv_sequence_length equals to q_sequence_length.
For cross attention, query and key might have different lengths.
This operator also covers the 3 following variants based on the number of heads:
Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762,
q_num_heads = kv_num_heads.Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245,
q_num_heads > kv_num_heads,q_num_heads % kv_num_heads == 0.Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150,
q_num_heads > kv_num_heads,kv_num_heads=1.
Attention bias to be added is calculated based on attn_mask input and is_causal attribute:
attn_mask: A boolean mask where a value ofTrueindicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.If
is_causalis set to1, attention scores above the causal frontier are masked out. For internal cache (past_key) this is the standard offset frompast_sequence_length; for external cache (nonpad_kv_seqlenwithoutpast_key) this is bottom-right aligned bynonpad_kv_seqlen - q_sequence_length.If both
attn_maskandis_causalare set, the valid positions are the intersection of both masks.
If a query row is fully masked after this intersection, its output row is zero.
With respect to KV cache update, this operator allows the following two use cases:
Cache update happens inside the Attention operator. In this case, the
KandVinputs contain only the incoming
tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present key and value, respectively. Note that this only works correctly for the special case where the past key and value do not contain padded tokens.
Cache update happens outside the Attention operator (for example, through the
TensorScatteroperator). In this
case, the K and V inputs correspond to the entire cache tensor, so the four optional inputs/outputs past and
present key and value should not be used. An additional input nonpad_kv_seqlen of shape (batch_size,) may be
provided to indicate the number of non-padding tokens in each sample of the batch to save unnecessary computation.
Here, the kv_sequence dimension of attn_mask can be shorter than K and V, but still needs to be at least as long
as the maximum value of nonpad_kv_seqlen.
Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided:
The following pattern is applied by this operator:
Q K V
| | |
Q*sqrt(scale) K*sqrt(scale) |
| | |
| Transpose |
| | |
---MatMul--- |
| |
softcap (if provided) |
| |
at_mask---Add |
| |
Softmax |
| |
-----MatMul------
|
Y
Inputs
Q (T1): Query tensor. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, head_size)or 3D tensor with shape(batch_size, q_sequence_length, q_hidden_size). For cases with a 3D input tensor,q_hidden_size = q_num_heads * head_sizeK (T1): Key tensor. 4D tensor with shape
(batch_size, kv_num_heads, kv_sequence_length, head_size)or 3D tensor with shape(batch_size, kv_sequence_length, k_hidden_size). For cases with a 3D input tensor,k_hidden_size = kv_num_heads * head_sizeV (T2): Value tensor. 4D tensor with shape
(batch_size, kv_num_heads, kv_sequence_length, v_head_size)or 3D tensor with shape(batch_size, kv_sequence_length, v_hidden_size). For cases with a 3D input tensor,v_hidden_size = kv_num_heads * v_head_sizeattn_mask (U): Attention mask. Shape must be broadcastable to
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.The last dimension can also be shorter thantotal_sequence_lengthand will be padded tototal_sequence_lengthwith negative infinity. Two types of masks are supported: a boolean mask where a value ofTrueindicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score.past_key (T1): past state cache for key with shape
(batch_size, kv_num_heads, past_sequence_length, head_size)past_value (T2): past state cache for value with shape
(batch_size, kv_num_heads, past_sequence_length, v_head_size)nonpad_kv_seqlen (tensor(int64)): A vector of integers of shape
(batch_size,)that indicates the number of valid (ie, non-padding) tokens in each sample. A padding mask can be derived from this. This should not be used together withpast_keyandpast_valueinputs orpresent_keyandpresent_valueoutputs (See the KV cache use cases in the operator description).
Outputs
Y (T1): The output tensor . 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, v_head_size)or 3D tensor with shape(batch_size, q_sequence_length, hidden_size). For cases with a 3D input tensor,hidden_size = q_num_heads * v_head_sizepresent_key (T1): Updated key cache with shape
(batch_size, kv_num_heads, total_sequence_length, head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.present_value (T2): Updated value cache with shape
(batch_size, kv_num_heads, total_sequence_length, v_head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.qk_matmul_output (T1): The output of QK matmul. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.
Type Constraints
T1: Constrain Q and K inputs types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).
T2: Constrain V input types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).
U: Constrain output ‘mask’ types to boolean tensors and input types. Allowed types: tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8).
Examples#
test_cc_attention_24_fullymasked_qk_matmul_output_mode3_zero
Node:
Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output)
Attributes:
qk_matmul_output_mode = 3
Inputs:
Q: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1. , 0. ],
[ 0. , 1. ]],
[[ 0.5, 0.5],
[ 1. , -1. ]]]]
K: shape=(1, 2, 3, 2), dtype=float32
[[[[ 1. , 0. ],
[ 0.5 , 0.5 ],
[ 0. , 1. ]],
[[-1. , 1. ],
[ 1. , 1. ],
[ 0.25, -0.5 ]]]]
V: shape=(1, 2, 3, 2), dtype=float32
[[[[ 1. , 0. ],
[ 0. , 1. ],
[-1. , 1. ]],
[[ 2. , -2. ],
[ 0.5 , 0.25],
[-0.5 , 0. ]]]]
attn_mask: shape=(1, 2, 2, 3), dtype=bool
[[[[False, False, False],
[False, False, False]],
[[False, False, False],
[False, False, False]]]]
Outputs:
Y: shape=(1, 2, 2, 2), dtype=float32
[[[[0., 0.],
[0., 0.]],
[[0., 0.],
[0., 0.]]]]
qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
[[[[nan, nan, nan],
[nan, nan, nan]],
[[nan, nan, nan],
[nan, nan, nan]]]]
test_cc_attention_24_qk_matmul_output_mode3_softmax_precision
Node:
Attention(Q, K, V) -> (Y, "", "", qk_matmul_output)
Attributes:
qk_matmul_output_mode = 3
Inputs:
Q: shape=(1, 2, 2, 4), dtype=float32
[[[[0.47307658, 0.23097962, 0.78772914, 0.28697807],
[0.5986557 , 0.8135814 , 0.02500451, 0.0589357 ]],
[[0.38251805, 0.2956621 , 0.3988651 , 0.37017328],
[0.10475397, 0.03714603, 0.93207705, 0.98540026]]]]
K: shape=(1, 2, 3, 4), dtype=float32
[[[[0.57229 , 0.4511435 , 0.07304686, 0.76316774],
[0.21401769, 0.22756338, 0.37048477, 0.5940939 ],
[0.6135922 , 0.00247645, 0.5526311 , 0.68208873]],
[[0.18123084, 0.7024574 , 0.383707 , 0.30577767],
[0.00781494, 0.6027138 , 0.66016024, 0.6627315 ],
[0.6305149 , 0.28021097, 0.60129017, 0.9475914 ]]]]
V: shape=(1, 2, 3, 4), dtype=float32
[[[[0.6715034 , 0.6713074 , 0.35836458, 0.23935741],
[0.8293797 , 0.64154536, 0.71596503, 0.12925214],
[0.84466636, 0.7092908 , 0.70639706, 0.9940042 ]],
[[0.25770772, 0.36776882, 0.8353369 , 0.6261551 ],
[0.21529329, 0.49399358, 0.09771872, 0.67859787],
[0.54878294, 0.3643933 , 0.18085933, 0.3052022 ]]]]
Outputs:
Y: shape=(1, 2, 2, 4), dtype=float32
[[[[0.78449166, 0.67586154, 0.59811294, 0.4810154 ],
[0.7750107 , 0.6742243 , 0.57878083, 0.44518116]],
[[0.3506411 , 0.407938 , 0.35311505, 0.52554405],
[0.35865024, 0.4100264 , 0.32065555, 0.51681125]]]]
qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
[[[[0.31976908, 0.31416643, 0.3660645 ],
[0.37512714, 0.30729562, 0.3175772 ]],
[[0.30486938, 0.32805195, 0.36707866],
[0.25747713, 0.34540045, 0.3971224 ]]]]
test_cc_attention_4d_causal_nonpad_attn_mask_composition
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[0.24988627, 0.38419312],
[0.49044216, 0.66443473]],
[[0.65402985, 0.4631328 ],
[0.5551101 , 0.06749588]]],
[[[0.0172379 , 0.35170764],
[0.5883945 , 0.8262432 ]],
[[0.9321008 , 0.47234195],
[0.42556155, 0.26788622]]]]
K: shape=(2, 2, 4, 2), dtype=float32
[[[[0.3490997 , 0.604357 ],
[0.7757599 , 0.1406244 ],
[0.26939183, 0.8771148 ],
[0.90059036, 0.6026541 ]],
[[0.24831206, 0.05852199],
[0.7421605 , 0.13815868],
[0.00857764, 0.13765335],
[0.8771915 , 0.58826363]]],
[[[0.95356995, 0.4522164 ],
[0.91463107, 0.71303976],
[0.9329594 , 0.6416764 ],
[0.88627857, 0.26962817]],
[[0.2583304 , 0.8503816 ],
[0.88740987, 0.3403653 ],
[0.96880656, 0.1806879 ],
[0.9674024 , 0.91090786]]]]
V: shape=(2, 2, 4, 2), dtype=float32
[[[[0.44831312, 0.8245209 ],
[0.06107759, 0.6168141 ],
[0.8847538 , 0.29109675],
[0.24607062, 0.13781232]],
[[0.4793862 , 0.76533633],
[0.8959265 , 0.45007414],
[0.08505452, 0.80296475],
[0.32882142, 0.90864104]]],
[[[0.1610483 , 0.34349614],
[0.35218954, 0.7289061 ],
[0.8512274 , 0.72585875],
[0.46584773, 0.6272389 ]],
[[0.31214702, 0.95401984],
[0.3151207 , 0.71161073],
[0.29118264, 0.21261078],
[0.39396596, 0.62221926]]]]
attn_mask: shape=(2, 1, 2, 4), dtype=float32
[[[[ 0. , 0. , -2. , -4. ],
[ 0. , 0. , -1. , -2. ]]],
[[[ 0. , -0.5, -2. , -4. ],
[ 0. , -0.5, -1.5, -3. ]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[4, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[0.30248764, 0.6936386 ],
[0.36174923, 0.6155303 ]],
[[0.6816254 , 0.5988387 ],
[0.6068539 , 0.63868487]]],
[[[0.23612623, 0.49488044],
[0.31726953, 0.5300721 ]],
[[0.3134448 , 0.84822863],
[0.3105105 , 0.7749761 ]]]]
test_cc_attention_4d_causal_nonpad_batch_prefill
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[0.24988627, 0.38419312],
[0.49044216, 0.66443473]],
[[0.65402985, 0.4631328 ],
[0.5551101 , 0.06749588]]],
[[[0.0172379 , 0.35170764],
[0.5883945 , 0.8262432 ]],
[[0.9321008 , 0.47234195],
[0.42556155, 0.26788622]]]]
K: shape=(2, 2, 4, 2), dtype=float32
[[[[0.3490997 , 0.604357 ],
[0.7757599 , 0.1406244 ],
[0.26939183, 0.8771148 ],
[0.90059036, 0.6026541 ]],
[[0.24831206, 0.05852199],
[0.7421605 , 0.13815868],
[0.00857764, 0.13765335],
[0.8771915 , 0.58826363]]],
[[[0.95356995, 0.4522164 ],
[0.91463107, 0.71303976],
[0.9329594 , 0.6416764 ],
[0.88627857, 0.26962817]],
[[0.2583304 , 0.8503816 ],
[0.88740987, 0.3403653 ],
[0.96880656, 0.1806879 ],
[0.9674024 , 0.91090786]]]]
V: shape=(2, 2, 4, 2), dtype=float32
[[[[0.44831312, 0.8245209 ],
[0.06107759, 0.6168141 ],
[0.8847538 , 0.29109675],
[0.24607062, 0.13781232]],
[[0.4793862 , 0.76533633],
[0.8959265 , 0.45007414],
[0.08505452, 0.80296475],
[0.32882142, 0.90864104]]],
[[[0.1610483 , 0.34349614],
[0.35218954, 0.7289061 ],
[0.8512274 , 0.72585875],
[0.46584773, 0.6272389 ]],
[[0.31214702, 0.95401984],
[0.3151207 , 0.71161073],
[0.29118264, 0.21261078],
[0.39396596, 0.62221926]]]]
attn_mask: shape=(2, 1, 2, 4), dtype=float32
[[[[ 0. , 0. , -2. , -4. ],
[ 0. , 0. , -1. , -2. ]]],
[[[ 0. , -0.5, -2. , -4. ],
[ 0. , -0.5, -1.5, -3. ]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[4, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[0.30248764, 0.6936386 ],
[0.36174923, 0.6155303 ]],
[[0.6816254 , 0.5988387 ],
[0.6068539 , 0.63868487]]],
[[[0.23612623, 0.49488044],
[0.31726953, 0.5300721 ]],
[[0.3134448 , 0.84822863],
[0.3105105 , 0.7749761 ]]]]
test_cc_attention_4d_causal_nonpad_continued_prefill
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[0.24988627, 0.38419312],
[0.49044216, 0.66443473]],
[[0.65402985, 0.4631328 ],
[0.5551101 , 0.06749588]]],
[[[0.0172379 , 0.35170764],
[0.5883945 , 0.8262432 ]],
[[0.9321008 , 0.47234195],
[0.42556155, 0.26788622]]]]
K: shape=(2, 2, 4, 2), dtype=float32
[[[[0.3490997 , 0.604357 ],
[0.7757599 , 0.1406244 ],
[0.26939183, 0.8771148 ],
[0.90059036, 0.6026541 ]],
[[0.24831206, 0.05852199],
[0.7421605 , 0.13815868],
[0.00857764, 0.13765335],
[0.8771915 , 0.58826363]]],
[[[0.95356995, 0.4522164 ],
[0.91463107, 0.71303976],
[0.9329594 , 0.6416764 ],
[0.88627857, 0.26962817]],
[[0.2583304 , 0.8503816 ],
[0.88740987, 0.3403653 ],
[0.96880656, 0.1806879 ],
[0.9674024 , 0.91090786]]]]
V: shape=(2, 2, 4, 2), dtype=float32
[[[[0.44831312, 0.8245209 ],
[0.06107759, 0.6168141 ],
[0.8847538 , 0.29109675],
[0.24607062, 0.13781232]],
[[0.4793862 , 0.76533633],
[0.8959265 , 0.45007414],
[0.08505452, 0.80296475],
[0.32882142, 0.90864104]]],
[[[0.1610483 , 0.34349614],
[0.35218954, 0.7289061 ],
[0.8512274 , 0.72585875],
[0.46584773, 0.6272389 ]],
[[0.31214702, 0.95401984],
[0.3151207 , 0.71161073],
[0.29118264, 0.21261078],
[0.39396596, 0.62221926]]]]
attn_mask: shape=(2, 1, 2, 4), dtype=float32
[[[[ 0. , 0. , -2. , -4. ],
[ 0. , 0. , -1. , -2. ]]],
[[[ 0. , -0.5, -2. , -4. ],
[ 0. , -0.5, -1.5, -3. ]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[4, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[0.30248764, 0.6936386 ],
[0.36174923, 0.6155303 ]],
[[0.6816254 , 0.5988387 ],
[0.6068539 , 0.63868487]]],
[[[0.23612623, 0.49488044],
[0.31726953, 0.5300721 ]],
[[0.3134448 , 0.84822863],
[0.3105105 , 0.7749761 ]]]]
test_cc_attention_4d_causal_nonpad_kv_continued_prefill
Node:
Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[0.6962669 , 0.07776612],
[0.08501613, 0.9095214 ]],
[[0.54328156, 0.16403002],
[0.49489892, 0.05037552]]],
[[[0.7477982 , 0.23961657],
[0.20933568, 0.9141033 ]],
[[0.27740717, 0.6019501 ],
[0.43859255, 0.7029143 ]]]]
K: shape=(2, 2, 4, 2), dtype=float32
[[[[0.7954803 , 0.29793 ],
[0.37033385, 0.38571107],
[0.15864354, 0.578012 ],
[0.8403792 , 0.58553374]],
[[0.97887236, 0.9464309 ],
[0.36310166, 0.22601879],
[0.35388404, 0.2672615 ],
[0.8902225 , 0.02329171]]],
[[[0.06205994, 0.75321126],
[0.40568942, 0.6124233 ],
[0.32807046, 0.9187455 ],
[0.31630176, 0.6255547 ]],
[[0.8385001 , 0.7832124 ],
[0.0923354 , 0.02002227],
[0.00747234, 0.32396793],
[0.15816778, 0.6564828 ]]]]
V: shape=(2, 2, 4, 2), dtype=float32
[[[[0.89469373, 0.5180939 ],
[0.65565157, 0.86190075],
[0.77400553, 0.99199396],
[0.18585944, 0.12069196]],
[[0.20994651, 0.6532453 ],
[0.51686764, 0.53793424],
[0.4303609 , 0.9325729 ],
[0.34185243, 0.34366912]]],
[[[0.26953828, 0.644491 ],
[0.8432479 , 0.62828964],
[0.24633849, 0.00292784],
[0.8958709 , 0.98316544]],
[[0.8923167 , 0.88685066],
[0.52004623, 0.39126772],
[0.3298484 , 0.3558908 ],
[0.58473134, 0.36779422]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[4, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[0.78347737, 0.7652148 ],
[0.6120686 , 0.6140527 ]],
[[0.36611217, 0.7018949 ],
[0.36323944, 0.60708874]]],
[[[0.57898587, 0.6357523 ],
[0.43766996, 0.40132707]],
[[0.7492422 , 0.6963835 ],
[0.6383399 , 0.61082983]]]]
test_cc_attention_4d_causal_nonpad_negative_offset_structural_empty
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[0.24988627, 0.38419312],
[0.49044216, 0.66443473]],
[[0.65402985, 0.4631328 ],
[0.5551101 , 0.06749588]]],
[[[0.0172379 , 0.35170764],
[0.5883945 , 0.8262432 ]],
[[0.9321008 , 0.47234195],
[0.42556155, 0.26788622]]]]
K: shape=(2, 2, 4, 2), dtype=float32
[[[[0.3490997 , 0.604357 ],
[0.7757599 , 0.1406244 ],
[0.26939183, 0.8771148 ],
[0.90059036, 0.6026541 ]],
[[0.24831206, 0.05852199],
[0.7421605 , 0.13815868],
[0.00857764, 0.13765335],
[0.8771915 , 0.58826363]]],
[[[0.95356995, 0.4522164 ],
[0.91463107, 0.71303976],
[0.9329594 , 0.6416764 ],
[0.88627857, 0.26962817]],
[[0.2583304 , 0.8503816 ],
[0.88740987, 0.3403653 ],
[0.96880656, 0.1806879 ],
[0.9674024 , 0.91090786]]]]
V: shape=(2, 2, 4, 2), dtype=float32
[[[[0.44831312, 0.8245209 ],
[0.06107759, 0.6168141 ],
[0.8847538 , 0.29109675],
[0.24607062, 0.13781232]],
[[0.4793862 , 0.76533633],
[0.8959265 , 0.45007414],
[0.08505452, 0.80296475],
[0.32882142, 0.90864104]]],
[[[0.1610483 , 0.34349614],
[0.35218954, 0.7289061 ],
[0.8512274 , 0.72585875],
[0.46584773, 0.6272389 ]],
[[0.31214702, 0.95401984],
[0.3151207 , 0.71161073],
[0.29118264, 0.21261078],
[0.39396596, 0.62221926]]]]
attn_mask: shape=(2, 1, 2, 4), dtype=float32
[[[[ 0. , 0. , -2. , -4. ],
[ 0. , 0. , -1. , -2. ]]],
[[[ 0. , -0.5, -2. , -4. ],
[ 0. , -0.5, -1.5, -3. ]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[4, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[0.30248764, 0.6936386 ],
[0.36174923, 0.6155303 ]],
[[0.6816254 , 0.5988387 ],
[0.6068539 , 0.63868487]]],
[[[0.23612623, 0.49488044],
[0.31726953, 0.5300721 ]],
[[0.3134448 , 0.84822863],
[0.3105105 , 0.7749761 ]]]]
test_cc_attention_4d_diff_heads_mask4d_nonpad_kv
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
scale = 0.5
Inputs:
Q: shape=(2, 2, 2, 2), dtype=float32
[[[[ 0.1 , 0.2 ],
[ 0.3 , 0.4 ]],
[[-0.1 , 0.05],
[ 0.2 , -0.3 ]]],
[[[ 0.5 , 0.5 ],
[ 0. , 1. ]],
[[ 1. , 0. ],
[ 0.5 , -0.5 ]]]]
K: shape=(2, 2, 3, 2), dtype=float32
[[[[ 0.5 , -0.5 ],
[ 0. , 0.5 ],
[ 1. , 0. ]],
[[-0.5 , 1. ],
[ 0.25, -0.25],
[ 0.5 , 0.5 ]]],
[[[ 1. , 1. ],
[-1. , 0. ],
[ 0. , -1. ]],
[[ 0.5 , 0.5 ],
[ 0.1 , 0.2 ],
[ 0.3 , 0.4 ]]]]
V: shape=(2, 2, 3, 2), dtype=float32
[[[[ 0.5, 0.5],
[-1. , 0. ],
[ 0. , 0.5]],
[[ 0.5, -0.5],
[ 1. , 1. ],
[-0.5, 0.5]]],
[[[ 0.2, -0.2],
[ 0.4, 0.6],
[-0.1, 0.3]],
[[ 0. , 1. ],
[ 1. , 0. ],
[-1. , -1. ]]]]
attn_mask: shape=(2, 2, 2, 3), dtype=float32
[[[[ 0. , -0.5, -1. ],
[ 0.2, 0. , -0.2]],
[[ 0.5, 0. , -0.1],
[ 0. , -0.3, 0.1]]],
[[[ 0. , 0. , 0. ],
[-0.4, -0.2, 0. ]],
[[ 0.1, -0.1, 0.2],
[ 0. , 0. , -0.5]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[2, 3]
Outputs:
Y: shape=(2, 2, 2, 2), dtype=float32
[[[[-9.29813758e-02, 3.02339554e-01],
[-2.21888170e-01, 2.59370595e-01]],
[[ 6.80762649e-01, 4.22879569e-02],
[ 7.45313048e-01, 2.35939145e-01]]],
[[[ 1.75710469e-01, 1.15763910e-01],
[ 1.92802399e-01, 1.78690210e-01]],
[[-1.23460844e-01, -2.79014523e-09],
[ 1.49503186e-01, 1.59121960e-01]]]]
test_cc_attention_4d_gqa_causal_nonpad_decode
Node:
Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(1, 4, 2, 2), dtype=float32
[[[[ 0.1 , 0.2 ],
[ 0.3 , 0.4 ]],
[[-0.1 , 0.05],
[ 0.2 , -0.3 ]],
[[ 0.5 , 0.5 ],
[ 0. , 1. ]],
[[ 1. , 0. ],
[ 0.5 , -0.5 ]]]]
K: shape=(1, 2, 3, 2), dtype=float32
[[[[ 1. , 0. ],
[ 0.5 , 0.5 ],
[ 0. , 1. ]],
[[-1. , 1. ],
[ 1. , 1. ],
[ 0.25, -0.5 ]]]]
V: shape=(1, 2, 3, 2), dtype=float32
[[[[ 1. , 0. ],
[ 0. , 1. ],
[-1. , 1. ]],
[[ 2. , -2. ],
[ 0.5 , 0.25],
[-0.5 , 0. ]]]]
nonpad_kv_seqlen: shape=(1,), dtype=int64
[3]
Outputs:
Y: shape=(1, 4, 2, 2), dtype=float32
[[[[ 0.4911621 , 0.50883794],
[-0.02356531, 0.6783799 ]],
[[ 0.48674485, 0.5132551 ],
[ 0.11724145, 0.6063233 ]],
[[ 0.9953577 , -0.4930365 ],
[ 0.9917567 , -0.74587834]],
[[ 0.79335546, -0.19003323],
[ 0.29831943, -0.26321504]]]]
test_cc_attention_4d_gqa_causal_nonpad_decode_fp16
Node:
Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
Inputs:
Q: shape=(2, 9, 4, 8), dtype=float16
[[[[0.0945 , 0.7554 , 0.3372 , ..., 0.2251 , 0.09625 , 0.5435 ],
[0.3733 , 0.892 , 0.391 , ..., 0.666 , 0.8926 , 0.8906 ],
[0.4827 , 0.1914 , 0.3293 , ..., 0.4897 , 0.82 , 0.8774 ],
[0.995 , 0.8486 , 0.4473 , ..., 0.9634 , 0.0887 , 0.977 ]],
[[0.8115 , 0.326 , 0.4934 , ..., 0.8677 , 0.611 , 0.8037 ],
[0.661 , 0.8496 , 0.1693 , ..., 0.001018, 0.5024 , 0.6113 ],
[0.6006 , 0.9126 , 0.63 , ..., 0.676 , 0.5825 , 0.7954 ],
[0.5156 , 0.7734 , 0.983 , ..., 0.2205 , 0.8726 , 0.3906 ]],
[[0.6196 , 0.9556 , 0.7896 , ..., 0.4795 , 0.0976 , 0.574 ],
[0.9536 , 0.7485 , 0.564 , ..., 0.313 , 0.1863 , 0.1627 ],
[0.8877 , 0.708 , 0.2732 , ..., 0.0961 , 0.77 , 0.11365 ],
[0.6196 , 0.1549 , 0.838 , ..., 0.2201 , 0.6826 , 0.02515 ]],
...,
[[0.8877 , 0.687 , 0.4976 , ..., 0.2415 , 0.888 , 0.3813 ],
[0.3013 , 0.3867 , 0.4314 , ..., 0.954 , 0.781 , 0.2961 ],
[0.854 , 0.2605 , 0.4016 , ..., 0.749 , 0.895 , 0.01886 ],
[0.993 , 0.8794 , 0.579 , ..., 0.2678 , 0.315 , 0.3975 ]],
[[0.4949 , 0.07935 , 0.587 , ..., 0.667 , 0.828 , 0.4207 ],
[0.9634 , 0.4626 , 0.001799, ..., 0.1188 , 0.4258 , 0.718 ],
[0.3308 , 0.8975 , 0.854 , ..., 0.3113 , 0.8013 , 0.05927 ],
[0.357 , 0.01604 , 0.1256 , ..., 0.948 , 0.852 , 0.8555 ]],
[[0.1307 , 0.7183 , 0.2534 , ..., 0.2489 , 0.1777 , 0.657 ],
[0.5684 , 0.668 , 0.6265 , ..., 0.4478 , 0.5815 , 0.657 ],
[0.9146 , 0.796 , 0.0865 , ..., 0.2957 , 0.5703 , 0.1877 ],
[0.2417 , 0.7974 , 0.4294 , ..., 0.5586 , 0.853 , 0.2217 ]]],
[[[0.951 , 0.6343 , 0.904 , ..., 0.5186 , 0.843 , 0.873 ],
[0.272 , 0.0337 , 0.711 , ..., 0.4724 , 0.1549 , 0.896 ],
[0.7607 , 0.988 , 0.506 , ..., 0.2335 , 0.1082 , 0.1857 ],
[0.803 , 0.2551 , 0.8965 , ..., 0.6304 , 0.2249 , 0.2778 ]],
[[0.1129 , 0.86 , 0.944 , ..., 0.007214, 0.73 , 0.8486 ],
[0.2307 , 0.591 , 0.662 , ..., 0.724 , 0.05148 , 0.2148 ],
[0.02669 , 0.5044 , 0.5186 , ..., 0.656 , 0.321 , 0.8345 ],
[0.1984 , 0.42 , 0.933 , ..., 0.6943 , 0.873 , 0.8047 ]],
[[0.772 , 0.425 , 0.7803 , ..., 0.2461 , 0.745 , 0.3655 ],
[0.601 , 0.371 , 0.8857 , ..., 0.7334 , 0.1779 , 0.3953 ],
[0.8677 , 0.3762 , 0.5303 , ..., 0.0942 , 0.1152 , 0.9146 ],
[0.583 , 0.3235 , 0.946 , ..., 0.2822 , 0.704 , 0.584 ]],
...,
[[0.506 , 0.712 , 0.2119 , ..., 0.3293 , 0.2125 , 0.4702 ],
[0.3206 , 0.346 , 0.6333 , ..., 0.977 , 0.1055 , 0.358 ],
[0.11206 , 0.04666 , 0.695 , ..., 0.631 , 0.2296 , 0.1803 ],
[0.1407 , 0.9473 , 0.8794 , ..., 0.4949 , 0.982 , 0.848 ]],
[[0.995 , 0.4465 , 0.12274 , ..., 0.288 , 0.4624 , 0.412 ],
[0.0913 , 0.7095 , 0.3145 , ..., 0.7456 , 0.974 , 0.565 ],
[0.1742 , 0.166 , 0.797 , ..., 0.1173 , 0.524 , 0.1388 ],
[0.8164 , 0.01196 , 0.8647 , ..., 0.67 , 0.321 , 0.607 ]],
[[0.919 , 0.7085 , 0.267 , ..., 0.1848 , 0.2778 , 0.583 ],
[0.2109 , 0.4834 , 0.7056 , ..., 0.4595 , 0.5103 , 0.321 ],
[0.7495 , 0.8286 , 0.3352 , ..., 0.807 , 0.8374 , 0.2168 ],
[0.4192 , 0.002748, 0.2634 , ..., 0.556 , 0.2805 , 0.3054 ]]]]
K: shape=(2, 3, 6, 8), dtype=float16
[[[[8.8623e-01, 7.0361e-01, 8.4521e-01, ..., 2.0764e-01, 7.2363e-01,
8.4766e-01],
[6.6455e-01, 3.1860e-01, 8.8867e-01, ..., 3.4546e-01, 9.4727e-01,
5.0342e-01],
[5.1367e-01, 6.5186e-01, 7.2949e-01, ..., 4.2432e-01, 2.4097e-01,
9.3213e-01],
[6.8115e-01, 6.8457e-01, 8.6914e-01, ..., 5.5811e-01, 8.3301e-01,
8.6084e-01],
[3.6938e-01, 7.8760e-01, 8.8867e-01, ..., 5.0098e-01, 2.8125e-01,
2.8174e-01],
[2.3267e-01, 2.2388e-01, 9.4727e-01, ..., 1.4807e-01, 9.7021e-01,
9.5312e-01]],
[[8.7939e-01, 6.4551e-01, 2.8125e-01, ..., 3.5065e-02, 6.1377e-01,
8.9990e-01],
[3.6987e-01, 5.9668e-01, 7.0557e-01, ..., 8.3838e-01, 7.5293e-01,
9.1016e-01],
[7.1875e-01, 1.0910e-02, 1.1322e-01, ..., 8.7451e-01, 2.5781e-01,
4.0222e-02],
[3.9355e-01, 7.1338e-01, 9.7559e-01, ..., 6.0107e-01, 8.2617e-01,
1.1041e-01],
[8.1445e-01, 9.1895e-01, 8.4229e-01, ..., 1.1676e-01, 4.8340e-01,
2.0715e-01],
[8.5449e-01, 7.3340e-01, 8.3984e-01, ..., 6.5967e-01, 8.4045e-02,
1.8262e-01]],
[[3.4131e-01, 6.5430e-01, 1.7468e-01, ..., 1.0962e-01, 8.1006e-01,
1.5454e-01],
[5.5371e-01, 6.7993e-02, 6.3037e-01, ..., 4.8633e-01, 6.7139e-01,
5.9700e-03],
[7.2510e-01, 7.5244e-01, 6.8237e-02, ..., 4.5068e-01, 6.8932e-03,
8.8623e-01],
[5.4199e-01, 3.7646e-01, 9.2871e-01, ..., 8.0371e-01, 9.8291e-01,
7.0898e-01],
[6.4258e-01, 9.9902e-01, 7.2949e-01, ..., 9.8779e-01, 9.3750e-02,
6.5576e-01],
[1.1945e-01, 5.6934e-01, 5.6738e-01, ..., 5.8447e-01, 6.6260e-01,
6.7139e-01]]],
[[[1.8356e-02, 4.2847e-01, 6.1572e-01, ..., 8.1787e-01, 3.4009e-01,
9.6729e-01],
[8.3740e-01, 8.0713e-01, 6.2793e-01, ..., 5.1575e-02, 6.0547e-01,
5.2100e-01],
[2.8870e-02, 3.2642e-01, 4.3408e-01, ..., 2.8955e-01, 2.6562e-01,
5.7520e-01],
[4.9683e-01, 4.9902e-01, 4.4312e-01, ..., 6.7725e-01, 9.5557e-01,
1.3770e-01],
[9.9731e-02, 2.2778e-01, 1.4148e-01, ..., 9.8050e-05, 6.3965e-01,
4.8267e-01],
[1.4832e-01, 3.0591e-01, 5.9424e-01, ..., 1.8433e-01, 1.0815e-01,
6.4990e-01]],
[[9.0674e-01, 9.1748e-01, 9.4482e-01, ..., 7.9688e-01, 4.8145e-01,
9.4385e-01],
[9.2407e-02, 1.3782e-01, 9.1357e-01, ..., 5.4541e-01, 7.0703e-01,
4.3628e-01],
[3.7573e-01, 4.3140e-01, 3.0103e-01, ..., 7.7832e-01, 6.1127e-02,
4.6313e-01],
[8.7988e-01, 1.5479e-01, 4.8364e-01, ..., 9.8340e-01, 6.4746e-01,
1.2622e-01],
[6.8213e-01, 5.4199e-02, 9.1748e-01, ..., 2.9102e-01, 8.9795e-01,
7.9346e-01],
[3.1226e-01, 7.6709e-01, 6.3574e-01, ..., 9.7070e-01, 7.3120e-02,
5.9863e-01]],
[[2.5244e-01, 3.2104e-01, 7.5098e-01, ..., 9.3311e-01, 7.6123e-01,
9.3994e-01],
[4.3945e-01, 6.3477e-01, 9.5215e-01, ..., 2.2998e-01, 3.7915e-01,
9.8193e-01],
[7.6172e-01, 1.7120e-02, 8.3466e-03, ..., 5.5273e-01, 6.7090e-01,
1.5527e-01],
[5.6250e-01, 6.6797e-01, 2.6514e-01, ..., 7.3389e-01, 2.0972e-01,
6.5527e-01],
[1.3635e-01, 1.7761e-01, 1.4709e-01, ..., 2.4695e-01, 8.9648e-01,
9.4385e-01],
[2.3218e-01, 2.6993e-02, 6.5674e-01, ..., 7.0557e-01, 4.5972e-01,
2.4805e-01]]]]
V: shape=(2, 3, 6, 8), dtype=float16
[[[[0.2004 , 0.4282 , 0.9204 , ..., 0.06033 , 0.3813 , 0.8096 ],
[0.1216 , 0.3513 , 0.8833 , ..., 0.4768 , 0.5493 , 0.9263 ],
[0.4238 , 0.597 , 0.812 , ..., 0.499 , 0.743 , 0.3118 ],
[0.3052 , 0.486 , 0.628 , ..., 0.5186 , 0.7974 , 0.564 ],
[0.291 , 0.605 , 0.704 , ..., 0.442 , 0.25 , 0.422 ],
[0.736 , 0.3054 , 0.3667 , ..., 0.1913 , 0.8413 , 0.265 ]],
[[0.3218 , 0.2054 , 0.3962 , ..., 0.205 , 0.01563 , 0.613 ],
[0.5586 , 0.6885 , 0.5254 , ..., 0.9365 , 0.4194 , 0.6265 ],
[0.2852 , 0.4028 , 0.94 , ..., 0.3547 , 0.403 , 0.607 ],
[0.1678 , 0.7637 , 0.7793 , ..., 0.445 , 0.52 , 0.4966 ],
[0.2012 , 0.451 , 0.6357 , ..., 0.708 , 0.02582 , 0.3772 ],
[0.208 , 0.91 , 0.05463 , ..., 0.659 , 0.3806 , 0.9727 ]],
[[0.339 , 0.8535 , 0.03452 , ..., 0.3306 , 0.747 , 0.1459 ],
[0.5737 , 0.7573 , 0.527 , ..., 0.7695 , 0.4905 , 0.402 ],
[0.2186 , 0.3645 , 0.937 , ..., 0.539 , 0.6797 , 0.3857 ],
[0.4094 , 0.1833 , 0.6226 , ..., 0.218 , 0.5874 , 0.383 ],
[0.6094 , 0.988 , 0.394 , ..., 0.9004 , 0.1876 , 0.82 ],
[0.1097 , 0.3174 , 0.01627 , ..., 0.696 , 0.6597 , 0.7627 ]]],
[[[0.5303 , 0.9775 , 0.7065 , ..., 0.2302 , 0.884 , 0.4202 ],
[0.319 , 0.5386 , 0.635 , ..., 0.1444 , 0.9453 , 0.6396 ],
[0.2515 , 0.8374 , 0.4248 , ..., 0.595 , 0.6313 , 0.4097 ],
[0.932 , 0.4756 , 0.6533 , ..., 0.7554 , 0.933 , 0.3604 ],
[0.2925 , 0.3213 , 0.3499 , ..., 0.312 , 0.544 , 0.2612 ],
[0.09314 , 0.007397, 0.4978 , ..., 0.969 , 0.3613 , 0.5234 ]],
[[0.4224 , 0.433 , 0.5327 , ..., 0.947 , 0.985 , 0.6973 ],
[0.1969 , 0.263 , 0.844 , ..., 0.4788 , 0.218 , 0.9756 ],
[0.661 , 0.427 , 0.274 , ..., 0.316 , 0.567 , 0.6904 ],
[0.888 , 0.621 , 0.618 , ..., 0.2245 , 0.7417 , 0.8154 ],
[0.278 , 0.8066 , 0.12445 , ..., 0.486 , 0.1539 , 0.4636 ],
[0.0608 , 0.7104 , 0.995 , ..., 0.398 , 0.42 , 0.3901 ]],
[[0.793 , 0.326 , 0.885 , ..., 0.7734 , 0.859 , 0.4893 ],
[0.86 , 0.4106 , 0.4014 , ..., 0.4407 , 0.992 , 0.297 ],
[0.975 , 0.989 , 0.606 , ..., 0.744 , 0.04483 , 0.49 ],
[0.6797 , 0.85 , 0.5127 , ..., 0.04388 , 0.4456 , 0.3848 ],
[0.8438 , 0.04922 , 0.589 , ..., 0.2152 , 0.3254 , 0.4387 ],
[0.165 , 0.4075 , 0.2546 , ..., 0.1492 , 0.01906 , 0.75 ]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[6, 5]
Outputs:
Y: shape=(2, 9, 4, 8), dtype=float16
[[[[0.2527, 0.4624, 0.872 , ..., 0.3376, 0.556 , 0.6753],
[0.2612, 0.464 , 0.803 , ..., 0.3818, 0.618 , 0.657 ],
[0.2632, 0.483 , 0.791 , ..., 0.3936, 0.563 , 0.627 ],
[0.338 , 0.4624, 0.727 , ..., 0.3643, 0.606 , 0.5615]],
[[0.2432, 0.4548, 0.875 , ..., 0.3325, 0.548 , 0.695 ],
[0.2607, 0.464 , 0.8086, ..., 0.3757, 0.6123, 0.6587],
[0.2651, 0.4836, 0.789 , ..., 0.3943, 0.568 , 0.6245],
[0.3433, 0.451 , 0.7173, ..., 0.3564, 0.607 , 0.568 ]],
[[0.2512, 0.4617, 0.874 , ..., 0.3254, 0.5493, 0.679 ],
[0.2605, 0.464 , 0.809 , ..., 0.378 , 0.613 , 0.6587],
[0.2605, 0.4832, 0.792 , ..., 0.3875, 0.55 , 0.6304],
[0.3384, 0.457 , 0.7217, ..., 0.36 , 0.5913, 0.5654]],
...,
[[0.3652, 0.64 , 0.5244, ..., 0.543 , 0.6436, 0.3142],
[0.3916, 0.501 , 0.541 , ..., 0.4402, 0.619 , 0.337 ],
[0.4385, 0.595 , 0.52 , ..., 0.5386, 0.522 , 0.4478],
[0.3894, 0.5645, 0.4656, ..., 0.5737, 0.533 , 0.5054]],
[[0.3855, 0.665 , 0.4956, ..., 0.554 , 0.634 , 0.3135],
[0.3716, 0.4895, 0.57 , ..., 0.4421, 0.629 , 0.3374],
[0.4404, 0.6143, 0.5024, ..., 0.543 , 0.5205, 0.4478],
[0.387 , 0.55 , 0.443 , ..., 0.56 , 0.542 , 0.4988]],
[[0.3594, 0.6367, 0.525 , ..., 0.5366, 0.6475, 0.312 ],
[0.3806, 0.4756, 0.563 , ..., 0.4282, 0.623 , 0.3398],
[0.4365, 0.633 , 0.496 , ..., 0.5527, 0.524 , 0.4446],
[0.3855, 0.5566, 0.4407, ..., 0.563 , 0.542 , 0.4993]]],
[[[0.4082, 0.724 , 0.665 , ..., 0.1807, 0.9194, 0.547 ],
[0.3777, 0.786 , 0.601 , ..., 0.3052, 0.8325, 0.4927],
[0.497 , 0.676 , 0.614 , ..., 0.3953, 0.868 , 0.4817],
[0.4795, 0.636 , 0.5713, ..., 0.4038, 0.8105, 0.4333]],
[[0.4163, 0.7407, 0.668 , ..., 0.1838, 0.917 , 0.5386],
[0.3752, 0.779 , 0.6016, ..., 0.3013, 0.8354, 0.497 ],
[0.505 , 0.714 , 0.613 , ..., 0.4094, 0.8555, 0.4624],
[0.4805, 0.6406, 0.5747, ..., 0.3943, 0.813 , 0.4329]],
[[0.4006, 0.708 , 0.6626, ..., 0.1775, 0.9214, 0.555 ],
[0.3706, 0.762 , 0.6025, ..., 0.2922, 0.842 , 0.5063],
[0.4902, 0.691 , 0.615 , ..., 0.386 , 0.866 , 0.4802],
[0.4753, 0.6274, 0.5747, ..., 0.3865, 0.816 , 0.439 ]],
...,
[[0.828 , 0.37 , 0.6323, ..., 0.5996, 0.9287, 0.389 ],
[0.8657, 0.5356, 0.641 , ..., 0.65 , 0.684 , 0.4229],
[0.8223, 0.6147, 0.61 , ..., 0.506 , 0.6216, 0.414 ],
[0.8228, 0.4868, 0.605 , ..., 0.4531, 0.609 , 0.4133]],
[[0.8276, 0.3699, 0.634 , ..., 0.6006, 0.928 , 0.3894],
[0.8643, 0.5317, 0.647 , ..., 0.6533, 0.6865, 0.4246],
[0.8267, 0.6123, 0.6016, ..., 0.5083, 0.6265, 0.411 ],
[0.825 , 0.5234, 0.6055, ..., 0.4617, 0.585 , 0.4165]],
[[0.829 , 0.3713, 0.625 , ..., 0.5947, 0.9307, 0.386 ],
[0.8643, 0.529 , 0.6396, ..., 0.647 , 0.6943, 0.421 ],
[0.821 , 0.626 , 0.611 , ..., 0.5005, 0.607 , 0.416 ],
[0.8286, 0.5254, 0.605 , ..., 0.4536, 0.549 , 0.4204]]]]
Differences with previous version (23)#
SchemaDiff: Attention (domain 'ai.onnx')
old version: 23
new version: 24
breaking: yes
Breaking reasons:
input ‘nonpad_kv_seqlen’ (added): at position 6; option=Single; type_str=’tensor(int64)’
Inputs:
[BREAKING] added ‘nonpad_kv_seqlen’: at position 6; option=Single; type_str=’tensor(int64)’
Documentation:
line similarity: 0.83 (+16/-2 lines)
--- Attention v23
+++ Attention v24
@@ -14,10 +14,24 @@
3) Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`.
Attention bias to be added is calculated based on `attn_mask` input and `is_causal` attribute:
-1) If `is_causal` is set to `1`, a query index i attends keys j <= i + past_sequence_length.
-2) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.
+1) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.
+2) If `is_causal` is set to `1`, attention scores above the causal frontier are masked out. For internal cache (`past_key`) this is the standard offset from `past_sequence_length`; for external cache (`nonpad_kv_seqlen` without `past_key`) this is bottom-right aligned by `nonpad_kv_seqlen - q_sequence_length`.
3) If both `attn_mask` and `is_causal` are set, the valid positions are the intersection of both masks.
If a query row is fully masked after this intersection, its output row is zero.
+
+With respect to KV cache update, this operator allows the following two use cases:
+
+1) Cache update happens inside the Attention operator. In this case, the `K` and `V` inputs contain only the incoming
+tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are
+all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present
+key and value, respectively. Note that this only works correctly for the special case where the past key and value
+do not contain padded tokens.
+2) Cache update happens outside the Attention operator (for example, through the `TensorScatter` operator). In this
+case, the `K` and `V` inputs correspond to the entire cache tensor, so the four optional inputs/outputs past and
+present key and value should not be used. An additional input `nonpad_kv_seqlen` of shape (batch_size,) may be
+provided to indicate the number of non-padding tokens in each sample of the batch to save unnecessary computation.
+Here, the kv_sequence dimension of `attn_mask` can be shorter than `K` and `V`, but still needs to be at least as long
+as the maximum value of `nonpad_kv_seqlen`.
Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them.
The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided: