| |
| |
| |
| |
| |
| |
|
|
| import torch |
|
|
| from ...ops.simple_gla.parallel import parallel_simple_gla |
|
|
|
|
| def parallel_retention( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| scale: float | None = None, |
| output_attentions: bool = False, |
| cu_seqlens: torch.LongTensor | None = None, |
| **kwargs, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| r""" |
| Args: |
| q (torch.Tensor): |
| queries of shape `[B, T, H, K]`. |
| k (torch.Tensor): |
| keys of shape `[B, T, H, K]`. |
| v (torch.Tensor): |
| values of shape `[B, T, H, V]`. |
| scale (Optional[float]): |
| Scale factor for attention scores. |
| If not provided, it will default to `1 / sqrt(K)`. Default: `None`. |
| output_attentions (bool): |
| Whether to output the materialized attention scores of shape `[B, H, T, T]`. Default: `False`. |
| cu_seqlens (torch.LongTensor): |
| Cumulative sequence lengths of shape `[N+1]` used for variable-length training, |
| consistent with the FlashAttention API. |
| |
| Returns: |
| o (torch.Tensor): |
| Outputs of shape `[B, T, H, V]`. |
| attn (torch.Tensor): |
| Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. |
| """ |
| if 'head_first' in kwargs: |
| raise DeprecationWarning( |
| "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", |
| ) |
| if cu_seqlens is not None and q.shape[0] != 1: |
| raise ValueError( |
| f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " |
| f"Please flatten variable-length inputs before processing.", |
| ) |
| s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() |
| g = s[None, None, :].expand(q.shape[0], q.shape[1], q.shape[2]) |
|
|
| o, attn = parallel_simple_gla( |
| q=q, |
| k=k, |
| v=v, |
| scale=scale, |
| g=g, |
| output_attentions=output_attentions, |
| cu_seqlens=cu_seqlens, |
| ) |
| return o, attn |
|
|
