wehub-resource-sync
250 lines
8.5 KiB
Python
250 lines
8.5 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed 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.
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
from einops import rearrange, repeat
|
|
|
|
import paddle
|
|
|
|
|
|
def construct_local_mask(
|
|
seqlen_q,
|
|
seqlen_k,
|
|
window_size=(-1, -1), # -1 means infinite window size
|
|
sink_token_length=0,
|
|
query_padding_mask=None,
|
|
key_padding_mask=None,
|
|
key_leftpad=None,
|
|
):
|
|
row_idx = rearrange(paddle.arange(seqlen_q, dtype=paddle.int64), "s -> s 1")
|
|
col_idx = paddle.arange(seqlen_k, dtype=paddle.int64)
|
|
if key_leftpad is not None:
|
|
key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
|
|
col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
|
|
col_idx = paddle.where(
|
|
col_idx >= key_leftpad, col_idx - key_leftpad, 2**32
|
|
)
|
|
sk = (
|
|
seqlen_k
|
|
if key_padding_mask is None
|
|
else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
|
|
)
|
|
sq = (
|
|
seqlen_q
|
|
if query_padding_mask is None
|
|
else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
|
|
)
|
|
if window_size[0] < 0:
|
|
return col_idx > row_idx + sk - sq + window_size[1]
|
|
else:
|
|
sk = (
|
|
paddle.full_like(col_idx, seqlen_k)
|
|
if key_padding_mask is None
|
|
else sk
|
|
)
|
|
return paddle.logical_or(
|
|
col_idx > paddle.minimum(row_idx + sk - sq + window_size[1], sk),
|
|
paddle.logical_and(
|
|
col_idx < row_idx + sk - sq - window_size[0],
|
|
col_idx >= sink_token_length,
|
|
),
|
|
)
|
|
|
|
|
|
def attention_ref(
|
|
q,
|
|
k,
|
|
v,
|
|
query_padding_mask=None,
|
|
key_padding_mask=None,
|
|
key_leftpad=None,
|
|
attn_bias=None,
|
|
dropout_p=0.0,
|
|
dropout_mask=None,
|
|
causal=False,
|
|
qv=None,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
window_size=(-1, -1), # -1 means infinite window size
|
|
attention_chunk=0,
|
|
sink_token_length=0,
|
|
softcap=0.0,
|
|
upcast=True,
|
|
reorder_ops=False,
|
|
intermediate_dtype=None,
|
|
):
|
|
"""
|
|
Arguments:
|
|
q: (batch_size, seqlen_q, nheads, head_dim)
|
|
k: (batch_size, seqlen_k, nheads, head_dim)
|
|
v: (batch_size, seqlen_k, nheads, head_dim_v)
|
|
qv: (batch_size, seqlen_q, nheads, head_dim_v)
|
|
query_padding_mask: (batch_size, seqlen_q)
|
|
key_padding_mask: (batch_size, seqlen_k)
|
|
attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k)
|
|
dropout_p: float
|
|
dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k)
|
|
causal: whether to apply causal masking
|
|
upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast
|
|
output back to fp16/bf16.
|
|
reorder_ops: whether to change the order of operations (scaling k instead of scaling k, etc.)
|
|
without changing the math. This is to estimate the numerical error from operation
|
|
reordering.
|
|
Output:
|
|
output: (batch_size, seqlen_q, nheads, head_dim_v)
|
|
attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout
|
|
"""
|
|
if causal:
|
|
window_size = (window_size[0], 0)
|
|
dtype_og = q.dtype
|
|
if upcast:
|
|
q = paddle.cast(q, paddle.float32)
|
|
k = paddle.cast(k, paddle.float32)
|
|
v = paddle.cast(v, paddle.float32)
|
|
if qv is not None:
|
|
qv = paddle.cast(qv, paddle.float32)
|
|
|
|
if q_descale is not None:
|
|
raise AssertionError
|
|
q_descale = repeat(
|
|
q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2]
|
|
)
|
|
q = (q.cast(paddle.float32) * q_descale).cast(q.dtype)
|
|
qv = (
|
|
(qv.cast(paddle.float32) * q_descale).cast(qv.dtype)
|
|
if qv is not None
|
|
else None
|
|
)
|
|
|
|
if k_descale is not None:
|
|
raise AssertionError
|
|
k = (
|
|
k.cast(paddle.float32) * rearrange(k_descale, "b h -> b 1 h 1")
|
|
).cast(k.dtype)
|
|
|
|
if v_descale is not None:
|
|
raise AssertionError
|
|
v = (
|
|
v.cast(paddle.float32) * rearrange(v_descale, "b h -> b 1 h 1")
|
|
).cast(v.dtype)
|
|
|
|
seqlen_q, seqlen_k = q.shape[1], k.shape[1]
|
|
|
|
# (batch_size, seqlen, nheads, head_dim) -> (batch_size, nheads, seqlen, head_dim)
|
|
q = paddle.transpose(q, [0, 2, 1, 3])
|
|
k = paddle.transpose(k, [0, 2, 1, 3])
|
|
v = paddle.transpose(v, [0, 2, 1, 3])
|
|
|
|
k = repeat(k, "b h s d -> b (h g) s d", g=q.shape[1] // k.shape[1])
|
|
v = repeat(v, "b h s d -> b (h g) s d", g=q.shape[1] // v.shape[1])
|
|
d = q.shape[-1]
|
|
dv = v.shape[-1]
|
|
softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv)
|
|
|
|
if not reorder_ops:
|
|
scores = paddle.matmul(q * softmax_scale, k, transpose_y=True)
|
|
else:
|
|
scores = paddle.matmul(q, k * softmax_scale, transpose_y=True)
|
|
|
|
if qv is not None:
|
|
raise AssertionError
|
|
scores = scores + paddle.matmul(qv * softmax_scale, v, transpose_y=True)
|
|
|
|
if softcap > 0:
|
|
raise AssertionError
|
|
scores = paddle.tanh(scores / softcap) * softcap
|
|
|
|
if key_padding_mask is not None:
|
|
raise AssertionError
|
|
scores.masked_fill_(
|
|
rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")
|
|
)
|
|
local_mask = None
|
|
|
|
if window_size[0] >= 0 or window_size[1] >= 0:
|
|
local_mask = construct_local_mask(
|
|
seqlen_q,
|
|
seqlen_k,
|
|
window_size,
|
|
sink_token_length,
|
|
query_padding_mask,
|
|
key_padding_mask,
|
|
key_leftpad=key_leftpad,
|
|
)
|
|
if attention_chunk > 0:
|
|
raise AssertionError
|
|
|
|
if local_mask is not None:
|
|
scores.masked_fill_(local_mask, float("-inf"))
|
|
if attn_bias is not None:
|
|
scores = scores + attn_bias
|
|
# when all values in a line of attn_bias are -inf, setting value in this line to a very small value
|
|
# to prevent softmax giving nan output
|
|
all_inf_mask = (attn_bias == -np.inf).all(axis=-1, keepdim=True)
|
|
scores = paddle.where(
|
|
all_inf_mask, paddle.full_like(scores, -1e9), scores
|
|
)
|
|
|
|
attention = paddle.nn.functional.softmax(scores, axis=-1).cast(v.dtype)
|
|
|
|
if attn_bias is not None:
|
|
# when all values in a line of attn_bias are -inf, we setting value in this line to a very small value
|
|
# to prevent softmax giving nan output, however, after softmax, values in this line become 1/seqlen,
|
|
# so setting them to 0 after softmax
|
|
attention = paddle.where(
|
|
all_inf_mask, paddle.zeros_like(attention), attention
|
|
)
|
|
|
|
# We want to mask here so that the attention matrix doesn't have any NaNs
|
|
# Otherwise we'll get NaN in dV
|
|
if query_padding_mask is not None:
|
|
raise AssertionError
|
|
attention = attention.masked_fill(
|
|
rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0
|
|
)
|
|
|
|
# Without this we might get NaN in dv
|
|
if key_padding_mask is not None:
|
|
raise AssertionError
|
|
attention = attention.masked_fill(
|
|
rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0
|
|
)
|
|
# Some rows might be completely masked out so we fill them with zero instead of NaN
|
|
if local_mask is not None:
|
|
attention = attention.masked_fill(
|
|
paddle.all(local_mask, axis=-1, keepdim=True), 0.0
|
|
)
|
|
dropout_scaling = 1.0 / (1 - dropout_p)
|
|
# attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling
|
|
# output = paddle.matmul(attention_drop, v, transpose_y=True)
|
|
if dropout_mask is not None:
|
|
raise AssertionError
|
|
attention_drop = attention.masked_fill(~dropout_mask, 0.0)
|
|
else:
|
|
attention_drop = attention
|
|
if intermediate_dtype is not None:
|
|
attention_drop = attention_drop.cast(intermediate_dtype).cast(
|
|
attention_drop.dtype
|
|
)
|
|
output = paddle.matmul(attention_drop, v * dropout_scaling)
|
|
output = paddle.transpose(output, [0, 2, 1, 3])
|
|
if query_padding_mask is not None:
|
|
output.masked_fill_(
|
|
rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0
|
|
)
|
|
return output.cast(dtype=dtype_og), attention.cast(dtype=dtype_og)
|