Attention¶
Attention - 25¶
Version¶
name: Attention (GitHub)
domain:
mainsince_version:
25function:
Truesupport_level:
SupportType.COMMONshape 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:
Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762,
q_num_heads = kv_num_heads.Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245,
q_num_heads > kv_num_heads,q_num_heads % kv_num_heads == 0.Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150,
q_num_heads > kv_num_heads,kv_num_heads=1.
Attention bias to be added is calculated based on attn_mask input and is_causal attribute:
attn_mask: A boolean mask where a value ofTrueindicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.If
is_causalis set to1, causal masking is applied with bottom-right (offset-aware) alignment: queryiattends keyjiffj <= 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:
Cache update happens inside the Attention operator. In this case, the
KandVinputs contain only the incoming tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present key and value, respectively. Note that this only works correctly for the special case where the past key and value do not contain padded tokens.Cache update happens outside the Attention operator (for example, through the
TensorScatteroperator). In this case, theKandVinputs correspond to the entire cache tensor, so the four optional inputs/outputs past and present key and value should not be used. An additional inputnonpad_kv_seqlenof 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 ofattn_maskcan be shorter thanKandV, but still needs to be at least as long as the maximum value ofnonpad_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 indexiattends keyjiffj <= i + offset, whereoffsetis the count of valid keys preceding the query block (past_sequence_lengthfor an internalpast_keycache, ornonpad_kv_seqlen - q_sequence_lengthper batch for an external cache). Whenoffset = 0this 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
0allows the current position but no preceding position, while-1leaves the left side unbounded. This bound is composed withis_causalandattn_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 mode3, a fully-masked query row (every key disallowed) is a zero row, consistent with the corresponding row of the primary outputY. The mode-3output is emitted at the operator’s output precision (T1); whensoftmax_precisiondiffers fromT1this is a cast of the softmax result toT1.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
0allows the current position but no following position, while-1leaves the right side unbounded. Setis_causal=0to 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, scaleQ,Kbysqrt(scale)before matmul.softcap - FLOAT (default is
0.0):Soft cap for attention logits, applied as
softcap * tanh(logits / softcap). Default value of0.0means 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_sizeK (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_sizeV (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_sizeattn_mask (optional, heterogeneous) - U:
Attention mask. Shape must be broadcastable to
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length. The last dimension can also be shorter thantotal_sequence_lengthand will be padded tototal_sequence_lengthwith negative infinity. Two types of masks are supported: a boolean mask where a value ofTrueindicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score.past_key (optional, heterogeneous) - T1:
Past state for key with shape
(batch_size, kv_num_heads, past_sequence_length, head_size). Must be used together withpast_valueinput.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 withpast_keyinput.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 withpast_keyandpast_valueinputs orpresent_keyandpresent_valueoutputs (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_sizepresent_key (optional, heterogeneous) - T1:
Updated key cache with shape
(batch_size, kv_num_heads, total_sequence_length, head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.present_value (optional, heterogeneous) - T2:
Updated value cache with shape
(batch_size, kv_num_heads, total_sequence_length, v_head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.qk_matmul_output (optional, heterogeneous) - T1:
The output of QK matmul. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.
Type Constraints¶
T1 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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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¶
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)],
)
Attention - 24¶
Version¶
name: Attention (GitHub)
domain:
mainsince_version:
24function:
Truesupport_level:
SupportType.COMMONshape 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:
Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762,
q_num_heads = kv_num_heads.Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245,
q_num_heads > kv_num_heads,q_num_heads % kv_num_heads == 0.Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150,
q_num_heads > kv_num_heads,kv_num_heads=1.
Attention bias to be added is calculated based on attn_mask input and is_causal attribute:
attn_mask: A boolean mask where a value ofTrueindicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.If
is_causalis set to1, causal masking is applied with bottom-right (offset-aware) alignment: queryiattends keyjiffj <= 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:
Cache update happens inside the Attention operator. In this case, the
KandVinputs contain only the incoming tokens for the current autoregressive step, and the four optional inputs/outputs past and present key and value are all needed. The Attention op performs a Concat operation on the past and incoming key and value to form the present key and value, respectively. Note that this only works correctly for the special case where the past key and value do not contain padded tokens.Cache update happens outside the Attention operator (for example, through the
TensorScatteroperator). In this case, theKandVinputs correspond to the entire cache tensor, so the four optional inputs/outputs past and present key and value should not be used. An additional inputnonpad_kv_seqlenof 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 ofattn_maskcan be shorter thanKandV, but still needs to be at least as long as the maximum value ofnonpad_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 indexiattends keyjiffj <= i + offset, whereoffsetis the count of valid keys preceding the query block (past_sequence_lengthfor an internalpast_keycache, ornonpad_kv_seqlen - q_sequence_lengthper batch for an external cache). Whenoffset = 0this 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 to1, qk_matmul_output is the output after the softcap operation (before mask addition). If set to2, qk_matmul_output includes the attention mask and softcap (if provided) applied to the output of qk matmul. If set to3, qk_matmul_output is the output after the softmax operation. In mode3, a fully-masked query row (every key disallowed) is a zero row, consistent with the corresponding row of the primary outputY: the fully-masked-row guard is applied before this output is produced. The mode-3output is emitted at the operator’s output precision (T1); whensoftmax_precisiondiffers fromT1this is a cast of the softmax result toT1. Default value is 0.scale - FLOAT :
Scaling factor applied to \(Q*K^T\). Default value is
1/sqrt(head_size). To prevent numerical overflow, scaleQ,Kbysqrt(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_sizeK (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_sizeV (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_sizeattn_mask (optional, heterogeneous) - U:
Attention mask. Shape must be broadcastable to
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.The last dimension can also be shorter thantotal_sequence_lengthand will be padded tototal_sequence_lengthwith negative infinity. Two types of masks are supported: a boolean mask where a value ofTrueindicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score.past_key (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 withpast_keyandpast_valueinputs orpresent_keyandpresent_valueoutputs (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_sizepresent_key (optional, heterogeneous) - T1:
Updated key cache with shape
(batch_size, kv_num_heads, total_sequence_length, head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.present_value (optional, heterogeneous) - T2:
Updated value cache with shape
(batch_size, kv_num_heads, total_sequence_length, v_head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.qk_matmul_output (optional, heterogeneous) - T1:
The output of QK matmul. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.
Type Constraints¶
T1 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.
Attention - 23¶
Version¶
name: Attention (GitHub)
domain:
mainsince_version:
23function:
Truesupport_level:
SupportType.COMMONshape 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:
Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762,
q_num_heads = kv_num_heads.Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245,
q_num_heads > kv_num_heads,q_num_heads % kv_num_heads == 0.Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150,
q_num_heads > kv_num_heads,kv_num_heads=1.
Attention bias to be added is calculated based on attn_mask input and is_causal attribute:
attn_mask: A boolean mask where a value ofTrueindicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.If
is_causalis set to1, causal masking is applied with bottom-right (offset-aware) alignment: queryiattends keyjiffj <= i + past_sequence_length(the count of cached keys inpast_key); for a square Q/K this is the standard lower-triangular mask. The causal frontier is computed independently of theattn_maskinput and is then composed with it additively, by summing their attention biases: a booleanattn_maskintersects the allowed set (its disallowed positions contribute-infto the bias), while a floatattn_maskis 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-Falsebooleanattn_maskrow) produces a zero output row (matching prevailing runtime practice), notNaN.
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: queryiattends keyjiffj <= i + past_sequence_length(the count of cached keys inpast_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 to1, qk_matmul_output is the output after the softcap operation (before mask addition). If set to2, qk_matmul_output includes the attention mask and softcap (if provided) applied to the output of qk matmul. If set to3, qk_matmul_output is the output after the softmax operation. In mode3, a fully-masked query row (every key disallowed, e.g. an all-Falsebooleanattn_maskrow) is a zero row, consistent with the corresponding row of the primary outputY: the fully-masked-row guard is applied before this output is produced. The mode-3output is emitted at the operator’s output precision (T1); whensoftmax_precisiondiffers fromT1this is a cast of the softmax result toT1. Default value is 0.scale - FLOAT :
Scaling factor applied to \(Q*K^T\). Default value is
1/sqrt(head_size). To prevent numerical overflow, scaleQ,Kbysqrt(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_sizeK (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_sizeV (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_sizeattn_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)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.Two types of masks are supported. A boolean mask where a value ofTrueindicates 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_sizepresent_key (optional, heterogeneous) - T1:
Updated key cache with shape
(batch_size, kv_num_heads, total_sequence_length, head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.present_value (optional, heterogeneous) - T2:
Updated value cache with shape
(batch_size, kv_num_heads, total_sequence_length, v_head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.qk_matmul_output (optional, heterogeneous) - T1:
The output of QK matmul. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.
Type Constraints¶
T1 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.