Attention#
Domain:
ai.onnxSince version: 25
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, the K and V inputs correspond to the entire cache tensor, so the four optional inputs/outputs past and
present key and value should not be used. An additional input nonpad_kv_seqlen of shape (batch_size,) may be
provided to indicate the number of non-padding tokens in each sample of the batch to save unnecessary computation.
Here, the kv_sequence dimension of attn_mask can be shorter than K and V, but still needs to be at least as long
as the maximum value of nonpad_kv_seqlen.
Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided:
The following pattern is applied by this operator:
Q K V
| | |
Q*sqrt(scale) K*sqrt(scale) |
| | |
| Transpose |
| | |
---MatMul--- |
| |
softcap (if provided) |
| |
at_mask---Add |
| |
Softmax |
| |
-----MatMul------
|
Y
Inputs
Q (T1): Query tensor. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, head_size)or 3D tensor with shape(batch_size, q_sequence_length, q_hidden_size). For cases with a 3D input tensor,q_hidden_size = q_num_heads * head_sizeK (T1): Key tensor. 4D tensor with shape
(batch_size, kv_num_heads, kv_sequence_length, head_size)or 3D tensor with shape(batch_size, kv_sequence_length, k_hidden_size). For cases with a 3D input tensor,k_hidden_size = kv_num_heads * head_sizeV (T2): Value tensor. 4D tensor with shape
(batch_size, kv_num_heads, kv_sequence_length, v_head_size)or 3D tensor with shape(batch_size, kv_sequence_length, v_hidden_size). For cases with a 3D input tensor,v_hidden_size = kv_num_heads * v_head_sizeattn_mask (U): Attention mask. Shape must be broadcastable to
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length. The last dimension can also be shorter thantotal_sequence_lengthand will be padded tototal_sequence_lengthwith negative infinity. Two types of masks are supported: a boolean mask where a value ofTrueindicates that the element should take part in attention, or a float mask of the same type as query, key, value that is added to the attention score.past_key (T1): Past state for key with shape
(batch_size, kv_num_heads, past_sequence_length, head_size). Must be used together withpast_valueinput.past_value (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 (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
Y (T1): The output tensor. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, v_head_size)or 3D tensor with shape(batch_size, q_sequence_length, hidden_size). For cases with a 3D input tensor,hidden_size = q_num_heads * v_head_sizepresent_key (T1): Updated key cache with shape
(batch_size, kv_num_heads, total_sequence_length, head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.present_value (T2): Updated value cache with shape
(batch_size, kv_num_heads, total_sequence_length, v_head_size)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.qk_matmul_output (T1): The output of QK matmul. 4D tensor with shape
(batch_size, q_num_heads, q_sequence_length, total_sequence_length)wheretotal_sequence_length = past_sequence_length + kv_sequence_length.
Attributes
is_causal (int): 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): 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): 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): 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): 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).
Type Constraints
T1: Constrain Q and K inputs types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).
T2: Constrain V input types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).
U: Constrain output ‘mask’ types to boolean tensors and input types. Allowed types: tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8).
Examples#
test_cc_attention_3d_local_window
Node:
Attention(Q, K, V) -> (Y)
Attributes:
q_num_heads = 4
kv_num_heads = 1
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 4, 32), dtype=float32
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]
K: shape=(2, 6, 8), dtype=float32
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]
V: shape=(2, 6, 6), dtype=float32
[[[ 0., 10., 20., 30., 40., 50.],
[ 1., 11., 21., 31., 41., 51.],
[ 2., 12., 22., 32., 42., 52.],
[ 3., 13., 23., 33., 43., 53.],
[ 4., 14., 24., 34., 44., 54.],
[ 5., 15., 25., 35., 45., 55.]],
[[ 0., 10., 20., 30., 40., 50.],
[ 1., 11., 21., 31., 41., 51.],
[ 2., 12., 22., 32., 42., 52.],
[ 3., 13., 23., 33., 43., 53.],
[ 4., 14., 24., 34., 44., 54.],
[ 5., 15., 25., 35., 45., 55.]]]
Outputs:
Y: shape=(2, 4, 24), dtype=float32
[[[ 0. , 10. , 20. , ..., 30. , 40. , 50. ],
[ 0.5, 10.5, 20.5, ..., 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , ..., 31. , 41. , 51. ],
[ 2. , 12. , 22. , ..., 32. , 42. , 52. ]],
[[ 0. , 10. , 20. , ..., 30. , 40. , 50. ],
[ 0.5, 10.5, 20.5, ..., 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , ..., 31. , 41. , 51. ],
[ 2. , 12. , 22. , ..., 32. , 42. , 52. ]]]
test_cc_attention_bidirectional_window
Node:
Attention(Q, K, V) -> (Y)
Attributes:
left_window_size = 1
right_window_size = 2
Inputs:
Q: shape=(1, 1, 5, 1), dtype=float32
[[[[0.],
[0.],
[0.],
[0.],
[0.]]]]
K: shape=(1, 1, 5, 1), dtype=float32
[[[[0.],
[0.],
[0.],
[0.],
[0.]]]]
V: shape=(1, 1, 5, 1), dtype=float32
[[[[0.],
[1.],
[2.],
[3.],
[4.]]]]
Outputs:
Y: shape=(1, 1, 5, 1), dtype=float32
[[[[1. ],
[1.5],
[2.5],
[3. ],
[3.5]]]]
test_cc_attention_local_window
Node:
Attention(Q, K, V) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 6, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 6, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]]]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 0. , 10. , 20. , ..., 50. , 60. , 70. ],
[ 0.5, 10.5, 20.5, ..., 50.5, 60.5, 70.5],
[ 1. , 11. , 21. , ..., 51. , 61. , 71. ],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ]],
[[100. , 110. , 120. , ..., 150. , 160. , 170. ],
[100.5, 110.5, 120.5, ..., 150.5, 160.5, 170.5],
[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102. , 112. , 122. , ..., 152. , 162. , 172. ]],
[[200. , 210. , 220. , ..., 250. , 260. , 270. ],
[200.5, 210.5, 220.5, ..., 250.5, 260.5, 270.5],
[201. , 211. , 221. , ..., 251. , 261. , 271. ],
[202. , 212. , 222. , ..., 252. , 262. , 272. ]]],
[[[ 0. , 10. , 20. , ..., 50. , 60. , 70. ],
[ 0.5, 10.5, 20.5, ..., 50.5, 60.5, 70.5],
[ 1. , 11. , 21. , ..., 51. , 61. , 71. ],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ]],
[[100. , 110. , 120. , ..., 150. , 160. , 170. ],
[100.5, 110.5, 120.5, ..., 150.5, 160.5, 170.5],
[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102. , 112. , 122. , ..., 152. , 162. , 172. ]],
[[200. , 210. , 220. , ..., 250. , 260. , 270. ],
[200.5, 210.5, 220.5, ..., 250.5, 260.5, 270.5],
[201. , 211. , 221. , ..., 251. , 261. , 271. ],
[202. , 212. , 222. , ..., 252. , 262. , 272. ]]]]
test_cc_attention_local_window_default
Node:
Attention(Q, K, V) -> (Y)
Attributes:
left_window_size = -1
right_window_size = -1
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 6, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 6, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]]]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5]],
[[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5]],
[[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5]]],
[[[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5]],
[[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5]],
[[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5]]]]
test_cc_attention_local_window_ext_cache_float16_mask
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float16
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 8, 8), dtype=float16
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 8, 8), dtype=float16
[[[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]]],
[[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]]]]
attn_mask: shape=(1, 8), dtype=float16
[[0., 0., 0., 0., 0., 0., 0., 0.]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[6, 7]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float16
[[[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]]],
[[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]],
[[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]]]]
test_cc_attention_local_window_ext_cache_rank2_mask
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 8, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 8, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]]]
attn_mask: shape=(1, 8), dtype=float32
[[ 0., -inf, 0., 0., 0., 0., 0., 0.]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[6, 7]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 1. , 11. , 21. , ..., 51. , 61. , 71. ],
[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ]],
[[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ]],
[[201. , 211. , 221. , ..., 251. , 261. , 271. ],
[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[203. , 213. , 223. , ..., 253. , 263. , 273. ],
[204. , 214. , 224. , ..., 254. , 264. , 274. ]]],
[[[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ],
[ 5. , 15. , 25. , ..., 55. , 65. , 75. ]],
[[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ],
[105. , 115. , 125. , ..., 155. , 165. , 175. ]],
[[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[203. , 213. , 223. , ..., 253. , 263. , 273. ],
[204. , 214. , 224. , ..., 254. , 264. , 274. ],
[205. , 215. , 225. , ..., 255. , 265. , 275. ]]]]
test_cc_attention_local_window_ext_cache_rank3_head_mask
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 8, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 8, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]]]
attn_mask: shape=(3, 4, 8), dtype=float32
[[[-inf, 0., 0., ..., 0., 0., 0.],
[-inf, 0., 0., ..., 0., 0., 0.],
[-inf, 0., 0., ..., 0., 0., 0.],
[-inf, 0., 0., ..., 0., 0., 0.]],
[[ 0., -inf, 0., ..., 0., 0., 0.],
[ 0., -inf, 0., ..., 0., 0., 0.],
[ 0., -inf, 0., ..., 0., 0., 0.],
[ 0., -inf, 0., ..., 0., 0., 0.]],
[[ 0., 0., -inf, ..., 0., 0., 0.],
[ 0., 0., -inf, ..., 0., 0., 0.],
[ 0., 0., -inf, ..., 0., 0., 0.],
[ 0., 0., -inf, ..., 0., 0., 0.]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[6, 7]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 1.5, 11.5, 21.5, ..., 51.5, 61.5, 71.5],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ]],
[[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ]],
[[200.5, 210.5, 220.5, ..., 250.5, 260.5, 270.5],
[202. , 212. , 222. , ..., 252. , 262. , 272. ],
[203.5, 213.5, 223.5, ..., 253.5, 263.5, 273.5],
[204. , 214. , 224. , ..., 254. , 264. , 274. ]]],
[[[ 2. , 12. , 22. , ..., 52. , 62. , 72. ],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ],
[ 5. , 15. , 25. , ..., 55. , 65. , 75. ]],
[[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ],
[105. , 115. , 125. , ..., 155. , 165. , 175. ]],
[[202. , 212. , 222. , ..., 252. , 262. , 272. ],
[203.5, 213.5, 223.5, ..., 253.5, 263.5, 273.5],
[204. , 214. , 224. , ..., 254. , 264. , 274. ],
[205. , 215. , 225. , ..., 255. , 265. , 275. ]]]]
test_cc_attention_local_window_ext_cache_rank4_batch_mask
Node:
Attention(Q, K, V, attn_mask, "", "", nonpad_kv_seqlen) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 8, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 8, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]]]
attn_mask: shape=(2, 1, 4, 8), dtype=float32
[[[[-inf, 0., 0., 0., 0., 0., 0., 0.],
[-inf, 0., 0., 0., 0., 0., 0., 0.],
[-inf, 0., 0., 0., 0., 0., 0., 0.],
[-inf, 0., 0., 0., 0., 0., 0., 0.]]],
[[[ 0., -inf, 0., 0., 0., 0., 0., 0.],
[ 0., -inf, 0., 0., 0., 0., 0., 0.],
[ 0., -inf, 0., 0., 0., 0., 0., 0.],
[ 0., -inf, 0., 0., 0., 0., 0., 0.]]]]
nonpad_kv_seqlen: shape=(2,), dtype=int64
[6, 7]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 1.5, 11.5, 21.5, ..., 51.5, 61.5, 71.5],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ]],
[[101.5, 111.5, 121.5, ..., 151.5, 161.5, 171.5],
[102. , 112. , 122. , ..., 152. , 162. , 172. ],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ]],
[[201.5, 211.5, 221.5, ..., 251.5, 261.5, 271.5],
[202. , 212. , 222. , ..., 252. , 262. , 272. ],
[203. , 213. , 223. , ..., 253. , 263. , 273. ],
[204. , 214. , 224. , ..., 254. , 264. , 274. ]]],
[[[ 2.5, 12.5, 22.5, ..., 52.5, 62.5, 72.5],
[ 3. , 13. , 23. , ..., 53. , 63. , 73. ],
[ 4. , 14. , 24. , ..., 54. , 64. , 74. ],
[ 5. , 15. , 25. , ..., 55. , 65. , 75. ]],
[[102.5, 112.5, 122.5, ..., 152.5, 162.5, 172.5],
[103. , 113. , 123. , ..., 153. , 163. , 173. ],
[104. , 114. , 124. , ..., 154. , 164. , 174. ],
[105. , 115. , 125. , ..., 155. , 165. , 175. ]],
[[202.5, 212.5, 222.5, ..., 252.5, 262.5, 272.5],
[203. , 213. , 223. , ..., 253. , 263. , 273. ],
[204. , 214. , 224. , ..., 254. , 264. , 274. ],
[205. , 215. , 225. , ..., 255. , 265. , 275. ]]]]
test_cc_attention_local_window_gqa_rank4_mask
Node:
Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output)
Attributes:
is_causal = 1
left_window_size = 2
softcap = 2.0
softmax_precision = 11
qk_matmul_output_mode = 3
Inputs:
Q: shape=(2, 4, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 2, 6, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 2, 6, 6), dtype=float32
[[[[ 0., 10., 20., 30., 40., 50.],
[ 1., 11., 21., 31., 41., 51.],
[ 2., 12., 22., 32., 42., 52.],
[ 3., 13., 23., 33., 43., 53.],
[ 4., 14., 24., 34., 44., 54.],
[ 5., 15., 25., 35., 45., 55.]],
[[100., 110., 120., 130., 140., 150.],
[101., 111., 121., 131., 141., 151.],
[102., 112., 122., 132., 142., 152.],
[103., 113., 123., 133., 143., 153.],
[104., 114., 124., 134., 144., 154.],
[105., 115., 125., 135., 145., 155.]]],
[[[ 0., 10., 20., 30., 40., 50.],
[ 1., 11., 21., 31., 41., 51.],
[ 2., 12., 22., 32., 42., 52.],
[ 3., 13., 23., 33., 43., 53.],
[ 4., 14., 24., 34., 44., 54.],
[ 5., 15., 25., 35., 45., 55.]],
[[100., 110., 120., 130., 140., 150.],
[101., 111., 121., 131., 141., 151.],
[102., 112., 122., 132., 142., 152.],
[103., 113., 123., 133., 143., 153.],
[104., 114., 124., 134., 144., 154.],
[105., 115., 125., 135., 145., 155.]]]]
attn_mask: shape=(2, 4, 4, 6), dtype=bool
[[[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]]],
[[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]],
[[False, False, False, False, False, False],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True],
[ True, True, True, True, True, True]]]]
Outputs:
Y: shape=(2, 4, 4, 6), dtype=float32
[[[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[ 0.5, 10.5, 20.5, 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , 31. , 41. , 51. ],
[ 2. , 12. , 22. , 32. , 42. , 52. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[ 0.5, 10.5, 20.5, 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , 31. , 41. , 51. ],
[ 2. , 12. , 22. , 32. , 42. , 52. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[100.5, 110.5, 120.5, 130.5, 140.5, 150.5],
[101. , 111. , 121. , 131. , 141. , 151. ],
[102. , 112. , 122. , 132. , 142. , 152. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[100.5, 110.5, 120.5, 130.5, 140.5, 150.5],
[101. , 111. , 121. , 131. , 141. , 151. ],
[102. , 112. , 122. , 132. , 142. , 152. ]]],
[[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[ 0.5, 10.5, 20.5, 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , 31. , 41. , 51. ],
[ 2. , 12. , 22. , 32. , 42. , 52. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[ 0.5, 10.5, 20.5, 30.5, 40.5, 50.5],
[ 1. , 11. , 21. , 31. , 41. , 51. ],
[ 2. , 12. , 22. , 32. , 42. , 52. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[100.5, 110.5, 120.5, 130.5, 140.5, 150.5],
[101. , 111. , 121. , 131. , 141. , 151. ],
[102. , 112. , 122. , 132. , 142. , 152. ]],
[[ 0. , 0. , 0. , 0. , 0. , 0. ],
[100.5, 110.5, 120.5, 130.5, 140.5, 150.5],
[101. , 111. , 121. , 131. , 141. , 151. ],
[102. , 112. , 122. , 132. , 142. , 152. ]]]]
qk_matmul_output: shape=(2, 4, 4, 6), dtype=float32
[[[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]]],
[[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]],
[[0. , 0. , 0. , 0. , 0. , 0. ],
[0.5 , 0.5 , 0. , 0. , 0. , 0. ],
[0.33333334, 0.33333334, 0.33333334, 0. , 0. , 0. ],
[0. , 0.33333334, 0.33333334, 0.33333334, 0. , 0. ]]]]
test_cc_attention_local_window_rank1_boolean_mask
Node:
Attention(Q, K, V, attn_mask) -> (Y)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 6, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 6, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
[ 3., 13., 23., ..., 53., 63., 73.],
[ 4., 14., 24., ..., 54., 64., 74.],
[ 5., 15., 25., ..., 55., 65., 75.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
[103., 113., 123., ..., 153., 163., 173.],
[104., 114., 124., ..., 154., 164., 174.],
[105., 115., 125., ..., 155., 165., 175.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
[203., 213., 223., ..., 253., 263., 273.],
[204., 214., 224., ..., 254., 264., 274.],
[205., 215., 225., ..., 255., 265., 275.]]]]
attn_mask: shape=(6,), dtype=bool
[ True, True, True, True, False, False]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 0. , 10. , 20. , ..., 50. , 60. , 70. ],
[ 0.5, 10.5, 20.5, ..., 50.5, 60.5, 70.5],
[ 1. , 11. , 21. , ..., 51. , 61. , 71. ],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ]],
[[100. , 110. , 120. , ..., 150. , 160. , 170. ],
[100.5, 110.5, 120.5, ..., 150.5, 160.5, 170.5],
[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102. , 112. , 122. , ..., 152. , 162. , 172. ]],
[[200. , 210. , 220. , ..., 250. , 260. , 270. ],
[200.5, 210.5, 220.5, ..., 250.5, 260.5, 270.5],
[201. , 211. , 221. , ..., 251. , 261. , 271. ],
[202. , 212. , 222. , ..., 252. , 262. , 272. ]]],
[[[ 0. , 10. , 20. , ..., 50. , 60. , 70. ],
[ 0.5, 10.5, 20.5, ..., 50.5, 60.5, 70.5],
[ 1. , 11. , 21. , ..., 51. , 61. , 71. ],
[ 2. , 12. , 22. , ..., 52. , 62. , 72. ]],
[[100. , 110. , 120. , ..., 150. , 160. , 170. ],
[100.5, 110.5, 120.5, ..., 150.5, 160.5, 170.5],
[101. , 111. , 121. , ..., 151. , 161. , 171. ],
[102. , 112. , 122. , ..., 152. , 162. , 172. ]],
[[200. , 210. , 220. , ..., 250. , 260. , 270. ],
[200.5, 210.5, 220.5, ..., 250.5, 260.5, 270.5],
[201. , 211. , 221. , ..., 251. , 261. , 271. ],
[202. , 212. , 222. , ..., 252. , 262. , 272. ]]]]
test_cc_attention_local_window_with_past
Node:
Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
Attributes:
is_causal = 1
left_window_size = 2
Inputs:
Q: shape=(2, 3, 4, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
K: shape=(2, 3, 2, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
V: shape=(2, 3, 2, 8), dtype=float32
[[[[ 8., 18., 28., ..., 58., 68., 78.],
[ 9., 19., 29., ..., 59., 69., 79.]],
[[108., 118., 128., ..., 158., 168., 178.],
[109., 119., 129., ..., 159., 169., 179.]],
[[208., 218., 228., ..., 258., 268., 278.],
[209., 219., 229., ..., 259., 269., 279.]]],
[[[ 8., 18., 28., ..., 58., 68., 78.],
[ 9., 19., 29., ..., 59., 69., 79.]],
[[108., 118., 128., ..., 158., 168., 178.],
[109., 119., 129., ..., 159., 169., 179.]],
[[208., 218., 228., ..., 258., 268., 278.],
[209., 219., 229., ..., 259., 269., 279.]]]]
past_key: shape=(2, 3, 8, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
past_value: shape=(2, 3, 8, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 5., 15., 25., ..., 55., 65., 75.],
[ 6., 16., 26., ..., 56., 66., 76.],
[ 7., 17., 27., ..., 57., 67., 77.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[105., 115., 125., ..., 155., 165., 175.],
[106., 116., 126., ..., 156., 166., 176.],
[107., 117., 127., ..., 157., 167., 177.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[205., 215., 225., ..., 255., 265., 275.],
[206., 216., 226., ..., 256., 266., 276.],
[207., 217., 227., ..., 257., 267., 277.]]]]
Outputs:
Y: shape=(2, 3, 4, 8), dtype=float32
[[[[ 7. , 17. , 27. , ..., 57. , 67. , 77. ],
[ 8. , 18. , 28. , ..., 58. , 68. , 78. ],
[ 8.5, 18.5, 28.5, ..., 58.5, 68.5, 78.5],
[ 9. , 19. , 29. , ..., 59. , 69. , 79. ]],
[[107. , 117. , 127. , ..., 157. , 167. , 177. ],
[108. , 118. , 128. , ..., 158. , 168. , 178. ],
[108.5, 118.5, 128.5, ..., 158.5, 168.5, 178.5],
[109. , 119. , 129. , ..., 159. , 169. , 179. ]],
[[207. , 217. , 227. , ..., 257. , 267. , 277. ],
[208. , 218. , 228. , ..., 258. , 268. , 278. ],
[208.5, 218.5, 228.5, ..., 258.5, 268.5, 278.5],
[209. , 219. , 229. , ..., 259. , 269. , 279. ]]],
[[[ 7. , 17. , 27. , ..., 57. , 67. , 77. ],
[ 8. , 18. , 28. , ..., 58. , 68. , 78. ],
[ 8.5, 18.5, 28.5, ..., 58.5, 68.5, 78.5],
[ 9. , 19. , 29. , ..., 59. , 69. , 79. ]],
[[107. , 117. , 127. , ..., 157. , 167. , 177. ],
[108. , 118. , 128. , ..., 158. , 168. , 178. ],
[108.5, 118.5, 128.5, ..., 158.5, 168.5, 178.5],
[109. , 119. , 129. , ..., 159. , 169. , 179. ]],
[[207. , 217. , 227. , ..., 257. , 267. , 277. ],
[208. , 218. , 228. , ..., 258. , 268. , 278. ],
[208.5, 218.5, 228.5, ..., 258.5, 268.5, 278.5],
[209. , 219. , 229. , ..., 259. , 269. , 279. ]]]]
present_key: shape=(2, 3, 10, 8), dtype=float32
[[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]],
[[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]],
[[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]]]]
present_value: shape=(2, 3, 10, 8), dtype=float32
[[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 7., 17., 27., ..., 57., 67., 77.],
[ 8., 18., 28., ..., 58., 68., 78.],
[ 9., 19., 29., ..., 59., 69., 79.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[107., 117., 127., ..., 157., 167., 177.],
[108., 118., 128., ..., 158., 168., 178.],
[109., 119., 129., ..., 159., 169., 179.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[207., 217., 227., ..., 257., 267., 277.],
[208., 218., 228., ..., 258., 268., 278.],
[209., 219., 229., ..., 259., 269., 279.]]],
[[[ 0., 10., 20., ..., 50., 60., 70.],
[ 1., 11., 21., ..., 51., 61., 71.],
[ 2., 12., 22., ..., 52., 62., 72.],
...,
[ 7., 17., 27., ..., 57., 67., 77.],
[ 8., 18., 28., ..., 58., 68., 78.],
[ 9., 19., 29., ..., 59., 69., 79.]],
[[100., 110., 120., ..., 150., 160., 170.],
[101., 111., 121., ..., 151., 161., 171.],
[102., 112., 122., ..., 152., 162., 172.],
...,
[107., 117., 127., ..., 157., 167., 177.],
[108., 118., 128., ..., 158., 168., 178.],
[109., 119., 129., ..., 159., 169., 179.]],
[[200., 210., 220., ..., 250., 260., 270.],
[201., 211., 221., ..., 251., 261., 271.],
[202., 212., 222., ..., 252., 262., 272.],
...,
[207., 217., 227., ..., 257., 267., 277.],
[208., 218., 228., ..., 258., 268., 278.],
[209., 219., 229., ..., 259., 269., 279.]]]]
Differences with previous version (24)#
SchemaDiff: Attention (domain 'ai.onnx')
old version: 24
new version: 25
breaking: no
Attributes:
added ‘is_causal’: type=INT; required=False; default=0
added ‘scale’: type=FLOAT; required=False; default=UNDEFINED
added ‘q_num_heads’: type=INT; required=False; default=UNDEFINED
added ‘kv_num_heads’: type=INT; required=False; default=UNDEFINED
added ‘softmax_precision’: type=INT; required=False; default=UNDEFINED
added ‘softcap’: type=FLOAT; required=False; default=0
added ‘qk_matmul_output_mode’: type=INT; required=False; default=0
added ‘left_window_size’: type=INT; required=False; default=-1
added ‘right_window_size’: type=INT; required=False; default=-1
Documentation:
line similarity: 0.64 (+60/-3 lines)
--- Attention v24
+++ Attention v25
@@ -15,9 +15,66 @@
Attention bias to be added is calculated based on `attn_mask` input and `is_causal` attribute:
1) `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.
-2) If `is_causal` is set to `1`, attention scores above the causal frontier are masked out. For internal cache (`past_key`) this is the standard offset from `past_sequence_length`; for external cache (`nonpad_kv_seqlen` without `past_key`) this is bottom-right aligned by `nonpad_kv_seqlen - q_sequence_length`.
-3) If both `attn_mask` and `is_causal` are set, the valid positions are the intersection of both masks.
-If a query row is fully masked after this intersection, its output row is zero.
+2) If `is_causal` is set to `1`, causal masking is applied with bottom-right (offset-aware) alignment: query `i` attends key `j` iff `j <= i + offset`, as illustrated below.
+
+```
+ 2D causal mask for Attention (PR onnx/onnx#8068)
+ S_q=4 queries, S_k=8 keys
+ Rule: query i attends key j iff j <= i + offset
+ offset = nonpad_kv_seqlen - S_q
+
+ nonpad_kv_seqlen=4, offset=4-4=0
+
+ k0 k1 k2 k3 k4 k5 k6 k7
+ +----+----+----+----+----+----+----+----+
+ q0 | ## | | | | | | | |
+ +----+----+----+----+----+----+----+----+
+ q1 | ## | ## | | | | | | |
+ +----+----+----+----+----+----+----+----+
+ q2 | ## | ## | ## | | | | | |
+ +----+----+----+----+----+----+----+----+
+ q3 | ## | ## | ## | ## | | | | |
+ +----+----+----+----+----+----+----+----+
+
+
+ nonpad_kv_seqlen=8, offset=8-4=4
+
+ k0 k1 k2 k3 k4 k5 k6 k7
+ +----+----+----+----+----+----+----+----+
+ q0 | ## | ## | ## | ## | ## | | | |
+ +----+----+----+----+----+----+----+----+
+ q1 | ## | ## | ## | ## | ## | ## | | |
+ +----+----+----+----+----+----+----+----+
+ q2 | ## | ## | ## | ## | ## | ## | ## | |
+ +----+----+----+----+----+----+----+----+
+ q3 | ## | ## | ## | ## | ## | ## | ## | ## |
+ +----+----+----+----+----+----+----+----+
+```
+
+With `nonpad_kv_seqlen=4` (offset=0), the mask is the standard lower-triangular. With `nonpad_kv_seqlen=8` (offset=4), the diagonal shifts right by 4, so each query sees the 4 additional valid cached keys.
+
+`offset` is the count of valid keys preceding the current query block: `offset = past_sequence_length` when `past_key` is provided; `offset = nonpad_kv_seqlen - q_sequence_length` (per batch) when an external cache is indicated by `nonpad_kv_seqlen` without `past_key`; `offset = 0` when neither is provided (the no-cache case, which reduces to the standard lower-triangular mask). When `offset < 0` (`nonpad_kv_seqlen < q_sequence_length`, i.e. more query tokens than cached keys) the leading query rows have an empty key set (no key satisfies `j <= i + offset`) and are fully masked. The causal frontier is computed independently of `attn_mask` and is then composed with it additively: a boolean `attn_mask` intersects the allowed set (its disallowed positions contribute `-inf` to the bias), while a float `attn_mask` is added to the attention scores rather than disabling positions. A fully-masked query row (no key attended, including the negative-offset leading rows) produces a zero output row, not `NaN`, for both `Y` and the mode-`3` `qk_matmul_output` debug output; the mode-`3` `qk_matmul_output` is emitted at the operator's output precision (`T1`).
+
+`left_window_size` and `right_window_size` independently restrict the keys visible to each query. A query at absolute position `p = offset + query_index` attends keys `j` satisfying `p - left_window_size <= j <= p + right_window_size` for each nonnegative bound. A value of `-1` leaves that side unbounded. For example, `(left_window_size=2, right_window_size=0)` is a causal left-looking window containing the current key and two preceding keys, while `(left_window_size=2, right_window_size=1)` is an asymmetric bidirectional window. Window bounds are composed with `is_causal` and `attn_mask`; when `is_causal=1`, the causal upper bound still excludes future keys.
+
+```
+ 2D sliding-window mask for Attention (opset 25)
+ S_q=4 queries, S_k=6 keys, left_window_size=2, right_window_size=1, offset=0
+
+ k0 k1 k2 k3 k4 k5
+ +----+----+----+----+----+----+
+ q0 | ## | ## | | | | |
+ +----+----+----+----+----+----+
+ q1 | ## | ## | ## | | | |
+ +----+----+----+----+----+----+
+ q2 | ## | ## | ## | ## | | |
+ +----+----+----+----+----+----+
+ q3 | | ## | ## | ## | ## | |
+ +----+----+----+----+----+----+
+
+ q0 attends {k0,k1}, q1 attends {k0,k1,k2}, q2 attends {k0,k1,k2,k3},
+ q3 attends {k1,k2,k3,k4}.
+```
With respect to KV cache update, this operator allows the following two use cases: