.. _op_ai_onnx_Attention: Attention ========= - **Domain**: ``ai.onnx`` - **Since 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: 1) Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, ``q_num_heads = kv_num_heads``. 2) 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``. 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) ``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: .. code-block:: 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_size`` - **K** (*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_size`` - **V** (*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_size`` - **attn_mask** (*U*): Attention mask. Shape must be broadcastable to ``(batch_size, q_num_heads, q_sequence_length, total_sequence_length)`` where ``total_sequence_length = past_sequence_length + kv_sequence_length.`` The last dimension can also be shorter than ``total_sequence_length`` and will be padded to ``total_sequence_length`` with negative infinity. Two types of masks are supported: 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. - **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 with ``past_key`` and ``past_value`` inputs or ``present_key`` and ``present_value`` outputs (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_size`` - **present_key** (*T1*): Updated key cache with shape ``(batch_size, kv_num_heads, total_sequence_length, head_size)`` where ``total_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)`` where ``total_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)`` where ``total_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** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output) Attributes: qk_matmul_output_mode = 3 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V) -> (Y, "", "", qk_matmul_output) Attributes: qk_matmul_output_mode = 3 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y) Attributes: scale = 0.5 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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** .. code-block:: text Node: Attention(Q, K, V, "", "", "", nonpad_kv_seqlen) -> (Y) Attributes: is_causal = 1 .. code-block:: text 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) .. code-block:: diff --- 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: Version History --------------- - :doc:`Version 23 `