(l-onnx-doc-Attention)= # Attention (l-onnx-op-attention-25)= ## Attention - 25 ### Version - **name**: [Attention (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Attention) - **domain**: `main` - **since_version**: `25` - **function**: `True` - **support_level**: `SupportType.COMMON` - **shape inference**: `True` This version of the operator has been available **since version 25**. ### Summary Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed. This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V. For self attention, `kv_sequence_length` equals to `q_sequence_length`. For cross attention, query and key might have different lengths. This operator also covers the 3 following variants based on the number of heads: 1) Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, `q_num_heads = kv_num_heads`. 2) Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, `q_num_heads > kv_num_heads`, `q_num_heads % kv_num_heads == 0`. 3) Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`. Attention bias to be added is calculated based on `attn_mask` input and `is_causal` attribute: 1) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score. 2) If `is_causal` is set to `1`, causal masking is applied with bottom-right (offset-aware) alignment: query `i` attends key `j` iff `j <= i + offset`, as illustrated below. ``` 2D causal mask for Attention (PR onnx/onnx#8068) S_q=4 queries, S_k=8 keys Rule: query i attends key j iff j <= i + offset offset = nonpad_kv_seqlen - S_q nonpad_kv_seqlen=4, offset=4-4=0 k0 k1 k2 k3 k4 k5 k6 k7 +----+----+----+----+----+----+----+----+ q0 | ## | | | | | | | | +----+----+----+----+----+----+----+----+ q1 | ## | ## | | | | | | | +----+----+----+----+----+----+----+----+ q2 | ## | ## | ## | | | | | | +----+----+----+----+----+----+----+----+ q3 | ## | ## | ## | ## | | | | | +----+----+----+----+----+----+----+----+ nonpad_kv_seqlen=8, offset=8-4=4 k0 k1 k2 k3 k4 k5 k6 k7 +----+----+----+----+----+----+----+----+ q0 | ## | ## | ## | ## | ## | | | | +----+----+----+----+----+----+----+----+ q1 | ## | ## | ## | ## | ## | ## | | | +----+----+----+----+----+----+----+----+ q2 | ## | ## | ## | ## | ## | ## | ## | | +----+----+----+----+----+----+----+----+ q3 | ## | ## | ## | ## | ## | ## | ## | ## | +----+----+----+----+----+----+----+----+ ``` With `nonpad_kv_seqlen=4` (offset=0), the mask is the standard lower-triangular. With `nonpad_kv_seqlen=8` (offset=4), the diagonal shifts right by 4, so each query sees the 4 additional valid cached keys. `offset` is the count of valid keys preceding the current query block: `offset = past_sequence_length` when `past_key` is provided; `offset = nonpad_kv_seqlen - q_sequence_length` (per batch) when an external cache is indicated by `nonpad_kv_seqlen` without `past_key`; `offset = 0` when neither is provided (the no-cache case, which reduces to the standard lower-triangular mask). When `offset < 0` (`nonpad_kv_seqlen < q_sequence_length`, i.e. more query tokens than cached keys) the leading query rows have an empty key set (no key satisfies `j <= i + offset`) and are fully masked. The causal frontier is computed independently of `attn_mask` and is then composed with it additively: a boolean `attn_mask` intersects the allowed set (its disallowed positions contribute `-inf` to the bias), while a float `attn_mask` is added to the attention scores rather than disabling positions. A fully-masked query row (no key attended, including the negative-offset leading rows) produces a zero output row, not `NaN`, for both `Y` and the mode-`3` `qk_matmul_output` debug output; the mode-`3` `qk_matmul_output` is emitted at the operator's output precision (`T1`). `left_window_size` and `right_window_size` independently restrict the keys visible to each query. A query at absolute position `p = offset + query_index` attends keys `j` satisfying `p - left_window_size <= j <= p + right_window_size` for each nonnegative bound. A value of `-1` leaves that side unbounded. For example, `(left_window_size=2, right_window_size=0)` is a causal left-looking window containing the current key and two preceding keys, while `(left_window_size=2, right_window_size=1)` is an asymmetric bidirectional window. Window bounds are composed with `is_causal` and `attn_mask`; when `is_causal=1`, the causal upper bound still excludes future keys. ``` 2D sliding-window mask for Attention (opset 25) S_q=4 queries, S_k=6 keys, left_window_size=2, right_window_size=1, offset=0 k0 k1 k2 k3 k4 k5 +----+----+----+----+----+----+ q0 | ## | ## | | | | | +----+----+----+----+----+----+ q1 | ## | ## | ## | | | | +----+----+----+----+----+----+ q2 | ## | ## | ## | ## | | | +----+----+----+----+----+----+ q3 | | ## | ## | ## | ## | | +----+----+----+----+----+----+ q0 attends {k0,k1}, q1 attends {k0,k1,k2}, q2 attends {k0,k1,k2,k3}, q3 attends {k1,k2,k3,k4}. ``` With respect to KV cache update, this operator allows the following two use cases: 1) Cache update happens inside the Attention operator. In this case, the `K` and `V` inputs contain only the incoming tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present key and value, respectively. Note that this only works correctly for the special case where the past key and value do not contain padded tokens. 2) Cache update happens outside the Attention operator (for example, through the `TensorScatter` operator). In this case, the `K` and `V` inputs correspond to the entire cache tensor, so the four optional inputs/outputs past and present key and value should not be used. An additional input `nonpad_kv_seqlen` of shape (batch_size,) may be provided to indicate the number of non-padding tokens in each sample of the batch to save unnecessary computation. Here, the kv_sequence dimension of `attn_mask` can be shorter than `K` and `V`, but still needs to be at least as long as the maximum value of `nonpad_kv_seqlen`. Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided: ``` 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 ``` ### Attributes * **is_causal - INT** (default is `0`): If set to `1`, causal masking is applied. For a square Q/K (no cache offset) this is a lower-triangular matrix. In general the mask is bottom-right (offset-aware): query in-block index `i` attends key `j` iff `j <= i + offset`, where `offset` is the count of valid keys preceding the query block (`past_sequence_length` for an internal `past_key` cache, or `nonpad_kv_seqlen - q_sequence_length` per batch for an external cache). When `offset = 0` this reduces to the lower-triangular (top-left) mask. * **kv_num_heads - INT** : Number of heads of key and value. Must be used with 3D inputs of Q, K and V. * **left_window_size - INT** (default is `-1`): Maximum number of positions to the left of the current absolute query position that may be attended. A value of `0` allows the current position but no preceding position, while `-1` leaves the left side unbounded. This bound is composed with `is_causal` and `attn_mask`. * **q_num_heads - INT** : Number of heads of query. Must be used with 3D inputs of Q, K and V. * **qk_matmul_output_mode - INT** (default is `0`): Determines what the optional 4th output contains: `0` (default): raw QK matmul result; `1`: after softcap (before bias addition); `2`: QK + softcap + bias; `3`: post-softmax probabilities (after fully-masked-row guard). In mode `3`, a fully-masked query row (every key disallowed) is a zero row, consistent with the corresponding row of the primary output `Y`. The mode-`3` output is emitted at the operator's output precision (`T1`); when `softmax_precision` differs from `T1` this is a cast of the softmax result to `T1`. * **right_window_size - INT** (default is `-1`): Maximum number of positions to the right of the current absolute query position that may be attended. A value of `0` allows the current position but no following position, while `-1` leaves the right side unbounded. Set `is_causal=0` to use a positive right window. * **scale - FLOAT** : Scaling factor applied to $Q*K^T$. Default value is `1/sqrt(head_size)`. To prevent [numerical overflow](https://tinyurl.com/sudb9s96), scale `Q`, `K` by `sqrt(scale)` before matmul. * **softcap - FLOAT** (default is `0.0`): Soft cap for attention logits, applied as `softcap * tanh(logits / softcap)`. Default value of `0.0` means no soft capping is applied. The soft cap is applied before mask / bias addition and softmax. * **softmax_precision - INT** : Specifies the precision for softmax computation. If provided, the attention weights will be cast to this type before softmax and then cast back to the original type. Supported values are: `1` (FLOAT), `10` (FLOAT16), `11` (DOUBLE), `16` (BFLOAT16). ### Inputs Between 3 and 7 inputs. - **Q** (heterogeneous) - **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** (heterogeneous) - **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** (heterogeneous) - **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** (optional, heterogeneous) - **U**: Attention mask. Shape must be broadcastable to `(batch_size, q_num_heads, q_sequence_length, total_sequence_length)` where `total_sequence_length = past_sequence_length + kv_sequence_length`. The last dimension can also be shorter than `total_sequence_length` and will be padded to `total_sequence_length` with negative infinity. Two types of masks are supported: a boolean mask where a value of `True` indicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score. - **past_key** (optional, heterogeneous) - **T1**: Past state for key with shape `(batch_size, kv_num_heads, past_sequence_length, head_size)`. Must be used together with `past_value` input. - **past_value** (optional, heterogeneous) - **T2**: Past state for value with shape `(batch_size, kv_num_heads, past_sequence_length, v_head_size)`. Must be used together with `past_key` input. - **nonpad_kv_seqlen** (optional, heterogeneous) - **tensor(int64)**: A vector of integers of shape `(batch_size,)` that indicates the number of valid (i.e., non-padding) tokens in each sample. A padding mask can be derived from this. This should not be used together with `past_key` and `past_value` inputs or `present_key` and `present_value` outputs (see the KV cache use cases in the operator description). ### Outputs Between 1 and 4 outputs. - **Y** (heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain Q and K inputs types to float tensors. * **T2** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain V input types to float tensors. * **U** in ( `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)` ): Constrain output 'mask' types to boolean tensors and input types. ### Examples #### _attention ```python import numpy as np import onnx node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"]) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_fp16 ```python import numpy as np import onnx node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"]) Q = np.random.rand(2, 3, 4, 8).astype(np.float16) K = np.random.rand(2, 3, 6, 8).astype(np.float16) V = np.random.rand(2, 3, 6, 8).astype(np.float16) Y, _, _, _ = _compute_attention(Q, K, V) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_fp16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa ```python import numpy as np import onnx node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"]) Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_gqa", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes ```python import numpy as np import onnx node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"]) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_diff_heads_sizes", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_scaled ```python import numpy as np import onnx scale = 1e-2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, scale=scale) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_scaled ```python import numpy as np import onnx scale = 1e-2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, ) Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, scale=scale) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_gqa_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_scaled ```python import numpy as np import onnx scale = 1e-2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, scale=scale) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_diff_heads_sizes_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_gqa_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_diff_heads_sizes_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_3d_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_3d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_3d_mask_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_3d_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_4d_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_4d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_4d_mask_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_4d_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_mask_bool ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(bool) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_bool", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_attn_mask_bool_4d ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6).astype(bool) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_bool_4d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_attn_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_gqa_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_attn_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_diff_heads_sizes_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_with_past_and_present ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_gqa_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_with_past_and_present_fp16 ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 9, 4, 8).astype(np.float16) K = np.random.rand(2, 3, 6, 8).astype(np.float16) V = np.random.rand(2, 3, 6, 8).astype(np.float16) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float16) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_gqa_with_past_and_present_fp16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_with_past_and_present ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_diff_heads_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_with_past_and_present_mask3D ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_diff_heads_with_past_and_present_mask3d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_with_past_and_present_mask4D ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_diff_heads_with_past_and_present_mask4d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_softcap ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=2.0, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, softcap=2.0) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_gqa_softcap ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=2.0, ) Q = np.random.rand(2, 9, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention(Q, K, V, softcap=2.0) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_gqa_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_diff_head_sizes_softcap ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=2.0, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, softcap=2.0, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_diff_heads_sizes_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_qk_matmul ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y", "", "", "qk_matmul_output"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, qk_matmul_output = _compute_attention(Q, K, V) expect( node, inputs=[Q, K, V], outputs=[Y, qk_matmul_output], name="test_attention_4d_with_qk_matmul", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_qk_matmul_bias ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], qk_matmul_output_mode=2, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, qk_matmul_output_mode=2, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_4d_with_qk_matmul_bias", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_qk_matmul_softcap ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], softcap=2.0, qk_matmul_output_mode=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, softcap=2.0, qk_matmul_output_mode=1, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_4d_with_qk_matmul_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_qk_matmul_softmax ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], qk_matmul_output_mode=3, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, qk_matmul_output_mode=3, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_4d_with_qk_matmul_softmax", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul_bias ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], qk_matmul_output_mode=2, ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, qk_matmul_output_mode=2, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul_bias", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul_bias_3d_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], qk_matmul_output_mode=2, ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, qk_matmul_output_mode=2, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul_bias_4d_mask ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], qk_matmul_output_mode=2, ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, qk_matmul_output_mode=2, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul_bias_3d_mask_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], qk_matmul_output_mode=2, is_causal=1, ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, qk_matmul_output_mode=2, is_causal=1, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul_bias_4d_mask_causal ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], qk_matmul_output_mode=2, is_causal=1, ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, qk_matmul_output_mode=2, is_causal=1, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_with_past_and_present_qk_matmul ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], ) past_sequence_length = 12 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_4d_with_past_and_present_qk_matmul", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa ```python import numpy as np import onnx q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_gqa", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_diff_heads_sizes", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_scaled ```python import numpy as np import onnx scale = 1e-2 q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa_scaled ```python import numpy as np import onnx scale = 1e-2 q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_gqa_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes_scaled ```python import numpy as np import onnx scale = 1e-2 q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, scale=scale, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_diff_heads_sizes_scaled", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_causal ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa_causal ```python import numpy as np import onnx q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_gqa_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes_causal ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_diff_heads_sizes_causal", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_attn_mask ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_3d_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa_attn_mask ```python import numpy as np import onnx q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_3d_gqa_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes_attn_mask ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) attn_mask = np.random.rand(4, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_3d_diff_heads_sizes_attn_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_softcap ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa_softcap ```python import numpy as np import onnx q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_gqa_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes_softcap ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, softcap=3.0, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_diff_heads_sizes_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_with_past_and_present ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_3d_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_gqa_with_past_and_present ```python import numpy as np import onnx q_num_heads, kv_num_heads = 9, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 72).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_3d_gqa_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_diff_head_sizes_with_past_and_present ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 30).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_3d_diff_heads_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_with_past_and_present_qk_matmul ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_3d_with_past_and_present_qk_matmul", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_with_past_and_present_qk_matmul_bias ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, qk_matmul_output_mode=2, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, qk_matmul_output_mode=2, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_3d_with_past_and_present_qk_matmul_bias", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_with_past_and_present_qk_matmul_softcap ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, softcap=2.0, qk_matmul_output_mode=1, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, softcap=2.0, qk_matmul_output_mode=1, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_3d_with_past_and_present_qk_matmul_softcap", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_with_past_and_present_qk_matmul_softmax ```python import numpy as np import onnx q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"], outputs=["Y", "present_key", "present_value", "qk_matmul_output"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, qk_matmul_output_mode=3, ) past_sequence_length = 12 Q = np.random.rand(2, 4, 24).astype(np.float32) K = np.random.rand(2, 6, 24).astype(np.float32) V = np.random.rand(2, 6, 24).astype(np.float32) attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, past_key=past_key, past_value=past_value, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, qk_matmul_output_mode=3, ) expect( node, inputs=[Q, K, V, attn_mask, past_key, past_value], outputs=[Y, present_key, present_value, qk_matmul_output], name="test_attention_3d_with_past_and_present_qk_matmul_softmax", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_transpose_verification ```python import numpy as np import onnx """Test case to verify correct 3D to 4D transpose behavior. This test verifies that 3D inputs are correctly reshaped and transposed according to the ONNX specification: [batch_size, seq_length, hidden_size] -> [batch_size, seq_length, num_heads, head_size] -> [batch_size, num_heads, seq_length, head_size] """ q_num_heads, kv_num_heads = 3, 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) # Test inputs that will clearly demonstrate the transpose behavior batch_size = 1 q_seq_length = 2 kv_seq_length = 2 head_size = 4 q_hidden_size = q_num_heads * head_size # 3 * 4 = 12 kv_hidden_size = kv_num_heads * head_size # 3 * 4 = 12 # Create structured inputs to verify correct transpose behavior # Q has a pattern where each position in hidden dimension has a specific value Q = np.zeros((batch_size, q_seq_length, q_hidden_size), dtype=np.float32) # Fill Q with pattern: head0=[1,1,1,1], head1=[2,2,2,2], head2=[3,3,3,3] for head in range(q_num_heads): start_idx = head * head_size end_idx = start_idx + head_size Q[0, :, start_idx:end_idx] = float(head + 1) K = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1 V = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1 Y, _, _, _ = _compute_attention( Q, K, V, q_num_heads=q_num_heads, kv_num_heads=kv_num_heads, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_transpose_verification", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_4d_diff_heads_mask4d_padded_kv ```python import numpy as np import onnx node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 10).astype(np.float32) attn_mask = np.random.rand(2, 3, 4, 4).astype(np.float32) nonpad_kv_seqlen = np.array([3, 4], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_diff_heads_mask4d_padded_kv", opset_imports=[onnx.helper.make_opsetid("", 24)], ) # --- bfloat16 / float16 tests for CastLike fixes --- ``` #### _attention_causal_bf16 ```python import numpy as np import onnx """is_causal=1 with bfloat16 inputs exercises causal mask CastLike fix.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(ml_dtypes.bfloat16) K = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) V = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_causal_bf16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_causal_fp16 ```python import numpy as np import onnx """is_causal=1 with float16 inputs.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float16) K = np.random.rand(2, 3, 6, 8).astype(np.float16) V = np.random.rand(2, 3, 6, 8).astype(np.float16) Y, _, _, _ = _compute_attention(Q, K, V, is_causal=1) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_4d_causal_fp16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_padded_kv_bf16 ```python import numpy as np import onnx """nonpad_kv_seqlen with bfloat16 inputs exercises padding mask CastLike fix.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], ) Q = np.random.rand(2, 3, 4, 8).astype(ml_dtypes.bfloat16) K = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) V = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) attn_mask = np.random.rand(2, 3, 4, 4).astype(ml_dtypes.bfloat16) nonpad_kv_seqlen = np.array([3, 4], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_padded_kv_bf16", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_causal_padded_kv_bf16 ```python import numpy as np import onnx """is_causal=1 + nonpad_kv_seqlen with bfloat16 exercises both CastLike fixes.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(ml_dtypes.bfloat16) K = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) V = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) attn_mask = np.random.rand(2, 3, 4, 4).astype(ml_dtypes.bfloat16) nonpad_kv_seqlen = np.array([3, 4], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, nonpad_kv_seqlen=nonpad_kv_seqlen, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_causal_padded_kv_bf16", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_attn_mask_causal_bf16 ```python import numpy as np import onnx """Float attn_mask + is_causal with bfloat16 — mask is bf16, causal mask must be cast.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(2, 3, 4, 8).astype(ml_dtypes.bfloat16) K = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) V = np.random.rand(2, 3, 6, 8).astype(ml_dtypes.bfloat16) attn_mask = np.random.rand(2, 1, 4, 6).astype(ml_dtypes.bfloat16) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_attn_mask_causal_bf16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_3d_causal_bf16 ```python import numpy as np import onnx """3D input with is_causal=1 and bfloat16.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, q_num_heads=3, kv_num_heads=3, ) Q = np.random.rand(2, 4, 24).astype(ml_dtypes.bfloat16) K = np.random.rand(2, 6, 24).astype(ml_dtypes.bfloat16) V = np.random.rand(2, 6, 24).astype(ml_dtypes.bfloat16) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, q_num_heads=3, kv_num_heads=3 ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_causal_bf16", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_softcap_with_neginf_mask ```python import numpy as np import onnx """Softcap + -inf mask: verifies softcap is applied BEFORE mask/bias. If ordering were wrong (mask then softcap), tanh(-inf/softcap) = -1, so softcap * tanh(-inf/softcap) = -softcap (finite). That leaks probability to masked positions. With correct ordering (softcap then mask), the -inf mask values survive to softmax and yield zero weight. """ np.random.seed(42) B, H, S_q, S_kv, D = 1, 1, 4, 6, 8 Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, S_kv, D).astype(np.float32) V = np.random.rand(B, H, S_kv, D).astype(np.float32) # All Q positions are blocked from KV positions 4 and 5. attn_mask = np.zeros((S_q, S_kv), dtype=np.float32) attn_mask[:, 4:] = -np.inf softcap = 0.5 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], softcap=softcap, ) Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask, softcap=softcap) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_softcap_neginf_mask", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_softcap_with_neginf_mask_poison ```python import numpy as np import onnx """Softcap + -inf mask + poison values at masked KV positions. V has value 1000 at the masked positions (4 and 5). With correct ordering the output stays in [0, 1] because the mask zeros out those positions. With wrong ordering the output explodes (> 50), making the failure obvious even with loose tolerances. """ np.random.seed(42) B, H, S_q, S_kv, D = 1, 1, 4, 6, 8 Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, S_kv, D).astype(np.float32) V = np.random.rand(B, H, S_kv, D).astype(np.float32) # Block all Q positions from KV positions 4 and 5. attn_mask = np.zeros((S_q, S_kv), dtype=np.float32) attn_mask[:, 4:] = -np.inf # Poison: if attention leaks to masked positions, output >> 1. V[:, :, 4:, :] = 1000.0 softcap = 0.5 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], softcap=softcap, ) Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask, softcap=softcap) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_4d_softcap_neginf_mask_poison", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_4d_gqa_causal_nonpad_decode ```python import numpy as np import onnx """External/static-cache decode (S_q=1) with per-batch valid lengths. K/V are the full static cache buffer; ``nonpad_kv_seqlen`` marks how many leading keys are valid per batch. With bottom-right (offset-aware) causal masking the single decode query attends keys ``0..nonpad[b]-1``. Under the old top-left alignment it would attend only key 0, so this test fails pre-fix and passes post-fix. """ np.random.seed(0) B, H_q, H_kv, L, D = 2, 4, 2, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H_q, 1, D).astype(np.float32) K = np.random.rand(B, H_kv, L, D).astype(np.float32) V = np.random.rand(B, H_kv, L, D).astype(np.float32) # Batch 0 has all 8 keys valid, batch 1 only the first 5. nonpad_kv_seqlen = np.array([8, 5], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) expect( node, inputs=[Q, K, V, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_gqa_causal_nonpad_decode", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_4d_gqa_causal_nonpad_decode_fp16 ```python import numpy as np import onnx """fp16 variant of the external-cache decode case (locks -inf dtype handling).""" np.random.seed(0) B, H_q, H_kv, L, D = 2, 4, 2, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H_q, 1, D).astype(np.float16) K = np.random.rand(B, H_kv, L, D).astype(np.float16) V = np.random.rand(B, H_kv, L, D).astype(np.float16) nonpad_kv_seqlen = np.array([8, 5], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) expect( node, inputs=[Q, K, V, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_gqa_causal_nonpad_decode_fp16", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_4d_causal_nonpad_continued_prefill ```python import numpy as np import onnx """Continued / chunked prefill (S_q=2) into a partially-filled static cache. With ``nonpad_kv_seqlen = [4]`` and ``S_q = 2`` the bottom-right offset is ``4 - 2 = 2``: query 0 attends keys ``{0,1,2}`` and query 1 attends ``{0,1,2,3}``. The old top-left alignment would mask everything past the diagonal (``{0}`` and ``{0,1}``), so this test fails pre-fix. """ np.random.seed(1) B, H, L, D = 1, 2, 4, 8 S_q = 2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, L, D).astype(np.float32) V = np.random.rand(B, H, L, D).astype(np.float32) nonpad_kv_seqlen = np.array([4], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) expect( node, inputs=[Q, K, V, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_causal_nonpad_continued_prefill", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_4d_causal_with_past_and_present ```python import numpy as np import onnx """Regression guard: internal (past_key) cache + is_causal. This exercises the unchanged scalar bottom-right path (offset = past_sequence_length). Its golden output must remain identical to the pre-fix behavior, proving the external-cache change does not touch the past_key path. """ np.random.seed(2) node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], is_causal=1, ) past_sequence_length = 3 Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 4, 8).astype(np.float32) V = np.random.rand(2, 3, 4, 8).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, past_key=past_key, past_value=past_value, is_causal=1, ) expect( node, inputs=[Q, K, V, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_4d_causal_with_past_and_present", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_causal_boolmask_nan_robustness ```python import numpy as np import onnx """Composed ``is_causal`` + boolean ``attn_mask`` NaN-robustness. The causal frontier (lower-triangular here, offset 0) and the boolean ``attn_mask`` are intersected: a key is attended only if allowed by both. This exercises two pre-fix NaN sources on the same forward pass: * **Bug-1 (allowed cells stay finite).** Query 0 is allowed key 0 by both the causal frontier (``{0}``) and the mask (``True`` at key 0). The old ``(1 - attn_mask) * -inf`` conversion computes ``0 * -inf = NaN`` at that allowed cell, poisoning the row. The select conversion ``where(attn_mask, 0, -inf)`` keeps it finite. * **Bug-2 (fully-masked row -> 0).** Query 1 is allowed keys ``{0, 1}`` by the causal frontier but the mask is ``False`` at both, so the combined constraint allows no key. ``softmax`` of an all-``-inf`` row is ``NaN``; the fully-masked-row guard zeros it before the ``P @ V`` contraction so the output row is ``0``. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(3) B, H, S, D = 1, 2, 2, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H, S, D).astype(np.float32) K = np.random.rand(B, H, S, D).astype(np.float32) V = np.random.rand(B, H, S, D).astype(np.float32) # Row 0: key 0 allowed (Bug-1 allowed cell). Row 1: no key allowed -> fully # masked once intersected with the causal frontier (Bug-2 empty row). attn_mask = np.array([[True, False], [False, False]], dtype=np.bool_) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, ) # Bug-1: allowed cells are finite (no NaN anywhere). Bug-2: the fully-masked # query row is exactly zero, not NaN. assert np.all(np.isfinite(Y)), "allowed cells must be finite (Bug-1)" assert np.array_equal(Y[:, :, 1, :], np.zeros_like(Y[:, :, 1, :])), ( "fully-masked row must be zero (Bug-2)" ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_causal_boolmask_nan_robustness", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_23_boolmask_fullymasked_row_nan_robustness ```python import numpy as np import onnx """Opset-23 fully-masked boolean ``attn_mask`` row -> zero (not ``NaN``). This locks the opset-23 / ``old.cc`` function-body fully-masked-row guard against future regressions. In opset 23 the only in-contract fully-masked row comes from an all-``False`` boolean ``attn_mask`` row (``is_causal`` is not set here): every key for that query is disallowed, so ``softmax`` over an all-``-inf`` bias row is ``NaN``. The guard zeros that row's probabilities before the ``P @ V`` contraction so the output row is exactly ``0``, while rows with at least one allowed key are unchanged. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(4) B, H, S, D = 1, 2, 2, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], ) Q = np.random.rand(B, H, S, D).astype(np.float32) K = np.random.rand(B, H, S, D).astype(np.float32) V = np.random.rand(B, H, S, D).astype(np.float32) # Row 0: no key allowed -> fully masked (Bug-2 empty row). Row 1: both keys # allowed -> finite, unchanged by the guard. attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_) Y, _, _, _ = _compute_attention(Q, K, V, attn_mask=attn_mask) # Fully-masked row 0 is exactly zero (not NaN); every other cell is finite. assert np.all(np.isfinite(Y)), "non-masked rows must be finite" assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), ( "fully-masked row must be zero (Bug-2)" ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_23_boolmask_fullymasked_row_nan_robustness", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_4d_causal_nonpad_negative_offset_structural_empty ```python import numpy as np import onnx """Negative bottom-right offset: structurally-empty early query rows -> zero. This is the onnx-node twin of the ORT gtest ``Attention_Causal_NonPadKVSeqLen_StructuralEmptyRow_Zero`` / ``StructuralEmptyRows_Zero_CUDA``. With ``nonpad_kv_seqlen = [2]`` and ``S_q = 4`` the bottom-right offset is ``2 - 4 = -2``: query row ``sq`` attends keys ``0..(sq - 2)``, so rows 0 and 1 have an empty key set. Their ``softmax`` over an all-``-inf`` bias row is ``NaN``; the fully-masked-row guard zeros those rows before the ``P @ V`` contraction so the output rows are exactly ``0``, while rows 2 and 3 (attending keys ``{0}`` and ``{0,1}``) stay finite and nonzero. A ``nonpad_kv_seqlen[b] < q_sequence_length`` input is out of the contract's intended use, but its result is still well-defined (zeroed rows) rather than ``NaN``; this test pins that defined behavior. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(7) B, H, L, D = 1, 2, 4, 8 S_q = 4 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, L, D).astype(np.float32) V = np.random.rand(B, H, L, D).astype(np.float32) # offset = nonpad - S_q = 2 - 4 = -2 -> rows 0,1 structurally empty. nonpad_kv_seqlen = np.array([2], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) # Structurally-empty early rows are exactly zero (not NaN); later rows finite. assert np.all(np.isfinite(Y)), "all output rows must be finite" assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), ( "structurally-empty row 0 must be zero" ) assert np.array_equal(Y[:, :, 1, :], np.zeros_like(Y[:, :, 1, :])), ( "structurally-empty row 1 must be zero" ) assert np.any(Y[:, :, 2, :] != 0) and np.any(Y[:, :, 3, :] != 0), ( "rows with a non-empty key set must be nonzero" ) expect( node, inputs=[Q, K, V, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_causal_nonpad_negative_offset_structural_empty", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_23_fullymasked_qk_matmul_output_mode3_zero ```python import numpy as np import onnx """Opset-23 ``qk_matmul_output_mode=3`` fully-masked row is a zero row. Mode ``3`` exposes the post-softmax matrix as the optional ``qk_matmul_output``. For a fully-masked query row (all-``False`` boolean ``attn_mask`` row), the fully-masked-row guard is applied before this output is produced, so the mode-3 row is zeroed, consistent with the primary output ``Y`` row (both are ``0``). This pins the mandated agreement between the guarded primary output and the mode-3 output at opset 23. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(13) B, H, S, D = 1, 2, 2, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], qk_matmul_output_mode=3, ) Q = np.random.rand(B, H, S, D).astype(np.float32) K = np.random.rand(B, H, S, D).astype(np.float32) V = np.random.rand(B, H, S, D).astype(np.float32) # Row 0: no key allowed -> fully masked. Row 1: both keys allowed -> finite. attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, qk_matmul_output_mode=3, ) # Primary output row 0 and the mode-3 row 0 are both guarded to zero. assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), ( "fully-masked primary output row must be zero" ) assert np.array_equal( qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :]) ), "mode-3 output row for a fully-masked query must be zero (consistent with Y)" assert np.all(np.isfinite(qk_matmul_output)), ( "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)" ) assert np.all(np.isfinite(Y)), ( "all Y rows are finite (the fully-masked row is guarded to 0.0)" ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_23_fullymasked_qk_matmul_output_mode3_zero", opset_imports=[onnx.helper.make_opsetid("", 23)], ) ``` #### _attention_24_fullymasked_qk_matmul_output_mode3_zero ```python import numpy as np import onnx """Opset-24 ``qk_matmul_output_mode=3`` fully-masked row is a zero row. The opset-24 twin of ``export_attention_23_fullymasked_qk_matmul_output_mode3_zero``. Mode ``3`` exposes the post-softmax matrix as the optional ``qk_matmul_output``. For a fully-masked query row (all-``False`` boolean ``attn_mask`` row), the fully-masked-row guard is applied before this output is produced, so the mode-3 row is zeroed, consistent with the primary output ``Y`` row (both are ``0``). This pins the mandated agreement between the guarded primary output and the mode-3 output at opset 24. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(13) B, H, S, D = 1, 2, 2, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], qk_matmul_output_mode=3, ) Q = np.random.rand(B, H, S, D).astype(np.float32) K = np.random.rand(B, H, S, D).astype(np.float32) V = np.random.rand(B, H, S, D).astype(np.float32) # Row 0: no key allowed -> fully masked. Row 1: both keys allowed -> finite. attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, qk_matmul_output_mode=3, ) # Primary output row 0 and the mode-3 row 0 are both guarded to zero. assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])), ( "fully-masked primary output row must be zero" ) assert np.array_equal( qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :]) ), "mode-3 output row for a fully-masked query must be zero (consistent with Y)" assert np.all(np.isfinite(qk_matmul_output)), ( "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)" ) assert np.all(np.isfinite(Y)), ( "all Y rows are finite (the fully-masked row is guarded to 0.0)" ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_24_fullymasked_qk_matmul_output_mode3_zero", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_24_qk_matmul_output_mode3_softmax_precision ```python import numpy as np import onnx """Mode-3 ``qk_matmul_output`` is emitted at the output precision ``T1``. ``qk_matmul_output_mode=3`` exposes the post-softmax probabilities. When ``softmax_precision`` differs from the operator's output type ``T1`` (here ``T1 = float16`` with softmax computed in ``float32``), the mode-3 output is cast back to ``T1`` -- matching the reference implementation, which casts the exposed matrix to ``Q.dtype``. This locks both the dtype contract and the fully-masked-row zeroing under a non-default ``softmax_precision``. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(17) B, H, S, D = 1, 2, 2, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], qk_matmul_output_mode=3, softmax_precision=int(onnx.TensorProto.FLOAT), ) # T1 = float16; softmax runs in float32, so the mode-3 output is cast back to # float16 on emission. Q = np.random.rand(B, H, S, D).astype(np.float16) K = np.random.rand(B, H, S, D).astype(np.float16) V = np.random.rand(B, H, S, D).astype(np.float16) # Row 0: fully masked. Row 1: both keys allowed -> finite. attn_mask = np.array([[False, False], [True, True]], dtype=np.bool_) Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, qk_matmul_output_mode=3, softmax_precision=int(onnx.TensorProto.FLOAT), ) # The mode-3 output is emitted at T1 (float16), not the float32 softmax # precision, matching the operator's output type. assert qk_matmul_output.dtype == np.float16, ( "mode-3 qk_matmul_output must be emitted at the output precision T1 (float16)" ) # The fully-masked row is still guarded to zero, consistent with Y. assert np.array_equal( qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :]) ), "mode-3 output row for a fully-masked query must be zero (consistent with Y)" assert np.all(np.isfinite(qk_matmul_output)), ( "all mode-3 rows are finite (the fully-masked row is guarded to 0.0)" ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_24_qk_matmul_output_mode3_softmax_precision", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_4d_causal_nonpad_attn_mask_composition ```python import numpy as np import onnx """Compose ``is_causal`` + ``nonpad_kv_seqlen`` + boolean ``attn_mask``. The existing nonpad tests use no ``attn_mask`` and the existing mask tests use no ``nonpad_kv_seqlen``; this is the first to activate all three constraints together on the external-cache path with ``batch > 1``. The three biases are summed additively and a key is attended only if allowed by all three. Crucially the inputs are designed so that **each constraint is independently necessary** -- removing any one changes the golden -- to avoid a degenerate test that a backend ignoring ``is_causal`` and/or ``nonpad_kv_seqlen`` could still pass: * **``is_causal`` binds.** Each batch has a key that the boolean mask allows (``True``) but the bottom-right causal frontier disallows (``j > i + offset``); only ``is_causal`` masks it (batch 0 row 0 key 2, batch 1 row 0 key 3). * **``attn_mask`` binds.** Each batch has a key the causal frontier and the padding bound both allow but the boolean mask sets ``False`` (batch 0 row 2 key 1, batch 1 row 2 key 2); only the mask masks it. * **``nonpad_kv_seqlen`` binds.** ``nonpad_kv_seqlen`` sets the per-batch causal *offset* (``offset = nonpad_kv_seqlen - q_sequence_length``), so dropping it collapses the frontier to top-left (``offset = 0``) and shifts which keys are attended. (Under ``is_causal=1`` the causal frontier already subsumes the ``j < nonpad`` padding bound, so ``nonpad_kv_seqlen`` binds through the offset it induces rather than through a redundant padding cut.) The mask is chosen to leave at least one allowed key on every query row, so this exercises the *intersection* of the three constraints with finite outputs (the fully-masked-row guard is covered by ``test_attention_4d_causal_nonpad_negative_offset_structural_empty`` and ``test_attention_24_fullymasked_qk_matmul_output_mode3_zero``). 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(11) B, H, L, D = 2, 2, 6, 8 S_q = 3 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, L, D).astype(np.float32) V = np.random.rand(B, H, L, D).astype(np.float32) nonpad_kv_seqlen = np.array([4, 5], dtype=np.int64) # offsets [1, 2] # Per-batch (B, 1, S_q, L) bool mask. Each batch is laid out so all three # constraints uniquely bind (see the docstring): a causal-only-masked key # (mask True, j > i + offset), a mask-only-masked key (mask False, causal + # nonpad allow it), and >=1 allowed key per row. attn_mask = np.array( [ [ [ [True, True, True, False, False, False], [True, True, True, False, False, False], [True, False, True, True, False, False], ] ], [ [ [True, True, True, True, False, False], [True, True, True, True, False, False], [True, True, False, True, True, False], ] ], ], dtype=np.bool_, ) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) # The chosen mask leaves >=1 allowed key per row, so the composition stays # finite (no fully-masked row in this case). assert np.all(np.isfinite(Y)), "composed-constraint output must be finite" # Non-degeneracy: each constraint is independently necessary. Removing any one # of the three (is_causal, attn_mask, nonpad_kv_seqlen) must change the result, # so a backend that ignores is_causal or nonpad_kv_seqlen cannot reproduce the # golden by applying only the most restrictive mask. y_no_causal, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=0 ) y_no_mask, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1 ) y_no_nonpad, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1 ) assert not np.allclose(Y, y_no_causal, equal_nan=True), ( "is_causal must bind: dropping it changes the result" ) assert not np.allclose(Y, y_no_mask, equal_nan=True), ( "attn_mask must bind: dropping it changes the result" ) assert not np.allclose(Y, y_no_nonpad, equal_nan=True), ( "nonpad_kv_seqlen must bind (via the causal offset): dropping it changes the result" ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_causal_nonpad_attn_mask_composition", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_4d_causal_nonpad_batch_prefill ```python import numpy as np import onnx """Batch>1 continued prefill with distinct per-batch bottom-right offsets. The batched generalization of the ``batch == 1`` continued-prefill case: with ``nonpad_kv_seqlen = [4, 5, 6]`` and ``S_q = 2`` the per-batch bottom-right offsets are ``[2, 3, 4]`` (all ``>= 0``), so each batch realigns its causal frontier to its own valid-key prefix. This pins that the per-batch offset is applied independently across the batch dimension. 4D Q/K/V is used so ``q_num_heads``/``kv_num_heads`` are omitted (passing them would make the function body treat the input as 3D). """ np.random.seed(12) B, H, L, D = 3, 2, 6, 8 S_q = 2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, L, D).astype(np.float32) V = np.random.rand(B, H, L, D).astype(np.float32) nonpad_kv_seqlen = np.array([4, 5, 6], dtype=np.int64) # offsets [2, 3, 4] Y, _, _, _ = _compute_attention( Q, K, V, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, ) assert np.all(np.isfinite(Y)), "per-batch prefill output must be finite" expect( node, inputs=[Q, K, V, nonpad_kv_seqlen], outputs=[Y], name="test_attention_4d_causal_nonpad_batch_prefill", opset_imports=[onnx.helper.make_opsetid("", 24)], ) ``` #### _attention_local_window ```python import numpy as np import onnx """Causal sliding window attention with two preceding positions.""" left_window_size = 2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, is_causal=1, left_window_size=left_window_size ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_local_window", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_bidirectional_window ```python import numpy as np import onnx """Asymmetric bidirectional window independent of causal masking.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], left_window_size=1, right_window_size=2, ) Q = np.zeros((1, 1, 5, 1), dtype=np.float32) K = np.zeros((1, 1, 5, 1), dtype=np.float32) V = np.arange(5, dtype=np.float32).reshape(1, 1, 5, 1) Y, _, _, _ = _compute_attention( Q, K, V, left_window_size=1, right_window_size=2, ) np.testing.assert_allclose( Y.reshape(-1), np.array([1.0, 1.5, 2.5, 3.0, 3.5], dtype=np.float32) ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_bidirectional_window", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_default ```python import numpy as np import onnx """Disabled window bounds behave identically to version 24.""" node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], left_window_size=-1, right_window_size=-1, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, left_window_size=-1, right_window_size=-1 ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_local_window_default", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_rank1_boolean_mask ```python import numpy as np import onnx """A rank-1 boolean mask retains standard right-aligned broadcasting.""" left_window_size = 2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 6, 8).astype(np.float32) V = np.random.rand(2, 3, 6, 8).astype(np.float32) attn_mask = np.array([True, True, True, True, False, False]) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y], name="test_attention_local_window_rank1_boolean_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_with_past ```python import numpy as np import onnx """Sliding window with internal KV cache (past_key/past_value).""" left_window_size = 2 past_sequence_length = 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "", "past_key", "past_value"], outputs=["Y", "present_key", "present_value"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(2, 3, 4, 8).astype(np.float32) K = np.random.rand(2, 3, 2, 8).astype(np.float32) V = np.random.rand(2, 3, 2, 8).astype(np.float32) past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32) Y, present_key, present_value, _ = _compute_attention( Q, K, V, past_key=past_key, past_value=past_value, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, past_key, past_value], outputs=[Y, present_key, present_value], name="test_attention_local_window_with_past", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_ext_cache_rank3_head_mask ```python import numpy as np import onnx """External cache with a legal rank-3 ``(heads, q, kv)`` mask.""" left_window_size = 2 B, H, S_q, S_kv, D = 2, 3, 4, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, S_kv, D).astype(np.float32) V = np.random.rand(B, H, S_kv, D).astype(np.float32) # Rank 3 is right-aligned as (heads, q, kv), not (batch, q, kv). attn_mask = np.random.rand(H, S_q, S_kv).astype(np.float32) # External cache: nonpad_kv_seqlen marks valid key count per batch nonpad_kv_seqlen = np.array([6, 7], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_local_window_ext_cache_rank3_head_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_ext_cache_rank4_batch_mask ```python import numpy as np import onnx """External cache with a batch-specific rank-4 ``(batch, 1, q, kv)`` mask.""" left_window_size = 2 B, H, S_q, S_kv, D = 2, 3, 4, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, S_kv, D).astype(np.float32) V = np.random.rand(B, H, S_kv, D).astype(np.float32) attn_mask = np.random.rand(B, 1, S_q, S_kv).astype(np.float32) nonpad_kv_seqlen = np.array([6, 7], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_local_window_ext_cache_rank4_batch_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_ext_cache_rank2_mask ```python import numpy as np import onnx """External cache with a conventional rank-2 ``(1, kv)`` mask.""" left_window_size = 2 B, H, S_q, S_kv, D = 2, 3, 4, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(B, H, S_q, D).astype(np.float32) K = np.random.rand(B, H, S_kv, D).astype(np.float32) V = np.random.rand(B, H, S_kv, D).astype(np.float32) attn_mask = np.random.rand(1, S_kv).astype(np.float32) nonpad_kv_seqlen = np.array([6, 7], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_local_window_ext_cache_rank2_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_ext_cache_float16_mask ```python import numpy as np import onnx """External cache with a float16 attention mask.""" left_window_size = 2 B, H, S_q, S_kv, D = 2, 3, 4, 8, 8 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"], outputs=["Y"], is_causal=1, left_window_size=left_window_size, ) Q = np.zeros((B, H, S_q, D), dtype=np.float16) K = np.zeros((B, H, S_kv, D), dtype=np.float16) V = np.ones((B, H, S_kv, D), dtype=np.float16) attn_mask = np.zeros((1, S_kv), dtype=np.float16) nonpad_kv_seqlen = np.array([6, 7], dtype=np.int64) Y, _, _, _ = _compute_attention( Q, K, V, attn_mask=attn_mask, nonpad_kv_seqlen=nonpad_kv_seqlen, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen], outputs=[Y], name="test_attention_local_window_ext_cache_float16_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_3d_local_window ```python import numpy as np import onnx """Sliding window with 3D MQA inputs and a distinct V head size.""" left_window_size = 2 node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V"], outputs=["Y"], q_num_heads=4, kv_num_heads=1, is_causal=1, left_window_size=left_window_size, ) Q = np.random.rand(2, 4, 32).astype(np.float32) K = np.random.rand(2, 6, 8).astype(np.float32) V = np.random.rand(2, 6, 6).astype(np.float32) Y, _, _, _ = _compute_attention( Q, K, V, q_num_heads=4, kv_num_heads=1, is_causal=1, left_window_size=left_window_size, ) expect( node, inputs=[Q, K, V], outputs=[Y], name="test_attention_3d_local_window", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` #### _attention_local_window_gqa_rank4_mask ```python import numpy as np import onnx """Windowed GQA with distinct V head size and a per-head boolean mask.""" np.random.seed(25) B, H_q, H_kv, S_q, S_kv, D_qk, D_v = 2, 4, 2, 4, 6, 8, 6 left_window_size = 2 softmax_precision = int(onnx.TensorProto.DOUBLE) node = onnx.helper.make_node( "Attention", inputs=["Q", "K", "V", "attn_mask"], outputs=["Y", "", "", "qk_matmul_output"], is_causal=1, left_window_size=left_window_size, softcap=2.0, softmax_precision=softmax_precision, qk_matmul_output_mode=3, ) Q = np.random.rand(B, H_q, S_q, D_qk).astype(np.float32) K = np.random.rand(B, H_kv, S_kv, D_qk).astype(np.float32) V = np.random.rand(B, H_kv, S_kv, D_v).astype(np.float32) attn_mask = np.ones((B, H_q, S_q, S_kv), dtype=np.bool_) attn_mask[:, :, 0, :] = False Y, _, _, qk_matmul_output = _compute_attention( Q, K, V, attn_mask=attn_mask, is_causal=1, left_window_size=left_window_size, softcap=2.0, softmax_precision=softmax_precision, qk_matmul_output_mode=3, ) assert Y.shape == (B, H_q, S_q, D_v) assert np.array_equal(Y[:, :, 0, :], np.zeros_like(Y[:, :, 0, :])) assert np.array_equal( qk_matmul_output[:, :, 0, :], np.zeros_like(qk_matmul_output[:, :, 0, :]) ) expect( node, inputs=[Q, K, V, attn_mask], outputs=[Y, qk_matmul_output], name="test_attention_local_window_gqa_rank4_mask", opset_imports=[onnx.helper.make_opsetid("", 25)], ) ``` ```{toctree} text_diff_Attention_24_25 ``` (l-onnx-op-attention-24)= ## Attention - 24 ### Version - **name**: [Attention (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Attention) - **domain**: `main` - **since_version**: `24` - **function**: `True` - **support_level**: `SupportType.COMMON` - **shape inference**: `True` This version of the operator has been available **since version 24**. ### Summary Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed. This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V. For self attention, `kv_sequence_length` equals to `q_sequence_length`. For cross attention, query and key might have different lengths. This operator also covers the 3 following variants based on the number of heads: 1) Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, `q_num_heads = kv_num_heads`. 2) Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, `q_num_heads > kv_num_heads`, `q_num_heads % kv_num_heads == 0`. 3) Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`. Attention bias to be added is calculated based on `attn_mask` input and `is_causal` attribute: 1) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score. 2) If `is_causal` is set to `1`, causal masking is applied with bottom-right (offset-aware) alignment: query `i` attends key `j` iff `j <= i + offset`, as illustrated below. ``` 2D causal mask for Attention (PR onnx/onnx#8068) S_q=4 queries, S_k=8 keys Rule: query i attends key j iff j <= i + offset offset = nonpad_kv_seqlen - S_q nonpad_kv_seqlen=4, offset=4-4=0 k0 k1 k2 k3 k4 k5 k6 k7 +----+----+----+----+----+----+----+----+ q0 | ## | | | | | | | | +----+----+----+----+----+----+----+----+ q1 | ## | ## | | | | | | | +----+----+----+----+----+----+----+----+ q2 | ## | ## | ## | | | | | | +----+----+----+----+----+----+----+----+ q3 | ## | ## | ## | ## | | | | | +----+----+----+----+----+----+----+----+ nonpad_kv_seqlen=8, offset=8-4=4 k0 k1 k2 k3 k4 k5 k6 k7 +----+----+----+----+----+----+----+----+ q0 | ## | ## | ## | ## | ## | | | | +----+----+----+----+----+----+----+----+ q1 | ## | ## | ## | ## | ## | ## | | | +----+----+----+----+----+----+----+----+ q2 | ## | ## | ## | ## | ## | ## | ## | | +----+----+----+----+----+----+----+----+ q3 | ## | ## | ## | ## | ## | ## | ## | ## | +----+----+----+----+----+----+----+----+ ``` With `nonpad_kv_seqlen=4` (offset=0), the mask is the standard lower-triangular. With `nonpad_kv_seqlen=8` (offset=4), the diagonal shifts right by 4, so each query sees the 4 additional valid cached keys. `offset` is the count of valid keys preceding the current query block: `offset = past_sequence_length` when `past_key` is provided; `offset = nonpad_kv_seqlen - q_sequence_length` (per batch) when an external cache is indicated by `nonpad_kv_seqlen` without `past_key`; `offset = 0` when neither is provided (the no-cache case, which reduces to the standard lower-triangular mask). When `offset < 0` (`nonpad_kv_seqlen < q_sequence_length`, i.e. more query tokens than cached keys) the leading query rows have an empty key set (no key satisfies `j <= i + offset`) and are fully masked. The causal frontier is computed independently of `attn_mask` and is then composed with it additively: a boolean `attn_mask` intersects the allowed set (its disallowed positions contribute `-inf` to the bias), while a float `attn_mask` is added to the attention scores rather than disabling positions. A fully-masked query row (no key attended, including the negative-offset leading rows) produces a zero output row, not `NaN`, for both `Y` and the mode-`3` `qk_matmul_output` debug output; the mode-`3` `qk_matmul_output` is emitted at the operator's output precision (`T1`). Errata (in-place behavioral correction, no opset bump): the reference implementation and backend tests were incorrect when `nonpad_kv_seqlen != q_sequence_length` (nonzero bottom-right offset, top-left instead of bottom-right causal alignment) and produced `NaN` for fully-masked rows; corrected in version 1.23. This fixed three behaviors described above: external-cache bottom-right causal alignment (`offset = nonpad_kv_seqlen - q_sequence_length`), zero (non-`NaN`) output for fully-masked rows including the mode-`3` `qk_matmul_output`, and the mode-`3` `qk_matmul_output` precision (`T1`). With respect to KV cache update, this operator allows the following two use cases: 1) Cache update happens inside the Attention operator. In this case, the `K` and `V` inputs contain only the incoming tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present key and value, respectively. Note that this only works correctly for the special case where the past key and value do not contain padded tokens. 2) Cache update happens outside the Attention operator (for example, through the `TensorScatter` operator). In this case, the `K` and `V` inputs correspond to the entire cache tensor, so the four optional inputs/outputs past and present key and value should not be used. An additional input `nonpad_kv_seqlen` of shape (batch_size,) may be provided to indicate the number of non-padding tokens in each sample of the batch to save unnecessary computation. Here, the kv_sequence dimension of `attn_mask` can be shorter than `K` and `V`, but still needs to be at least as long as the maximum value of `nonpad_kv_seqlen`. Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided: ``` 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 ``` ### Attributes * **is_causal - INT** (default is `0`): If set to `1`, causal masking is applied. For a square Q/K (no cache offset) this is a lower-triangular matrix. In general the mask is bottom-right (offset-aware): query in-block index `i` attends key `j` iff `j <= i + offset`, where `offset` is the count of valid keys preceding the query block (`past_sequence_length` for an internal `past_key` cache, or `nonpad_kv_seqlen - q_sequence_length` per batch for an external cache). When `offset = 0` this reduces to the lower-triangular (top-left) mask. * **kv_num_heads - INT** : Number of heads of key and value. Must be used with 3D inputs of Q, K and V. * **q_num_heads - INT** : Number of heads of query. Must be used with 3D inputs of Q, K and V. * **qk_matmul_output_mode - INT** (default is `0`): If set to `0`, qk_matmul_output is the output of qk matmul. If set to `1`, qk_matmul_output is the output after the softcap operation (before mask addition). If set to `2`, qk_matmul_output includes the attention mask and softcap (if provided) applied to the output of qk matmul. If set to `3`, qk_matmul_output is the output after the softmax operation. In mode `3`, a fully-masked query row (every key disallowed) is a zero row, consistent with the corresponding row of the primary output `Y`: the fully-masked-row guard is applied before this output is produced. The mode-`3` output is emitted at the operator's output precision (`T1`); when `softmax_precision` differs from `T1` this is a cast of the softmax result to `T1`. Default value is 0. * **scale - FLOAT** : Scaling factor applied to $Q*K^T$. Default value is `1/sqrt(head_size)`. To prevent [numerical overflow](https://tinyurl.com/sudb9s96), scale `Q`, `K` by `sqrt(scale)` before matmul. * **softcap - FLOAT** (default is `0.0`): Softcap value for attention weights. Default value is 0. * **softmax_precision - INT** : The floating-point precision used in softmax computation. If softmax precision is not provided, the same precision as the input of softmax (Q and K) is used. ### Inputs Between 3 and 7 inputs. - **Q** (heterogeneous) - **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** (heterogeneous) - **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** (heterogeneous) - **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** (optional, heterogeneous) - **U**: Attention mask. Shape must be broadcastable to `(batch_size, q_num_heads, q_sequence_length, total_sequence_length)` where `total_sequence_length = past_sequence_length + kv_sequence_length.` The last dimension can also be shorter than `total_sequence_length` and will be padded to `total_sequence_length` with negative infinity. Two types of masks are supported: a boolean mask where a value of `True` indicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score. - **past_key** (optional, heterogeneous) - **T1**: past state cache for key with shape `(batch_size, kv_num_heads, past_sequence_length, head_size)` - **past_value** (optional, heterogeneous) - **T2**: past state cache for value with shape `(batch_size, kv_num_heads, past_sequence_length, v_head_size)` - **nonpad_kv_seqlen** (optional, heterogeneous) - **tensor(int64)**: A vector of integers of shape `(batch_size,)` that indicates the number of valid (ie, non-padding) tokens in each sample. A padding mask can be derived from this. This should not be used together with `past_key` and `past_value` inputs or `present_key` and `present_value` outputs (See the KV cache use cases in the operator description). ### Outputs Between 1 and 4 outputs. - **Y** (heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain Q and K inputs types to float tensors. * **T2** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain V input types to float tensors. * **U** in ( `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)` ): Constrain output 'mask' types to boolean tensors and input types. ```{toctree} text_diff_Attention_23_25 text_diff_Attention_23_24 ``` (l-onnx-op-attention-23)= ## Attention - 23 ### Version - **name**: [Attention (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Attention) - **domain**: `main` - **since_version**: `23` - **function**: `True` - **support_level**: `SupportType.COMMON` - **shape inference**: `True` This version of the operator has been available **since version 23**. ### Summary Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed. This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V. For self attention, `kv_sequence_length` equals to `q_sequence_length`. For cross attention, query and key might have different lengths. This operator also covers the 3 following variants based on the number of heads: 1) Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, `q_num_heads = kv_num_heads`. 2) Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, `q_num_heads > kv_num_heads`, `q_num_heads % kv_num_heads == 0`. 3) Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`. Attention bias to be added is calculated based on `attn_mask` input and `is_causal` attribute: 1) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score. 2) If `is_causal` is set to `1`, causal masking is applied with bottom-right (offset-aware) alignment: query `i` attends key `j` iff `j <= i + past_sequence_length` (the count of cached keys in `past_key`); for a square Q/K this is the standard lower-triangular mask. The causal frontier is computed independently of the `attn_mask` input and is then composed with it additively, by summing their attention biases: a boolean `attn_mask` intersects the allowed set (its disallowed positions contribute `-inf` to the bias), while a float `attn_mask` is added to the attention scores rather than disabling positions. A fully-masked query row (every key's combined additive bias is `-inf`, e.g. an all-`False` boolean `attn_mask` row) produces a zero output row (matching prevailing runtime practice), not `NaN`. Errata (in-place behavioral correction, no opset bump): a fully-masked query row (e.g. an all-`False` boolean `attn_mask` row) now produces a zero output row instead of `NaN`, and the same zero-row guard applies to the mode-`3` `qk_matmul_output` debug output; this only replaces previously-`NaN` outputs. The mode-`3` `qk_matmul_output` is also now emitted at the operator's output precision (`T1`), matching the reference implementation, which affects only its dtype and only when `softmax_precision` differs from `T1`. No numerically useful, well-defined result of the released opset is otherwise changed. Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided: ``` The following pattern is applied by this operator: Q K V | | | Q*sqrt(scale) K*sqrt(scale) | | | | | Transpose | | | | ---MatMul--- | | | softcap (if provided) | | | at_mask---Add | | | Softmax | | | -----MatMul------ | Y ``` ### Attributes * **is_causal - INT** (default is `0`): If set to `1`, causal masking is applied with bottom-right (offset-aware) alignment: query `i` attends key `j` iff `j <= i + past_sequence_length` (the count of cached keys in `past_key`); for a square Q/K this is the standard lower-triangular mask. * **kv_num_heads - INT** : Number of heads of key and value. Must be used with 3D inputs of Q, K and V. * **q_num_heads - INT** : Number of heads of query. Must be used with 3D inputs of Q, K and V. * **qk_matmul_output_mode - INT** (default is `0`): If set to `0`, qk_matmul_output is the output of qk matmul. If set to `1`, qk_matmul_output is the output after the softcap operation (before mask addition). If set to `2`, qk_matmul_output includes the attention mask and softcap (if provided) applied to the output of qk matmul. If set to `3`, qk_matmul_output is the output after the softmax operation. In mode `3`, a fully-masked query row (every key disallowed, e.g. an all-`False` boolean `attn_mask` row) is a zero row, consistent with the corresponding row of the primary output `Y`: the fully-masked-row guard is applied before this output is produced. The mode-`3` output is emitted at the operator's output precision (`T1`); when `softmax_precision` differs from `T1` this is a cast of the softmax result to `T1`. Default value is 0. * **scale - FLOAT** : Scaling factor applied to $Q*K^T$. Default value is `1/sqrt(head_size)`. To prevent [numerical overflow](https://tinyurl.com/sudb9s96), scale `Q`, `K` by `sqrt(scale)` before matmul. * **softcap - FLOAT** (default is `0.0`): Softcap value for attention weights. Default value is 0. * **softmax_precision - INT** : The floating-point precision used in softmax computation. If softmax precision is not provided, the same precision as the input of softmax (Q and K) is used. ### Inputs Between 3 and 6 inputs. - **Q** (heterogeneous) - **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** (heterogeneous) - **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** (heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **T1**: past state cache for key with shape `(batch_size, kv_num_heads, past_sequence_length, head_size)` - **past_value** (optional, heterogeneous) - **T2**: past state cache for value with shape `(batch_size, kv_num_heads, past_sequence_length, v_head_size)` ### Outputs Between 1 and 4 outputs. - **Y** (heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** (optional, heterogeneous) - **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** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain Q and K inputs types to float tensors. * **T2** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ): Constrain V input types to float tensors. * **U** in ( `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)` ): Constrain output 'mask' types to boolean tensors and input types.