wehub-resource-sync
126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
"""Attention operator in python"""
|
|
|
|
import numpy as np
|
|
|
|
from .softmax_python import softmax_python
|
|
|
|
|
|
def attention_python(
|
|
q: np.ndarray,
|
|
k: np.ndarray,
|
|
v: np.ndarray,
|
|
bias: np.ndarray | None,
|
|
qk_scale: float,
|
|
causal: str,
|
|
window_size: int | None = None,
|
|
layout: str = "BSNH",
|
|
): # pylint: disable=too-many-arguments, too-many-locals, invalid-name
|
|
"""Attention operator in python
|
|
|
|
Parameters
|
|
----------
|
|
q : np.ndarray
|
|
Query tensor with shape [batch, seq_length, num_heads, head_dim] in the layout specified by
|
|
`layout`.
|
|
k : np.ndarray
|
|
Key tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim] in the layout specified
|
|
by `layout`.
|
|
v : np.ndarray
|
|
Value tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim_v] in the layout
|
|
specified by `layout`.
|
|
bias : np.ndarray
|
|
Bias tensor with shape [batch, num_heads, seq_length, seq_length]
|
|
qk_scale : float
|
|
Scale factor for the query-key product.
|
|
causal : str
|
|
The type of causal mask to apply. Can be "none", "TopLeft", or "BottomRight".
|
|
window_size : Optional[int]
|
|
The window size for the causal mask.
|
|
layout : str
|
|
The layout of the input tensors, e.g. "BSNH" or "BNSH".
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
The output tensor with shape [batch, seq_length, num_heads, head_dim_v] in the layout
|
|
specified by `layout`.
|
|
"""
|
|
assert layout in ["BSNH", "BNSH", "SBNH"]
|
|
|
|
dim_b = layout.find("B")
|
|
dim_s = layout.find("S")
|
|
dim_n = layout.find("N")
|
|
dim_h = layout.find("H")
|
|
|
|
q = q.transpose(dim_b, dim_n, dim_s, dim_h) # b, n, s, h
|
|
k = k.transpose(dim_b, dim_n, dim_s, dim_h) # b, n, s_kv, h
|
|
kt = k.transpose(0, 1, 3, 2) # b, n, h, s_kv
|
|
v = v.transpose(dim_b, dim_n, dim_s, dim_h)
|
|
|
|
num_heads = q.shape[1]
|
|
num_kv_heads = k.shape[1]
|
|
s = q.shape[2]
|
|
s_kv = k.shape[2]
|
|
|
|
if num_heads != num_kv_heads:
|
|
assert num_heads % num_kv_heads == 0
|
|
factor = num_heads // num_kv_heads
|
|
kt = np.repeat(kt, factor, axis=1)
|
|
v = np.repeat(v, factor, axis=1)
|
|
|
|
if not qk_scale == "none":
|
|
score = q @ kt * qk_scale # b, n, s, s_kv
|
|
else:
|
|
score = q @ kt / np.sqrt(q.shape[-1]) # b, n, s, s_kv
|
|
if bias is not None:
|
|
score = score + bias # b, n, s, s_kv
|
|
if causal == "none":
|
|
attn = softmax_python(score, -1)
|
|
else:
|
|
if causal == "TopLeft":
|
|
offset = 0
|
|
elif causal == "BottomRight":
|
|
offset = abs(s - s_kv)
|
|
else:
|
|
raise ValueError(f"Unsupported causal type: {causal}")
|
|
score_masked = np.tril(score, k=offset)
|
|
|
|
if window_size:
|
|
score_masked = np.triu(
|
|
score_masked,
|
|
-window_size + 1, # pylint: disable=invalid-unary-operand-type
|
|
)
|
|
|
|
score_masked_exp = np.tril(
|
|
np.exp(score_masked - np.max(score_masked, axis=-1, keepdims=True)), k=offset
|
|
)
|
|
|
|
if window_size:
|
|
score_masked_exp = np.triu(
|
|
score_masked_exp,
|
|
-window_size + 1, # pylint: disable=invalid-unary-operand-type
|
|
)
|
|
|
|
score_masked_sum = np.sum(score_masked_exp, axis=-1, keepdims=True)
|
|
attn = np.divide(score_masked_exp, score_masked_sum)
|
|
|
|
out = attn @ v # b, n, s, h_v
|
|
return out.transpose(*np.argsort([dim_b, dim_n, dim_s, dim_h]).tolist())
|