add sink_attention (#2461)

Co-authored-by: umiswing <[email protected]>

Author: xxyux

paddleformers/nn/attention/__init__.py

@@ -24,13 +24,15 @@
     "interface": ["AttentionInterface", "ALL_ATTENTION_FUNCTIONS"],
     "sdpa_attention": ["sdpa_attention_forward"],
     "utils": ["repeat_kv"],
+    "sink_impl": ["sink_attention_forward"],
 }
 
 if TYPE_CHECKING:
     from .eager_attention import *
     from .flashmask_attention import *
     from .interface import *
     from .sdpa_attention import *
+    from .sink_impl import *
     from .utils import *
 else:
     sys.modules[__name__] = _LazyModule(

paddleformers/nn/attention/eager_attention.py

@@ -40,6 +40,7 @@ def eager_attention_forward(
     query = paddle.transpose(x=query, perm=perm)
     key = paddle.transpose(x=key, perm=perm)
     value = paddle.transpose(x=value, perm=perm)
+
     attn_weights = paddle.matmul(query, key.transpose([0, 1, 3, 2])) * scaling
     if attention_mask is not None:
         causal_mask = attention_mask[:, :, :, : key.shape[-2]]

paddleformers/nn/attention/flashmask_attention.py

@@ -18,30 +18,41 @@
 import paddle.nn as nn
 from paddle.nn.functional.flash_attention import flashmask_attention
 
+from .sink_impl import sink_attention_forward
+
 
 def flashmask_attention_forward(
     module: nn.Layer,
     query: paddle.Tensor,
     key: paddle.Tensor,
     value: paddle.Tensor,
-    attention_mask: Optional[paddle.Tensor] = None,
-    attn_mask_start_row_indices=None,
+    attn_mask_start_row_indices: paddle.Tensor,
     dropout: float = 0.0,
+    sink: Optional[paddle.Tensor] = None,
     scaling: Optional[float] = None,
     is_causal: Optional[bool] = None,
     **kwargs
 ):
-    if attn_mask_start_row_indices is not None:
-        attn_mask_start_row_indices = attn_mask_start_row_indices.unsqueeze(-1)
-
     # b,l,h,d
-    out = flashmask_attention(
-        query,
-        key,
-        value,
-        startend_row_indices=attn_mask_start_row_indices,
-        causal=True,
-    )
+    if sink is None:
+        out = flashmask_attention(
+            query,
+            key,
+            value,
+            startend_row_indices=attn_mask_start_row_indices,
+            causal=True,
+        )
+    else:
+        out = sink_attention_forward(
+            query,
+            key,
+            value,
+            sink,
+            startend_row_indices=attn_mask_start_row_indices,
+            dropout_p=dropout,
+            softmax_scale=scaling,
+            causal=is_causal,
+        )
     out = paddle.reshape(x=out, shape=[0, 0, out.shape[2] * out.shape[3]])
 
     return out, None

paddleformers/nn/attention/sdpa_attention.py

@@ -18,6 +18,7 @@
 import paddle.nn as nn
 
 from ...utils.masking_utils import _gen_from_sparse_attn_mask_indices
+from .sink_impl import sink_attention_forward
 
 
 def sdpa_attention_forward(
@@ -28,20 +29,33 @@ def sdpa_attention_forward(
     attention_mask: Optional[paddle.Tensor] = None,
     attn_mask_start_row_indices=None,
     dropout: float = 0.0,
+    sink: Optional[paddle.Tensor] = None,
     scaling: Optional[float] = None,
     is_causal: Optional[bool] = None,
     **kwargs,
 ):
     # query: b l h d
-
     if is_causal is None and attn_mask_start_row_indices is None:
         is_causal = query.shape[1] > 1 and attention_mask is None and getattr(module, "is_causal", True)
     elif attn_mask_start_row_indices is not None:
         is_causal = False
         attention_mask = _gen_from_sparse_attn_mask_indices(attn_mask_start_row_indices, query.dtype)
 
-    attn_output = nn.functional.scaled_dot_product_attention(
-        query, key, value, attention_mask, dropout, is_causal=is_causal, training=module.training
-    )
+    if sink is None:
+        attn_output = nn.functional.scaled_dot_product_attention(
+            query, key, value, attention_mask, dropout, is_causal=is_causal, training=module.training
+        )
+    else:
+        attn_output = sink_attention_forward(
+            query,
+            key,
+            value,
+            sink,
+            attention_mask=attention_mask,
+            startend_row_indices=None,
+            dropout_p=dropout,
+            softmax_scale=scaling,
+            causal=is_causal,
+        )
     attn_output = paddle.reshape(x=attn_output, shape=[0, 0, attn_output.shape[2] * attn_output.shape[3]])
     return attn_output, None