:nosearch: .. _op_ai_onnx_Attention-23: Attention - version 23 ====================== This page documents version **23** of operator **Attention**. See :doc:`Attention` for the latest version (since version 24). - **Domain**: ``ai.onnx`` - **Since version**: 23 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) 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. 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. 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 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.`` Two types of masks are supported. A boolean mask where a value of ``True`` indicates that the element should take part in attention. Also supports 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)`` **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_23_boolmask_fullymasked_row_nan_robustness** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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], [ True, True, True]], [[False, False, False], [ True, True, True]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0. , 0. ], [-0.16495374, 0.7457248 ]], [[ 0. , 0. ], [ 0.18005186, -0.17207974]]]] **test_cc_attention_23_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_3d** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.23092115, 0.5444725 , 0.6482418 , -0.37858847], [-0.23092115, 0.77539366, 0.04638294, -0.08028026]]] **test_cc_attention_3d_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] attn_mask: shape=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.458196 , 0.41020232, 0.9921614 , -0.7460807 ], [ 0.0697852 , 0.6150191 , 0.3994022 , -0.34422377]]] **test_cc_attention_3d_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0. , 1. , 1. , -1. ], [ 0.5 , 0.5 , 0.25, 0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0.412521 , 0.587479 , 0.79335546, -0.19003323], [ 0. , 0.6666667 , 0.7425821 , -0.5205673 ]]] **test_cc_attention_3d_diff_heads_sizes** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] Outputs: Y: shape=(1, 2, 6), dtype=float32 [[[ 0.23092115, 0.5444725 , 0.296508 , 0.6482418 , -0.37858847, 0.3571369 ], [-0.23092115, 0.77539366, 0.64288974, 0.04638294, -0.08028026, 0.57520705]]] **test_cc_attention_3d_diff_heads_sizes_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] attn_mask: shape=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 6), dtype=float32 [[[ 0.458196 , 0.41020232, 0.03320454, 0.9921614 , -0.7460807 , 0.4670852 ], [ 0.0697852 , 0.6150191 , 0.3722638 , 0.3994022 , -0.34422377, 0.5532167 ]]] **test_cc_attention_3d_diff_heads_sizes_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0. , 1. , 1. , -1. ], [ 0.5 , 0.5 , 0.25, 0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] Outputs: Y: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0.412521 , 0.587479 , 0.762437 , 0.79335546, -0.19003323, -0.0055371 ], [ 0. , 0.6666667 , 0.5 , 0.7425821 , -0.5205673 , 0.43631145]]] **test_cc_attention_3d_diff_heads_sizes_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.009999999776482582 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] Outputs: Y: shape=(1, 2, 6), dtype=float32 [[[ 0.00333332, 0.6649986 , 0.49749377, 0.6665943 , -0.58079195, 0.5803768 ], [-0.00333332, 0.6683319 , 0.50249374, 0.65489024, -0.5724745 , 0.5816217 ]]] **test_cc_attention_3d_diff_heads_sizes_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 1.0 softcap = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] Outputs: Y: shape=(1, 2, 6), dtype=float32 [[[ 0.15169363, 0.6033754 , 0.44272968, 0.67558366, -0.4552321 , 0.4223395 ], [-0.15169363, 0.755069 , 0.67027014, 0.29168454, -0.3030643 , 0.60680556]]] **test_cc_attention_3d_diff_heads_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 6), dtype=float32 [[[ 1. , 0. , -1. , 2. , -2. , 1. ], [ 0. , 1. , 2. , 0.5 , 0.25, -0.25], [-1. , 1. , 0.5 , -0.5 , 0. , 1. ]]] attn_mask: shape=(2, 5), dtype=float32 [[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.5 , 0.5 , -1. ], [ 0. , 0.25, 0.5 ]], [[ 0. , 0.5 , 0.5 ], [-0.5 , 0.75, -0.25]]]] Outputs: Y: shape=(1, 2, 6), dtype=float32 [[[ 0.04906689, 0.6637103 , 0.40457496, 0.17558974, 0.17481029, 0.24073282], [-0.03033203, 0.6089026 , 0.27225977, 0.01557301, 0.22852218, 0.47354427]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 3), dtype=float32 [[[[ 0.5 , 0.5 , -1. ], [ 0. , 0.25, 0.5 ], [ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 0. , 0.5 , 0.5 ], [-0.5 , 0.75, -0.25], [ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] **test_cc_attention_3d_gqa** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 4 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 8), dtype=float32 [[[-0.02356532, 0.6783799 , -0.03533878, 0.6841799 , 0.6482418 , -0.37858847, 0.37784207, -0.12898168], [-0.02356531, 0.6783799 , 0.11724145, 0.6063233 , 0.9917567 , -0.74587834, 0.29831943, -0.26321504]]] **test_cc_attention_3d_gqa_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: q_num_heads = 4 kv_num_heads = 2 scale = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] attn_mask: shape=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 8), dtype=float32 [[[ 0.30531266, 0.50214726, 0.29782698, 0.50647473, 0.9921614 , -0.7460807 , 0.74361753, -0.47596362], [ 0.2229336 , 0.5335671 , 0.31907293, 0.4797093 , 1.1873223 , -1.0039468 , 0.6715176 , -0.5857588 ]]] **test_cc_attention_3d_gqa_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 4 kv_num_heads = 2 is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 3, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ], [ 0.2 , -0.1 , 0.25, 0. , 0.5 , 0.5 , -1. , 1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 3, 8), dtype=float32 [[[ 1. , 0. , 1. , 0. , 2. , -2. , 2. , -2. ], [ 0.4911621 , 0.50883794, 0.5440794 , 0.45592055, 1.25 , -0.875 , 0.9953577 , -0.4930365 ], [ 0.07057842, 0.63075423, 0.05884897, 0.636809 , 0.6482418 , -0.37858847, 1.478919 , -1.3989784 ]]] **test_cc_attention_3d_gqa_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 4 kv_num_heads = 2 scale = 0.009999999776482582 .. code-block:: text Inputs: Q: shape=(1, 2, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 8), dtype=float32 [[[-3.3333333e-04, 6.6683334e-01, -4.9999997e-04, 6.6691661e-01, 6.6659433e-01, -5.8079195e-01, 6.6071773e-01, -5.7536310e-01], [-3.3333330e-04, 6.6683334e-01, 1.6666650e-03, 6.6583300e-01, 6.7248535e-01, -5.8624268e-01, 6.6077113e-01, -5.7789290e-01]]] **test_cc_attention_3d_gqa_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 4 kv_num_heads = 2 scale = 1.0 softcap = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 8), dtype=float32 [[[-0.03040495, 0.6819007 , -0.04947086, 0.69103485, 0.67558366, -0.4552321 , 0.40459138, -0.23742941], [-0.02114749, 0.67739695, 0.15203254, 0.5866078 , 0.94510204, -0.72255105, 0.35040644, -0.32931072]]] **test_cc_attention_3d_gqa_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) Attributes: q_num_heads = 4 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 8), dtype=float32 [[[ 0.1 , 0.2 , -0.1 , 0.05, 0.5 , 0.5 , 1. , 0. ], [ 0.3 , 0.4 , 0.2 , -0.3 , 0. , 1. , 0.5 , -0.5 ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 8), dtype=float32 [[[-0.12450287, 0.51090705, -0.13178737, 0.5070071 , 0.48057237, -0.20996696, 0.27224937, 0.02852735], [-0.12746866, 0.51840895, 0.00804363, 0.47238812, 0.7363033 , -0.5218315 , 0.23465633, -0.06618155]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_3d_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.009999999776482582 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.00333332, 0.6649986 , 0.6665943 , -0.58079195], [-0.00333332, 0.6683319 , 0.65489024, -0.5724745 ]]] **test_cc_attention_3d_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 1.0 softcap = 0.5 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.15169363, 0.6033754 , 0.67558366, -0.4552321 ], [-0.15169363, 0.755069 , 0.29168454, -0.3030643 ]]] **test_cc_attention_3d_transpose_verification** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: q_num_heads = 3 kv_num_heads = 3 .. code-block:: text Inputs: Q: shape=(1, 2, 12), dtype=float32 [[[1., 1., 1., 1., 2., 2., 2., 2., 3., 3., 3., 3.], [1., 1., 1., 1., 2., 2., 2., 2., 3., 3., 3., 3.]]] K: shape=(1, 2, 12), dtype=float32 [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]] V: shape=(1, 2, 12), dtype=float32 [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]] Outputs: Y: shape=(1, 2, 12), dtype=float32 [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]] **test_cc_attention_3d_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) Attributes: q_num_heads = 2 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.10764056, 0.45608166, 0.48057237, -0.20996696], [-0.31939295, 0.57814157, 0.05822894, 0.11373253]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_3d_with_past_and_present_qk_matmul** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: q_num_heads = 2 kv_num_heads = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.10764056, 0.45608166, 0.48057237, -0.20996696], [-0.31939295, 0.57814157, 0.05822894, 0.11373253]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.35355338, 0. , 0.70710677, 0.35355338, 0. ], [-0.35355338, 0.35355338, 0. , 0.35355338, 0.70710677]], [[ 0.35355338, 0.17677669, 0. , 0.70710677, -0.08838835], [ 0.70710677, -1.0606601 , -1.4142135 , 0. , 0.53033006]]]] **test_cc_attention_3d_with_past_and_present_qk_matmul_bias** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 0.5 qk_matmul_output_mode = 1 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] attn_mask: shape=(2, 5), dtype=float32 [[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[-0.07067307, 0.63377535, 0.30545416, 0.01247976], [-0.23752189, 0.5571051 , 0.09629637, 0.127618 ]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_3d_with_past_and_present_qk_matmul_softcap** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: q_num_heads = 2 kv_num_heads = 2 scale = 1.0 softcap = 0.5 qk_matmul_output_mode = 2 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.05362739, 0.4880938 , 0.4812978 , -0.23981026], [-0.27949443, 0.5497092 , 0.22774768, -0.08467997]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.3807971 , 0. , 0.4820138 , 0.3807971 , 0. ], [-0.3807971 , 0.3807971 , 0. , 0.3807971 , 0.4820138 ]], [[ 0.3807971 , 0.23105858, 0. , 0.4820138 , -0.12245933], [ 0.4820138 , -0.4975274 , -0.49966466, 0. , 0.45257413]]]] **test_cc_attention_3d_with_past_and_present_qk_matmul_softmax** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: q_num_heads = 2 kv_num_heads = 2 qk_matmul_output_mode = 3 .. code-block:: text Inputs: Q: shape=(1, 2, 4), dtype=float32 [[[ 1. , 0. , 0.5, 0.5], [ 0. , 1. , 1. , -1. ]]] K: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , -1. , 1. ], [ 0.5 , 0.5 , 1. , 1. ], [ 0. , 1. , 0.25, -0.5 ]]] V: shape=(1, 3, 4), dtype=float32 [[[ 1. , 0. , 2. , -2. ], [ 0. , 1. , 0.5 , 0.25], [-1. , 1. , -0.5 , 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 4), dtype=float32 [[[ 0.10764056, 0.45608166, 0.48057237, -0.20996696], [-0.31939295, 0.57814157, 0.05822894, 0.11373253]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[0.20710382, 0.14542593, 0.2949405 , 0.20710382, 0.14542593], [0.10673924, 0.21647945, 0.15200938, 0.21647945, 0.3082925 ]], [[0.21705809, 0.18188749, 0.15241571, 0.30911657, 0.13952214], [0.38144314, 0.06511761, 0.04572483, 0.18807767, 0.31963673]]]] **test_cc_attention_4d** .. code-block:: text Node: Attention(Q, K, V) -> (Y) .. 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, 2, 2), dtype=float32 [[[[ 1., 0.], [ 0., 1.]], [[ 1., 1.], [-1., 1.]]]] V: shape=(1, 2, 2, 2), dtype=float32 [[[[ 1., 2.], [ 3., 4.]], [[-1., 0.], [ 0., 1.]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 1.6604769 , 2.660477 ], [ 2.339523 , 3.339523 ]], [[-0.66976154, 0.33023846], [-0.80442965, 0.19557032]]]] **test_cc_attention_4d_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.458196 , 0.41020232], [ 0.0697852 , 0.6150191 ]], [[ 0.9921614 , -0.7460807 ], [ 0.3994022 , -0.34422377]]]] **test_cc_attention_4d_attn_mask_3d** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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, 3), dtype=float32 [[[ 0. , -1. , 0.5], [ 0.2, 0. , -0.4]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0. , 0.5626513 ], [ 0.03219185, 0.6617174 ]], [[ 0.48452467, -0.58578634], [ 0.3757273 , -0.26752764]]]] **test_cc_attention_4d_attn_mask_3d_causal** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0. , 1. ], [ 0.5 , 0.5 ]], [[-1. , 1. ], [ 1. , -1. ], [ 0.25, 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. ]]]] attn_mask: shape=(1, 3, 3), dtype=float32 [[[ 0. , -1. , 0.5], [ 0.2, 0. , -0.4], [ 0.1, -0.3, 0. ]]] Outputs: Y: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0.4875026 , 0.5124974 ], [ 0.03695408, 0.61167425]], [[ 2. , -2. ], [ 0.9650383 , -0.44755742], [ 0.7986912 , -0.689716 ]]]] **test_cc_attention_4d_attn_mask_4d** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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=float32 [[[[ 0. , -0.5, -1. ], [ 0.2, 0. , -0.1]], [[-0.2, 0.3, 0. ], [ 0. , -0.1, 0.4]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.458196 , 0.41020232], [-0.06765194, 0.6944395 ]], [[ 0.5724663 , -0.27137595], [ 0.02987868, -0.14798252]]]] **test_cc_attention_4d_attn_mask_4d_causal** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0. , 1. ], [ 0.5 , 0.5 ]], [[-1. , 1. ], [ 1. , -1. ], [ 0.25, 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. ]]]] attn_mask: shape=(1, 2, 3, 3), dtype=float32 [[[[ 0. , -0.5, -1. ], [ 0.2, 0. , -0.1], [ 0.5, -0.2, 0. ]], [[-0.2, 0.3, 0. ], [ 0. , -0.1, 0.4], [ 0.1, 0. , -0.3]]]] Outputs: Y: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0.4875026 , 0.5124974 ], [ 0.18708876, 0.52451503]], [[ 2. , -2. ], [ 0.93357575, -0.40036362], [ 0.8560081 , -0.63305765]]]] **test_cc_attention_4d_attn_mask_bool** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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 [[[[ True, True, False], [ True, False, True]], [[ True, False, True], [ True, True, False]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5621765 , 0.4378235 ], [-0.24491866, 0.62245935]], [[ 0.7890498 , -1.0312399 ], [ 0.9034121 , -0.3551182 ]]]] **test_cc_attention_4d_attn_mask_bool_4d** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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 [[[[ True, True, False], [ True, False, True]], [[ True, False, True], [ True, True, False]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5621765 , 0.4378235 ], [-0.24491866, 0.62245935]], [[ 0.7890498 , -1.0312399 ], [ 0.9034121 , -0.3551182 ]]]] **test_cc_attention_4d_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0. , 1. ], [ 0.5 , 0.5 ]], [[-1. , 1. ], [ 1. , -1. ], [ 0.25, 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. ]]]] Outputs: Y: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0.412521 , 0.587479 ], [ 0. , 0.6666667 ]], [[ 2. , -2. ], [ 0.79335546, -0.19003323], [ 0.7425821 , -0.5205673 ]]]] **test_cc_attention_4d_causal_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) Attributes: is_causal = 1 .. 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. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.39085424, 0.15993309], [-0.01604789, 0.39012018]], [[ 0.7178145 , -0.5209762 ], [ 0.3204866 , 0.16716442]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_diff_heads_mask4d_padded_kv** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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. ]]]] attn_mask: shape=(1, 1, 2, 3), dtype=float32 [[[[ 0., 0., -10000.], [ 0., 0., -10000.]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[ 0.49375033, 0.50624967], [ 0.49375033, 0.50624967]], [[ 0.4906261 , 0.5093739 ], [ 0.53120935, 0.46879062]], [[ 1.066311 , -0.5994665 ], [ 1.25 , -0.875 ]], [[ 0.9034121 , -0.3551182 ], [ 1.066311 , -0.5994665 ]]]] **test_cc_attention_4d_diff_heads_sizes** .. code-block:: text Node: Attention(Q, K, V) -> (Y) .. 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, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] Outputs: Y: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.23092115, 0.5444725 , 0.296508 ], [-0.23092115, 0.77539366, 0.64288974]], [[ 0.6482418 , -0.37858847, 0.3571369 ], [ 0.04638294, -0.08028026, 0.57520705]]]] **test_cc_attention_4d_diff_heads_sizes_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] attn_mask: shape=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.458196 , 0.41020232, 0.03320454], [ 0.0697852 , 0.6150191 , 0.3722638 ]], [[ 0.9921614 , -0.7460807 , 0.4670852 ], [ 0.3994022 , -0.34422377, 0.5532167 ]]]] **test_cc_attention_4d_diff_heads_sizes_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 2, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0. , 1. ], [ 0.5 , 0.5 ]], [[-1. , 1. ], [ 1. , -1. ], [ 0.25, 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, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] Outputs: Y: shape=(1, 2, 3, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0.412521 , 0.587479 , 0.762437 ], [ 0. , 0.6666667 , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.79335546, -0.19003323, -0.0055371 ], [ 0.7425821 , -0.5205673 , 0.43631145]]]] **test_cc_attention_4d_diff_heads_sizes_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 0.009999999776482582 .. 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, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] Outputs: Y: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.00333332, 0.6649986 , 0.49749377], [-0.00333332, 0.6683319 , 0.50249374]], [[ 0.6665943 , -0.58079195, 0.5803768 ], [ 0.65489024, -0.5724745 , 0.5816217 ]]]] **test_cc_attention_4d_diff_heads_sizes_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 1.0 softcap = 0.5 .. 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, 3), dtype=float32 [[[[ 1. , 0. , -1. ], [ 0. , 1. , 2. ], [-1. , 1. , 0.5 ]], [[ 2. , -2. , 1. ], [ 0.5 , 0.25, -0.25], [-0.5 , 0. , 1. ]]]] Outputs: Y: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.15169363, 0.6033754 , 0.44272968], [-0.15169363, 0.755069 , 0.67027014]], [[ 0.67558366, -0.4552321 , 0.4223395 ], [ 0.29168454, -0.3030643 , 0.60680556]]]] **test_cc_attention_4d_diff_heads_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) .. 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. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.12450287, 0.51090705], [-0.12746866, 0.51840895]], [[-0.13178737, 0.5070071 ], [ 0.00804363, 0.47238812]], [[ 0.48057237, -0.20996696], [ 0.7363033 , -0.5218315 ]], [[ 0.27224937, 0.02852735], [ 0.23465633, -0.06618155]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_diff_heads_with_past_and_present_mask3d** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value) Attributes: scale = 0.5 .. 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. ]]]] attn_mask: shape=(1, 2, 5), dtype=float32 [[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.19640994, 0.6580199 ], [-0.09580244, 0.52036375]], [[-0.19427064, 0.65238005], [ 0.00927786, 0.49179137]], [[ 0.30545416, 0.01247976], [ 0.560947 , -0.33633795]], [[ 0.20446573, 0.11752708], [ 0.21916693, -0.00962794]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_diff_heads_with_past_and_present_mask4d** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value) Attributes: scale = 0.5 .. 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. ]]]] attn_mask: shape=(1, 1, 2, 5), dtype=float32 [[[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.19640994, 0.6580199 ], [-0.09580244, 0.52036375]], [[-0.19427064, 0.65238005], [ 0.00927786, 0.49179137]], [[ 0.30545416, 0.01247976], [ 0.560947 , -0.33633795]], [[ 0.20446573, 0.11752708], [ 0.21916693, -0.00962794]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_fp16** .. code-block:: text Node: Attention(Q, K, V) -> (Y) .. code-block:: text Inputs: Q: shape=(2, 3, 4, 8), dtype=float16 [[[[0.768 , 0.2769 , 0.381 , ..., 0.3386 , 0.9316 , 0.6865 ], [0.797 , 0.695 , 0.511 , ..., 0.941 , 0.3232 , 0.735 ], [0.0287 , 0.1525 , 0.3313 , ..., 0.9663 , 0.2722 , 0.742 ], [0.9644 , 0.906 , 0.924 , ..., 0.9604 , 0.619 , 0.382 ]], [[0.1196 , 0.4138 , 0.012245, ..., 0.9214 , 0.09467 , 0.907 ], [0.5264 , 0.84 , 0.4602 , ..., 0.802 , 0.821 , 0.6484 ], [0.232 , 0.549 , 0.772 , ..., 0.005775, 0.809 , 0.514 ], [0.7993 , 0.6104 , 0.593 , ..., 0.891 , 0.4631 , 0.4915 ]], [[0.8076 , 0.2961 , 0.208 , ..., 0.2695 , 0.07605 , 0.145 ], [0.35 , 0.581 , 0.0441 , ..., 0.9043 , 0.3318 , 0.1179 ], [0.537 , 0.6436 , 0.168 , ..., 0.01337 , 0.805 , 0.6313 ], [0.994 , 0.3652 , 0.787 , ..., 0.767 , 0.462 , 0.985 ]]], [[[0.3628 , 0.83 , 0.7495 , ..., 0.7886 , 0.1573 , 0.0558 ], [0.799 , 0.8257 , 0.04434 , ..., 0.655 , 0.1383 , 0.7993 ], [0.4753 , 0.3433 , 0.3005 , ..., 0.3953 , 0.542 , 0.75 ], [0.0801 , 0.0763 , 0.2883 , ..., 0.993 , 0.8965 , 0.5186 ]], [[0.3167 , 0.922 , 0.4185 , ..., 0.885 , 0.6196 , 0.296 ], [0.4043 , 0.48 , 0.2421 , ..., 0.4587 , 0.521 , 0.3496 ], [0.578 , 0.5537 , 0.951 , ..., 0.5576 , 0.3003 , 0.526 ], [0.588 , 0.65 , 0.878 , ..., 0.9766 , 0.04785 , 0.749 ]], [[0.2006 , 0.4778 , 0.9956 , ..., 0.965 , 0.7446 , 0.5215 ], [0.6973 , 0.4497 , 0.4187 , ..., 0.723 , 0.7617 , 0.4243 ], [0.376 , 0.1816 , 0.9004 , ..., 0.9067 , 0.3618 , 0.6157 ], [0.04916 , 0.992 , 0.3367 , ..., 0.12354 , 0.1976 , 0.3323 ]]]] K: shape=(2, 3, 6, 8), dtype=float16 [[[[0.6216 , 0.861 , 0.5938 , ..., 0.0877 , 0.545 , 0.1114 ], [0.04868 , 0.9824 , 0.7256 , ..., 0.923 , 0.00809 , 0.6504 ], [0.9204 , 0.1576 , 0.1031 , ..., 0.1365 , 0.8228 , 0.3477 ], [0.4634 , 0.08734 , 0.598 , ..., 0.8447 , 0.2439 , 0.26 ], [0.171 , 0.927 , 0.973 , ..., 0.8657 , 0.666 , 0.2727 ], [0.3013 , 0.7837 , 0.381 , ..., 0.222 , 0.1247 , 0.599 ]], [[0.378 , 0.2192 , 0.8647 , ..., 0.6387 , 0.7964 , 0.281 ], [0.691 , 0.7476 , 0.8584 , ..., 0.04382 , 0.4983 , 0.688 ], [0.795 , 0.3357 , 0.6875 , ..., 0.0688 , 0.1875 , 0.0173 ], [0.010956, 0.9043 , 0.0662 , ..., 0.548 , 0.4624 , 0.7954 ], [0.926 , 0.827 , 0.601 , ..., 0.01913 , 0.06143 , 0.376 ], [0.8906 , 0.4302 , 0.289 , ..., 0.2229 , 0.3643 , 0.4412 ]], [[0.5244 , 0.494 , 0.01675 , ..., 0.603 , 0.891 , 0.5015 ], [0.2092 , 0.751 , 0.06223 , ..., 0.8066 , 0.8022 , 0.395 ], [0.5947 , 0.3872 , 0.5938 , ..., 0.184 , 0.399 , 0.788 ], [0.09515 , 0.542 , 0.1709 , ..., 0.2878 , 0.623 , 0.6763 ], [0.3906 , 0.8076 , 0.2429 , ..., 0.375 , 0.5576 , 0.8813 ], [0.9272 , 0.7295 , 0.6504 , ..., 0.905 , 0.9253 , 0.554 ]]], [[[0.4163 , 0.3062 , 0.1239 , ..., 0.03952 , 0.59 , 0.675 ], [0.3347 , 0.489 , 0.7856 , ..., 0.1453 , 0.05606 , 0.5493 ], [0.4255 , 0.683 , 0.647 , ..., 0.2908 , 0.9688 , 0.313 ], [0.1964 , 0.2465 , 0.1116 , ..., 0.748 , 0.6133 , 0.43 ], [0.4219 , 0.9897 , 0.4724 , ..., 0.492 , 0.416 , 0.191 ], [0.641 , 0.6777 , 0.4138 , ..., 0.701 , 0.4443 , 0.2162 ]], [[0.6597 , 0.527 , 0.5103 , ..., 0.2566 , 0.9053 , 0.953 ], [0.04782 , 0.708 , 0.7275 , ..., 0.2428 , 0.647 , 0.1777 ], [0.6426 , 0.8447 , 0.921 , ..., 0.4475 , 0.6587 , 0.494 ], [0.04538 , 0.5137 , 0.3374 , ..., 0.8613 , 0.5693 , 0.834 ], [0.1244 , 0.7446 , 0.11426 , ..., 0.1786 , 0.1487 , 0.9575 ], [0.5127 , 0.5205 , 0.7793 , ..., 0.2957 , 0.808 , 0.954 ]], [[0.1094 , 0.277 , 0.7505 , ..., 0.8125 , 0.09845 , 0.7393 ], [0.5786 , 0.4038 , 0.8374 , ..., 0.533 , 0.2117 , 0.558 ], [0.851 , 0.7427 , 0.7397 , ..., 0.9634 , 0.4805 , 0.4832 ], [0.6226 , 0.0897 , 0.548 , ..., 0.8125 , 0.8774 , 0.915 ], [0.854 , 0.693 , 0.2424 , ..., 0.4934 , 0.517 , 0.0836 ], [0.2725 , 0.754 , 0.1952 , ..., 0.6216 , 0.1534 , 0.5454 ]]]] V: shape=(2, 3, 6, 8), dtype=float16 [[[[0.4753 , 0.4448 , 0.806 , ..., 0.837 , 0.1581 , 0.5366 ], [0.3005 , 0.2703 , 0.9395 , ..., 0.905 , 0.693 , 0.566 ], [0.812 , 0.1626 , 0.875 , ..., 0.3066 , 0.3735 , 0.953 ], [0.963 , 0.269 , 0.273 , ..., 0.7285 , 0.8687 , 0.1381 ], [0.2224 , 0.4397 , 0.9346 , ..., 0.8096 , 0.2372 , 0.638 ], [0.076 , 0.728 , 0.302 , ..., 0.642 , 0.4287 , 0.5503 ]], [[0.524 , 0.8896 , 0.9575 , ..., 0.2717 , 0.7837 , 0.04794 ], [0.582 , 0.885 , 0.1246 , ..., 0.1964 , 0.533 , 0.8843 ], [0.782 , 0.3752 , 0.1667 , ..., 0.868 , 0.2988 , 0.8896 ], [0.672 , 0.2277 , 0.0884 , ..., 0.1917 , 0.593 , 0.4727 ], [0.314 , 0.01035 , 0.03458 , ..., 0.0249 , 0.3174 , 0.1203 ], [0.787 , 0.4954 , 0.791 , ..., 0.6787 , 0.2668 , 0.8975 ]], [[0.686 , 0.1578 , 0.284 , ..., 0.418 , 0.6245 , 0.947 ], [0.6196 , 0.676 , 0.0801 , ..., 0.959 , 0.4658 , 0.99 ], [0.714 , 0.4312 , 0.887 , ..., 0.9727 , 0.2556 , 0.8267 ], [0.11017 , 0.007835, 0.05338 , ..., 0.5825 , 0.3499 , 0.8335 ], [0.4648 , 0.6934 , 0.0673 , ..., 0.865 , 0.4956 , 0.4666 ], [0.4504 , 0.979 , 0.05844 , ..., 0.3506 , 0.3293 , 0.7593 ]]], [[[0.255 , 0.05832 , 0.2969 , ..., 0.5215 , 0.8794 , 0.823 ], [0.08167 , 0.3281 , 0.693 , ..., 0.3137 , 0.0643 , 0.349 ], [0.6504 , 0.8887 , 0.2983 , ..., 0.6167 , 0.1927 , 0.10443 ], [0.6953 , 0.04318 , 0.8047 , ..., 0.773 , 0.4648 , 0.4355 ], [0.4678 , 0.798 , 0.04443 , ..., 0.07715 , 0.4705 , 0.7666 ], [0.2332 , 0.3625 , 0.491 , ..., 0.2786 , 0.691 , 0.1003 ]], [[0.7734 , 0.6494 , 0.758 , ..., 0.0784 , 0.9976 , 0.517 ], [0.886 , 0.7744 , 0.0998 , ..., 0.632 , 0.1531 , 0.6753 ], [0.884 , 0.5566 , 0.9116 , ..., 0.04623 , 0.3098 , 0.3125 ], [0.09576 , 0.01733 , 0.2286 , ..., 0.8833 , 0.511 , 0.743 ], [0.3657 , 0.8823 , 0.2277 , ..., 0.6562 , 0.1917 , 0.3606 ], [0.5537 , 0.07935 , 0.725 , ..., 0.3335 , 0.676 , 0.1355 ]], [[0.0349 , 0.24 , 0.679 , ..., 0.835 , 0.6787 , 0.1176 ], [0.2008 , 0.1987 , 0.6875 , ..., 0.5337 , 0.806 , 0.934 ], [0.958 , 0.4932 , 0.4885 , ..., 0.7573 , 0.05563 , 0.291 ], [0.2291 , 0.864 , 0.4624 , ..., 0.285 , 0.3137 , 0.4717 ], [0.4517 , 0.755 , 0.1288 , ..., 0.8496 , 0.858 , 0.838 ], [0.07965 , 0.813 , 0.6484 , ..., 0.637 , 0.8574 , 0.3298 ]]]] Outputs: Y: shape=(2, 3, 4, 8), dtype=float16 [[[[0.4634, 0.3958, 0.698 , ..., 0.7026, 0.4287, 0.576 ], [0.449 , 0.3945, 0.6997, ..., 0.725 , 0.4514, 0.5586], [0.458 , 0.394 , 0.684 , ..., 0.7275, 0.466 , 0.5444], [0.4526, 0.3914, 0.7095, ..., 0.7305, 0.4438, 0.5566]], [[0.6167, 0.5005, 0.366 , ..., 0.3538, 0.4849, 0.5576], [0.6084, 0.5273, 0.372 , ..., 0.3508, 0.4912, 0.5522], [0.613 , 0.5356, 0.3718, ..., 0.3594, 0.4858, 0.5713], [0.6074, 0.502 , 0.3713, ..., 0.3584, 0.4749, 0.552 ]], [[0.509 , 0.5264, 0.2328, ..., 0.679 , 0.4148, 0.7964], [0.512 , 0.531 , 0.217 , ..., 0.6816, 0.4224, 0.8027], [0.5093, 0.5264, 0.228 , ..., 0.6885, 0.4185, 0.792 ], [0.5093, 0.5566, 0.2291, ..., 0.671 , 0.4097, 0.788 ]]], [[[0.4038, 0.4631, 0.42 , ..., 0.4077, 0.4414, 0.405 ], [0.4033, 0.448 , 0.425 , ..., 0.414 , 0.4512, 0.4102], [0.4038, 0.4336, 0.4287, ..., 0.429 , 0.4556, 0.418 ], [0.4204, 0.4429, 0.4329, ..., 0.4365, 0.4502, 0.4038]], [[0.5933, 0.4773, 0.4858, ..., 0.4468, 0.4697, 0.4766], [0.5947, 0.4795, 0.498 , ..., 0.4353, 0.4797, 0.4631], [0.6094, 0.4758, 0.5244, ..., 0.4106, 0.4888, 0.4495], [0.5913, 0.4678, 0.509 , ..., 0.4304, 0.4807, 0.4573]], [[0.358 , 0.552 , 0.5205, ..., 0.6343, 0.546 , 0.4905], [0.3645, 0.5645, 0.5103, ..., 0.6333, 0.5483, 0.5015], [0.353 , 0.5474, 0.5225, ..., 0.6353, 0.552 , 0.4941], [0.348 , 0.556 , 0.515 , ..., 0.651 , 0.5835, 0.5044]]]] **test_cc_attention_4d_gqa** .. code-block:: text Node: Attention(Q, K, V) -> (Y) .. 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. ]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.02356532, 0.6783799 ], [-0.02356531, 0.6783799 ]], [[-0.03533878, 0.6841799 ], [ 0.11724145, 0.6063233 ]], [[ 0.6482418 , -0.37858847], [ 0.9917567 , -0.74587834]], [[ 0.37784207, -0.12898168], [ 0.29831943, -0.26321504]]]] **test_cc_attention_4d_gqa_attn_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 .. 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. ]]]] attn_mask: shape=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[ 0.30531266, 0.50214726], [ 0.2229336 , 0.5335671 ]], [[ 0.29782698, 0.50647473], [ 0.31907293, 0.4797093 ]], [[ 0.9921614 , -0.7460807 ], [ 1.1873223 , -1.0039468 ]], [[ 0.74361753, -0.47596362], [ 0.6715176 , -0.5857588 ]]]] **test_cc_attention_4d_gqa_causal** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: is_causal = 1 .. code-block:: text Inputs: Q: shape=(1, 4, 3, 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 ], [ 0.25, 0.1 ]], [[-0.5 , 0.5 ], [-0.25, 0.75], [ 0.1 , -0.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. ]]]] Outputs: Y: shape=(1, 4, 3, 2), dtype=float32 [[[[ 1. , 0. ], [ 0.4911621 , 0.50883794], [-0.03533878, 0.6841799 ]], [[ 1. , 0. ], [ 0.5 , 0.5 ], [-0.23092115, 0.77539366]], [[ 2. , -2. ], [ 0.9953577 , -0.4930365 ], [ 0.60666823, -0.46362755]], [[ 2. , -2. ], [ 1.3812186 , -1.0718278 ], [ 0.5847607 , -0.50853795]]]] **test_cc_attention_4d_gqa_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 0.009999999776482582 .. 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. ]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-3.3333333e-04, 6.6683334e-01], [-3.3333330e-04, 6.6683334e-01]], [[-4.9999997e-04, 6.6691661e-01], [ 1.6666650e-03, 6.6583300e-01]], [[ 6.6659433e-01, -5.8079195e-01], [ 6.7248535e-01, -5.8624268e-01]], [[ 6.6071773e-01, -5.7536310e-01], [ 6.6077113e-01, -5.7789290e-01]]]] **test_cc_attention_4d_gqa_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 1.0 softcap = 0.5 .. 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. ]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.03040495, 0.6819007 ], [-0.02114749, 0.67739695]], [[-0.04947086, 0.69103485], [ 0.15203254, 0.5866078 ]], [[ 0.67558366, -0.4552321 ], [ 0.94510204, -0.72255105]], [[ 0.40459138, -0.23742941], [ 0.35040644, -0.32931072]]]] **test_cc_attention_4d_gqa_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) .. 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. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 4, 2, 2), dtype=float32 [[[[-0.12450287, 0.51090705], [-0.12746866, 0.51840895]], [[-0.13178737, 0.5070071 ], [ 0.00804363, 0.47238812]], [[ 0.48057237, -0.20996696], [ 0.7363033 , -0.5218315 ]], [[ 0.27224937, 0.02852735], [ 0.23465633, -0.06618155]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_gqa_with_past_and_present_fp16** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value) .. 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 ]]]] attn_mask: shape=(4, 18), dtype=float16 [[0.1581 , 0.3494 , 0.621 , ..., 0.509 , 0.999 , 0.0327 ], [0.4724 , 0.548 , 0.07904 , ..., 0.764 , 0.4182 , 0.7754 ], [0.39 , 0.1998 , 0.6245 , ..., 0.004356, 0.9424 , 0.09125 ], [0.673 , 0.6772 , 0.1316 , ..., 0.4263 , 0.4575 , 0.7085 ]] past_key: shape=(2, 3, 12, 8), dtype=float16 [[[[0.5303 , 0.1553 , 0.91 , ..., 0.3093 , 0.8994 , 0.2177 ], [0.3384 , 0.662 , 0.66 , ..., 0.573 , 0.0795 , 0.928 ], [0.534 , 0.2969 , 0.188 , ..., 0.6865 , 0.2976 , 0.03244 ], ..., [0.9775 , 0.08466 , 0.649 , ..., 0.0716 , 0.2751 , 0.7896 ], [0.8716 , 0.4856 , 0.07556 , ..., 0.383 , 0.003033, 0.71 ], [0.912 , 0.622 , 0.933 , ..., 0.0661 , 0.3047 , 0.7373 ]], [[0.775 , 0.03592 , 0.011 , ..., 0.01388 , 0.54 , 0.3003 ], [0.777 , 0.5967 , 0.3645 , ..., 0.6978 , 0.5845 , 0.53 ], [0.3755 , 0.001608, 0.8145 , ..., 0.6665 , 0.3289 , 0.7603 ], ..., [0.2415 , 0.002747, 0.394 , ..., 0.9844 , 0.2812 , 0.1603 ], [0.9604 , 0.6587 , 0.03964 , ..., 0.3857 , 0.5825 , 0.8164 ], [0.2223 , 0.2042 , 0.1219 , ..., 0.03165 , 0.807 , 0.1245 ]], [[0.3904 , 0.1189 , 0.618 , ..., 0.753 , 0.503 , 0.9507 ], [0.4692 , 0.2097 , 0.7705 , ..., 0.958 , 0.732 , 0.8633 ], [0.604 , 0.1118 , 0.588 , ..., 0.1339 , 0.573 , 0.6343 ], ..., [0.8223 , 0.4124 , 0.403 , ..., 0.846 , 0.493 , 0.575 ], [0.499 , 0.04944 , 0.713 , ..., 0.6885 , 0.938 , 0.381 ], [0.3647 , 0.981 , 0.7153 , ..., 0.899 , 0.906 , 0.634 ]]], [[[0.4712 , 0.623 , 0.5747 , ..., 0.245 , 0.1318 , 0.637 ], [0.51 , 0.7207 , 0.7217 , ..., 0.07306 , 0.8657 , 0.3967 ], [0.313 , 0.846 , 0.3518 , ..., 0.8076 , 0.3745 , 0.1225 ], ..., [0.8135 , 0.532 , 0.52 , ..., 0.3748 , 0.2534 , 0.503 ], [0.5806 , 0.877 , 0.939 , ..., 0.010124, 0.4126 , 0.8726 ], [0.4336 , 0.1708 , 0.557 , ..., 0.3865 , 0.9443 , 0.7344 ]], [[0.1107 , 0.767 , 0.725 , ..., 0.209 , 0.771 , 0.8286 ], [0.9917 , 0.348 , 0.0619 , ..., 0.7617 , 0.3303 , 0.5996 ], [0.5967 , 0.4238 , 0.95 , ..., 0.4062 , 0.0765 , 0.694 ], ..., [0.3086 , 0.581 , 0.588 , ..., 0.2893 , 0.9067 , 0.4124 ], [0.299 , 0.3604 , 0.562 , ..., 0.0693 , 0.3503 , 0.7603 ], [0.522 , 0.9927 , 0.4907 , ..., 0.213 , 0.266 , 0.895 ]], [[0.4028 , 0.7695 , 0.913 , ..., 0.3638 , 0.7637 , 0.9946 ], [0.00923 , 0.3105 , 0.635 , ..., 0.7427 , 0.469 , 0.9404 ], [0.5483 , 0.063 , 0.2261 , ..., 0.649 , 0.0228 , 0.8174 ], ..., [0.4016 , 0.7773 , 0.4514 , ..., 0.3088 , 0.797 , 0.7725 ], [0.7485 , 0.718 , 0.4255 , ..., 0.5845 , 0.09503 , 0.512 ], [0.7236 , 0.6655 , 0.36 , ..., 0.09753 , 0.215 , 0.5835 ]]]] past_value: shape=(2, 3, 12, 8), dtype=float16 [[[[0.608 , 0.9365 , 0.646 , ..., 0.83 , 0.3618 , 0.2002 ], [0.1482 , 0.741 , 0.8584 , ..., 0.368 , 0.3079 , 0.9976 ], [0.6387 , 0.824 , 0.777 , ..., 0.1534 , 0.9316 , 0.599 ], ..., [0.608 , 0.4153 , 0.738 , ..., 0.949 , 0.0758 , 0.6973 ], [0.05075, 0.3218 , 0.04764, ..., 0.5913 , 0.5225 , 0.9336 ], [0.329 , 0.7646 , 0.171 , ..., 0.8096 , 0.669 , 0.00903]], [[0.1675 , 0.921 , 0.1498 , ..., 0.0378 , 0.908 , 0.7837 ], [0.4946 , 0.6074 , 0.3533 , ..., 0.601 , 0.2115 , 0.6333 ], [0.9414 , 0.781 , 0.9795 , ..., 0.2142 , 0.6284 , 0.9487 ], ..., [0.11365, 0.607 , 0.717 , ..., 0.0315 , 0.268 , 0.3342 ], [0.6226 , 0.02585, 0.2186 , ..., 0.6733 , 0.742 , 0.3315 ], [0.0741 , 0.378 , 0.6885 , ..., 0.8813 , 0.7886 , 0.5176 ]], [[0.631 , 0.527 , 0.1058 , ..., 0.7773 , 0.739 , 0.94 ], [0.658 , 0.9775 , 0.7773 , ..., 0.373 , 0.5015 , 0.1617 ], [0.2258 , 0.3743 , 0.838 , ..., 0.072 , 0.06244, 0.7617 ], ..., [0.7495 , 0.3655 , 0.4368 , ..., 0.7153 , 0.6587 , 0.4268 ], [0.4885 , 0.4294 , 0.8574 , ..., 0.6147 , 0.5117 , 0.755 ], [0.5273 , 0.07324, 0.651 , ..., 0.3203 , 0.06027, 0.872 ]]], [[[0.4666 , 0.848 , 0.733 , ..., 0.6416 , 0.5728 , 0.01241], [0.1087 , 0.946 , 0.3486 , ..., 0.2783 , 0.8955 , 0.927 ], [0.961 , 0.6987 , 0.5703 , ..., 0.3208 , 0.9526 , 0.381 ], ..., [0.985 , 0.784 , 0.1171 , ..., 0.5947 , 0.967 , 0.318 ], [0.11755, 0.6265 , 0.1833 , ..., 0.00925, 0.55 , 0.548 ], [0.1578 , 0.945 , 0.2764 , ..., 0.7207 , 0.2021 , 0.4163 ]], [[0.1431 , 0.978 , 0.25 , ..., 0.2252 , 0.1288 , 0.3447 ], [0.3142 , 0.6055 , 0.2864 , ..., 0.9106 , 0.1136 , 0.272 ], [0.6147 , 0.847 , 0.396 , ..., 0.554 , 0.01836, 0.1515 ], ..., [0.2874 , 0.9575 , 0.9766 , ..., 0.2634 , 0.9106 , 0.3525 ], [0.978 , 0.71 , 0.4148 , ..., 0.4512 , 0.576 , 0.05423], [0.4338 , 0.08575, 0.784 , ..., 0.676 , 0.933 , 0.494 ]], [[0.1296 , 0.01025, 0.8755 , ..., 0.1215 , 0.1252 , 0.281 ], [0.7437 , 0.05063, 0.8096 , ..., 0.8633 , 0.8735 , 0.5415 ], [0.6567 , 0.9146 , 0.533 , ..., 0.3657 , 0.9775 , 0.416 ], ..., [0.1262 , 0.978 , 0.2334 , ..., 0.3152 , 0.2095 , 0.8735 ], [0.539 , 0.775 , 0.77 , ..., 0.5337 , 0.308 , 0.617 ], [0.8247 , 0.5903 , 0.3923 , ..., 0.78 , 0.9717 , 0.4485 ]]]] Outputs: Y: shape=(2, 9, 4, 8), dtype=float16 [[[[0.4565, 0.4722, 0.523 , ..., 0.5005, 0.4812, 0.5786], [0.458 , 0.4365, 0.5674, ..., 0.5024, 0.5005, 0.6333], [0.4712, 0.4397, 0.572 , ..., 0.5312, 0.4714, 0.591 ], [0.4604, 0.448 , 0.5513, ..., 0.5303, 0.4814, 0.601 ]], [[0.4646, 0.4583, 0.5293, ..., 0.5044, 0.4863, 0.581 ], [0.455 , 0.4417, 0.5703, ..., 0.513 , 0.487 , 0.628 ], [0.4663, 0.4407, 0.5615, ..., 0.529 , 0.4727, 0.5933], [0.479 , 0.4465, 0.554 , ..., 0.5254, 0.4841, 0.581 ]], [[0.4517, 0.4746, 0.525 , ..., 0.507 , 0.478 , 0.571 ], [0.459 , 0.4436, 0.5767, ..., 0.5186, 0.4827, 0.6274], [0.4695, 0.4456, 0.5645, ..., 0.5303, 0.4695, 0.588 ], [0.4832, 0.456 , 0.5654, ..., 0.533 , 0.4768, 0.5767]], ..., [[0.5464, 0.4841, 0.568 , ..., 0.4487, 0.5005, 0.618 ], [0.561 , 0.4775, 0.5386, ..., 0.4868, 0.5537, 0.5376], [0.5625, 0.4692, 0.5664, ..., 0.4783, 0.517 , 0.585 ], [0.5317, 0.4702, 0.552 , ..., 0.487 , 0.5317, 0.5674]], [[0.5435, 0.4827, 0.5664, ..., 0.4473, 0.498 , 0.613 ], [0.5713, 0.4663, 0.55 , ..., 0.4805, 0.5615, 0.53 ], [0.567 , 0.4734, 0.5645, ..., 0.4763, 0.5137, 0.587 ], [0.5337, 0.4688, 0.547 , ..., 0.4805, 0.5396, 0.5693]], [[0.553 , 0.4917, 0.5664, ..., 0.4568, 0.5054, 0.613 ], [0.567 , 0.4707, 0.5464, ..., 0.4817, 0.5566, 0.54 ], [0.5703, 0.466 , 0.567 , ..., 0.4705, 0.5186, 0.5737], [0.5293, 0.471 , 0.543 , ..., 0.4805, 0.5312, 0.5728]]], [[[0.46 , 0.603 , 0.424 , ..., 0.4155, 0.527 , 0.458 ], [0.483 , 0.6294, 0.4792, ..., 0.416 , 0.6055, 0.444 ], [0.4775, 0.6045, 0.4565, ..., 0.4143, 0.5664, 0.4448], [0.4692, 0.6245, 0.4358, ..., 0.4404, 0.598 , 0.4285]], [[0.4578, 0.599 , 0.4304, ..., 0.4202, 0.53 , 0.4531], [0.4937, 0.62 , 0.4722, ..., 0.424 , 0.611 , 0.4446], [0.475 , 0.6143, 0.4624, ..., 0.4182, 0.5596, 0.4385], [0.4631, 0.626 , 0.4348, ..., 0.4346, 0.5884, 0.4365]], [[0.465 , 0.6035, 0.4255, ..., 0.4248, 0.5234, 0.4507], [0.4963, 0.626 , 0.4705, ..., 0.4219, 0.6035, 0.444 ], [0.4622, 0.615 , 0.4634, ..., 0.4053, 0.5674, 0.4463], [0.453 , 0.623 , 0.4297, ..., 0.4285, 0.5957, 0.4434]], ..., [[0.5327, 0.534 , 0.5625, ..., 0.55 , 0.5537, 0.4788], [0.4932, 0.507 , 0.556 , ..., 0.5186, 0.5034, 0.519 ], [0.5234, 0.5312, 0.5586, ..., 0.562 , 0.5537, 0.493 ], [0.4788, 0.4834, 0.611 , ..., 0.4995, 0.466 , 0.526 ]], [[0.5376, 0.542 , 0.5605, ..., 0.548 , 0.5537, 0.4817], [0.4856, 0.4968, 0.557 , ..., 0.51 , 0.4897, 0.5293], [0.525 , 0.521 , 0.55 , ..., 0.555 , 0.546 , 0.494 ], [0.4934, 0.511 , 0.613 , ..., 0.514 , 0.4817, 0.5186]], [[0.5312, 0.538 , 0.5625, ..., 0.55 , 0.5547, 0.4792], [0.4773, 0.4963, 0.5547, ..., 0.5195, 0.4883, 0.5244], [0.5156, 0.548 , 0.5483, ..., 0.5635, 0.547 , 0.506 ], [0.51 , 0.5337, 0.582 , ..., 0.5166, 0.4783, 0.5244]]]] present_key: shape=(2, 3, 18, 8), dtype=float16 [[[[5.3027e-01, 1.5527e-01, 9.1016e-01, ..., 3.0933e-01, 8.9941e-01, 2.1765e-01], [3.3838e-01, 6.6211e-01, 6.6016e-01, ..., 5.7324e-01, 7.9529e-02, 9.2822e-01], [5.3418e-01, 2.9688e-01, 1.8799e-01, ..., 6.8652e-01, 2.9761e-01, 3.2440e-02], ..., [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]], [[7.7490e-01, 3.5919e-02, 1.1002e-02, ..., 1.3878e-02, 5.4004e-01, 3.0029e-01], [7.7686e-01, 5.9668e-01, 3.6450e-01, ..., 6.9775e-01, 5.8447e-01, 5.2979e-01], [3.7549e-01, 1.6079e-03, 8.1445e-01, ..., 6.6650e-01, 3.2886e-01, 7.6025e-01], ..., [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.9038e-01, 1.1890e-01, 6.1816e-01, ..., 7.5293e-01, 5.0293e-01, 9.5068e-01], [4.6924e-01, 2.0972e-01, 7.7051e-01, ..., 9.5801e-01, 7.3193e-01, 8.6328e-01], [6.0400e-01, 1.1182e-01, 5.8789e-01, ..., 1.3391e-01, 5.7324e-01, 6.3428e-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]]], [[[4.7119e-01, 6.2305e-01, 5.7471e-01, ..., 2.4500e-01, 1.3184e-01, 6.3721e-01], [5.0977e-01, 7.2070e-01, 7.2168e-01, ..., 7.3059e-02, 8.6572e-01, 3.9673e-01], [3.1299e-01, 8.4619e-01, 3.5181e-01, ..., 8.0762e-01, 3.7451e-01, 1.2250e-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]], [[1.1072e-01, 7.6709e-01, 7.2510e-01, ..., 2.0898e-01, 7.7100e-01, 8.2861e-01], [9.9170e-01, 3.4790e-01, 6.1890e-02, ..., 7.6172e-01, 3.3032e-01, 5.9961e-01], [5.9668e-01, 4.2383e-01, 9.5020e-01, ..., 4.0625e-01, 7.6477e-02, 6.9385e-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]], [[4.0283e-01, 7.6953e-01, 9.1309e-01, ..., 3.6377e-01, 7.6367e-01, 9.9463e-01], [9.2316e-03, 3.1055e-01, 6.3477e-01, ..., 7.4268e-01, 4.6899e-01, 9.4043e-01], [5.4834e-01, 6.2988e-02, 2.2607e-01, ..., 6.4893e-01, 2.2797e-02, 8.1738e-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]]]] present_value: shape=(2, 3, 18, 8), dtype=float16 [[[[0.608 , 0.9365 , 0.646 , ..., 0.83 , 0.3618 , 0.2002 ], [0.1482 , 0.741 , 0.8584 , ..., 0.368 , 0.3079 , 0.9976 ], [0.6387 , 0.824 , 0.777 , ..., 0.1534 , 0.9316 , 0.599 ], ..., [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.1675 , 0.921 , 0.1498 , ..., 0.0378 , 0.908 , 0.7837 ], [0.4946 , 0.6074 , 0.3533 , ..., 0.601 , 0.2115 , 0.6333 ], [0.9414 , 0.781 , 0.9795 , ..., 0.2142 , 0.6284 , 0.9487 ], ..., [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.631 , 0.527 , 0.1058 , ..., 0.7773 , 0.739 , 0.94 ], [0.658 , 0.9775 , 0.7773 , ..., 0.373 , 0.5015 , 0.1617 ], [0.2258 , 0.3743 , 0.838 , ..., 0.072 , 0.06244 , 0.7617 ], ..., [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.4666 , 0.848 , 0.733 , ..., 0.6416 , 0.5728 , 0.01241 ], [0.1087 , 0.946 , 0.3486 , ..., 0.2783 , 0.8955 , 0.927 ], [0.961 , 0.6987 , 0.5703 , ..., 0.3208 , 0.9526 , 0.381 ], ..., [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.1431 , 0.978 , 0.25 , ..., 0.2252 , 0.1288 , 0.3447 ], [0.3142 , 0.6055 , 0.2864 , ..., 0.9106 , 0.1136 , 0.272 ], [0.6147 , 0.847 , 0.396 , ..., 0.554 , 0.01836 , 0.1515 ], ..., [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.1296 , 0.01025 , 0.8755 , ..., 0.1215 , 0.1252 , 0.281 ], [0.7437 , 0.05063 , 0.8096 , ..., 0.8633 , 0.8735 , 0.5415 ], [0.6567 , 0.9146 , 0.533 , ..., 0.3657 , 0.9775 , 0.416 ], ..., [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 ]]]] **test_cc_attention_4d_scaled** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 0.009999999776482582 .. 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, 2, 2), dtype=float32 [[[[ 1., 0.], [ 0., 1.]], [[ 1., 1.], [-1., 1.]]]] V: shape=(1, 2, 2, 2), dtype=float32 [[[[ 1., 2.], [ 3., 4.]], [[-1., 0.], [ 0., 1.]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 1.995 , 2.9950001 ], [ 2.0049999 , 3.0049999 ]], [[-0.5025 , 0.49750003], [-0.5049998 , 0.49500015]]]] **test_cc_attention_4d_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y) Attributes: scale = 1.0 softcap = 0.5 .. 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. ]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.15169363, 0.6033754 ], [-0.15169363, 0.755069 ]], [[ 0.67558366, -0.4552321 ], [ 0.29168454, -0.3030643 ]]]] **test_cc_attention_4d_softcap_neginf_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: softcap = 0.5 .. code-block:: text Inputs: Q: shape=(1, 1, 4, 8), dtype=float32 [[[[0.7167289 , 0.6484601 , 0.46497524, 0.34117728, 0.55324554, 0.47324938, 0.8506881 , 0.04282242], [0.4230454 , 0.07251757, 0.7479862 , 0.1587475 , 0.6062188 , 0.806189 , 0.17963564, 0.45366198], [0.13773894, 0.34766454, 0.74359787, 0.6698144 , 0.26058674, 0.5744466 , 0.1911546 , 0.04261738], [0.12887597, 0.70876426, 0.2550329 , 0.06627256, 0.7609782 , 0.708198 , 0.7275301 , 0.50939566]]]] K: shape=(1, 1, 6, 8), dtype=float32 [[[[0.81594235, 0.868624 , 0.75029296, 0.81736696, 0.16860753, 0.88723135, 0.19616836, 0.57798064], [0.65411955, 0.7793319 , 0.9017522 , 0.47066295, 0.6826957 , 0.4715004 , 0.6312656 , 0.7740394 ], [0.3452173 , 0.23894429, 0.18115634, 0.68568075, 0.17885476, 0.65862894, 0.77072376, 0.40022814], [0.18269259, 0.8124025 , 0.6827437 , 0.437518 , 0.08335429, 0.7401209 , 0.15409368, 0.22070706], [0.8506275 , 0.14277291, 0.5163776 , 0.90439737, 0.4630888 , 0.33560514, 0.58655113, 0.4527613 ], [0.923729 , 0.45124698, 0.30754632, 0.9676665 , 0.52557784, 0.8928356 , 0.6388969 , 0.266801 ]]]] V: shape=(1, 1, 6, 8), dtype=float32 [[[[0.91515577, 0.08878785, 0.03561068, 0.29355663, 0.7839695 , 0.30121332, 0.5416486 , 0.11313885], [0.8851937 , 0.48614627, 0.05551815, 0.7825784 , 0.75917256, 0.1368118 , 0.08289552, 0.0944168 ], [0.55269563, 0.13022405, 0.6187148 , 0.7015471 , 0.09712279, 0.74281126, 0.35029292, 0.7578389 ], [0.2365092 , 0.9160407 , 0.11045456, 0.80876344, 0.40573037, 0.77204376, 0.58065724, 0.93201846], [0.45860708, 0.3792408 , 0.2441163 , 0.12815326, 0.74603045, 0.76796633, 0.5889299 , 0.96119374], [0.4440869 , 0.19609386, 0.56120396, 0.44926733, 0.1321832 , 0.664661 , 0.26817727, 0.20522803]]]] attn_mask: shape=(4, 6), dtype=float32 [[ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf]] Outputs: Y: shape=(1, 1, 4, 8), dtype=float32 [[[[0.65484625, 0.40108806, 0.20163937, 0.64396936, 0.5180408 , 0.48061964, 0.38587224, 0.4638726 ], [0.6566198 , 0.40271884, 0.19770749, 0.64299566, 0.5223832 , 0.47754967, 0.38596648, 0.46013588], [0.65369874, 0.40479586, 0.19766673, 0.6429548 , 0.5210476 , 0.48034573, 0.38840714, 0.4638093 ], [0.65444815, 0.40250054, 0.20119615, 0.64485633, 0.518264 , 0.48043442, 0.38539407, 0.46407378]]]] **test_cc_attention_4d_softcap_neginf_mask_poison** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: softcap = 0.5 .. code-block:: text Inputs: Q: shape=(1, 1, 4, 8), dtype=float32 [[[[0.7167289 , 0.6484601 , 0.46497524, 0.34117728, 0.55324554, 0.47324938, 0.8506881 , 0.04282242], [0.4230454 , 0.07251757, 0.7479862 , 0.1587475 , 0.6062188 , 0.806189 , 0.17963564, 0.45366198], [0.13773894, 0.34766454, 0.74359787, 0.6698144 , 0.26058674, 0.5744466 , 0.1911546 , 0.04261738], [0.12887597, 0.70876426, 0.2550329 , 0.06627256, 0.7609782 , 0.708198 , 0.7275301 , 0.50939566]]]] K: shape=(1, 1, 6, 8), dtype=float32 [[[[0.81594235, 0.868624 , 0.75029296, 0.81736696, 0.16860753, 0.88723135, 0.19616836, 0.57798064], [0.65411955, 0.7793319 , 0.9017522 , 0.47066295, 0.6826957 , 0.4715004 , 0.6312656 , 0.7740394 ], [0.3452173 , 0.23894429, 0.18115634, 0.68568075, 0.17885476, 0.65862894, 0.77072376, 0.40022814], [0.18269259, 0.8124025 , 0.6827437 , 0.437518 , 0.08335429, 0.7401209 , 0.15409368, 0.22070706], [0.8506275 , 0.14277291, 0.5163776 , 0.90439737, 0.4630888 , 0.33560514, 0.58655113, 0.4527613 ], [0.923729 , 0.45124698, 0.30754632, 0.9676665 , 0.52557784, 0.8928356 , 0.6388969 , 0.266801 ]]]] V: shape=(1, 1, 6, 8), dtype=float32 [[[[9.15155768e-01, 8.87878537e-02, 3.56106758e-02, 2.93556631e-01, 7.83969522e-01, 3.01213324e-01, 5.41648626e-01, 1.13138855e-01], [8.85193706e-01, 4.86146271e-01, 5.55181503e-02, 7.82578409e-01, 7.59172559e-01, 1.36811793e-01, 8.28955173e-02, 9.44167972e-02], [5.52695632e-01, 1.30224049e-01, 6.18714809e-01, 7.01547086e-01, 9.71227884e-02, 7.42811263e-01, 3.50292921e-01, 7.57838905e-01], [2.36509204e-01, 9.16040719e-01, 1.10454559e-01, 8.08763444e-01, 4.05730367e-01, 7.72043765e-01, 5.80657244e-01, 9.32018459e-01], [1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03], [1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03]]]] attn_mask: shape=(4, 6), dtype=float32 [[ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf], [ 0., 0., 0., 0., -inf, -inf]] Outputs: Y: shape=(1, 1, 4, 8), dtype=float32 [[[[0.65484625, 0.40108806, 0.20163937, 0.64396936, 0.5180408 , 0.48061964, 0.38587224, 0.4638726 ], [0.6566198 , 0.40271884, 0.19770749, 0.64299566, 0.5223832 , 0.47754967, 0.38596648, 0.46013588], [0.65369874, 0.40479586, 0.19766673, 0.6429548 , 0.5210476 , 0.48034573, 0.38840714, 0.4638093 ], [0.65444815, 0.40250054, 0.20119615, 0.64485633, 0.518264 , 0.48043442, 0.38539407, 0.46407378]]]] **test_cc_attention_4d_with_past_and_present** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value) .. 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. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.10764056, 0.45608166], [-0.31939295, 0.57814157]], [[ 0.48057237, -0.20996696], [ 0.05822894, 0.11373253]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul** .. code-block:: text Node: Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) .. 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. ]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.10764056, 0.45608166], [-0.31939295, 0.57814157]], [[ 0.48057237, -0.20996696], [ 0.05822894, 0.11373253]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.35355338, 0. , 0.70710677, 0.35355338, 0. ], [-0.35355338, 0.35355338, 0. , 0.35355338, 0.70710677]], [[ 0.35355338, 0.17677669, 0. , 0.70710677, -0.08838835], [ 0.70710677, -1.0606601 , -1.4142135 , 0. , 0.53033006]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul_bias** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: scale = 0.5 qk_matmul_output_mode = 1 .. 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=(2, 5), dtype=float32 [[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[-0.07067307, 0.63377535], [-0.23752189, 0.5571051 ]], [[ 0.30545416, 0.01247976], [ 0.09629637, 0.127618 ]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: scale = 0.5 qk_matmul_output_mode = 1 .. 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, 5), dtype=float32 [[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[-0.07067307, 0.63377535], [-0.23752189, 0.5571051 ]], [[ 0.30545416, 0.01247976], [ 0.09629637, 0.127618 ]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: scale = 0.5 is_causal = 1 qk_matmul_output_mode = 1 .. 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, 5), dtype=float32 [[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.2571047 , 0.2571047 ], [ 0.03885095, 0.39657053]], [[ 0.46144247, -0.18698354], [ 0.29362845, 0.16985056]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: scale = 0.5 qk_matmul_output_mode = 1 .. 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, 5), dtype=float32 [[[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]], [[ 0.1, 0.2, -0.3, 0. , 0.4], [-0.4, 0. , 0.3, -0.2, 0.1]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[-0.07067307, 0.63377535], [-0.23752189, 0.5571051 ]], [[ 0.34838727, -0.15952764], [ 0.18545583, -0.10488772]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal** .. code-block:: text Node: Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output) Attributes: scale = 0.5 is_causal = 1 qk_matmul_output_mode = 1 .. 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, 5), dtype=float32 [[[[ 0. , -0.5, -1. , 0.2, 0. ], [ 0.5, 0. , -0.2, -0.1, 0. ]], [[ 0.1, 0.2, -0.3, 0. , 0.4], [-0.4, 0. , 0.3, -0.2, 0.1]]]] past_key: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, -0.5], [ 0. , 0.5]], [[ 1. , 0. ], [-0.5, 1. ]]]] past_value: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.5, 0.5], [-1. , 0. ]], [[ 0. , 0.5], [ 0.5, -0.5]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.2571047 , 0.2571047 ], [ 0.03885095, 0.39657053]], [[ 0.61334735, -0.4131356 ], [ 0.5664716 , -0.16319034]]]] present_key: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , -0.5 ], [ 0. , 0.5 ], [ 1. , 0. ], [ 0.5 , 0.5 ], [ 0. , 1. ]], [[ 1. , 0. ], [-0.5 , 1. ], [-1. , 1. ], [ 1. , 1. ], [ 0.25, -0.5 ]]]] present_value: shape=(1, 2, 5, 2), dtype=float32 [[[[ 0.5 , 0.5 ], [-1. , 0. ], [ 1. , 0. ], [ 0. , 1. ], [-1. , 1. ]], [[ 0. , 0.5 ], [ 0.5 , -0.5 ], [ 2. , -2. ], [ 0.5 , 0.25], [-0.5 , 0. ]]]] qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32 [[[[ 0.25 , 0. , 0.5 , 0.25 , 0. ], [-0.25 , 0.25 , 0. , 0.25 , 0.5 ]], [[ 0.25 , 0.125 , 0. , 0.5 , -0.0625], [ 0.5 , -0.75 , -1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_qk_matmul** .. code-block:: text Node: Attention(Q, K, V) -> (Y, "", "", qk_matmul_output) .. 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. ]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.23092115, 0.5444725 ], [-0.23092115, 0.77539366]], [[ 0.6482418 , -0.37858847], [ 0.04638294, -0.08028026]]]] qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.70710677, 0.35355338, 0. ], [ 0. , 0.35355338, 0.70710677]], [[ 0. , 0.70710677, -0.08838835], [-1.4142135 , 0. , 0.53033006]]]] **test_cc_attention_4d_with_qk_matmul_bias** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output) Attributes: scale = 0.5 qk_matmul_output_mode = 1 .. 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=(2, 3), dtype=float32 [[ 0. , -0.5, -1. ], [ 0.5, 0. , -0.2]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.458196 , 0.41020232], [ 0.0697852 , 0.6150191 ]], [[ 0.9921614 , -0.7460807 ], [ 0.3994022 , -0.34422377]]]] qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.5 , 0.25 , 0. ], [ 0. , 0.25 , 0.5 ]], [[ 0. , 0.5 , -0.0625], [-1. , 0. , 0.375 ]]]] **test_cc_attention_4d_with_qk_matmul_softcap** .. code-block:: text Node: Attention(Q, K, V) -> (Y, "", "", qk_matmul_output) Attributes: scale = 1.0 softcap = 0.5 qk_matmul_output_mode = 2 .. 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. ]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.15169363, 0.6033754 ], [-0.15169363, 0.755069 ]], [[ 0.67558366, -0.4552321 ], [ 0.29168454, -0.3030643 ]]]] qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32 [[[[ 0.4820138 , 0.3807971 , 0. ], [ 0. , 0.3807971 , 0.4820138 ]], [[ 0. , 0.4820138 , -0.12245933], [-0.49966466, 0. , 0.45257413]]]] **test_cc_attention_4d_with_qk_matmul_softmax** .. 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, 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. ]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0.23092115, 0.5444725 ], [-0.23092115, 0.77539366]], [[ 0.6482418 , -0.37858847], [ 0.04638294, -0.08028026]]]] qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32 [[[[0.45552748, 0.31986618, 0.22460635], [0.22460635, 0.31986618, 0.45552748]], [[0.25358054, 0.5142905 , 0.23212898], [0.08261943, 0.33983436, 0.5775462 ]]]] **test_cc_attention_causal_boolmask_nan_robustness** .. code-block:: text Node: Attention(Q, K, V, attn_mask) -> (Y) Attributes: scale = 0.5 is_causal = 1 .. 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], [ True, True, True]], [[False, False, False], [ True, True, True]]]] Outputs: Y: shape=(1, 2, 2, 2), dtype=float32 [[[[ 0. , 0. ], [ 0.4378235, 0.5621765]], [[ 0. , 0. ], [ 0.9034121, -0.3551182]]]] **test_cc_shape_inference_tiny_llm** .. code-block:: text Node: Attention(query, key, value, attn_bias, past_key, past_value) -> (attn_out, present_key, present_value) Attributes: q_num_heads = 4 kv_num_heads = 4 is_causal = 1