| |
|
|
| import warnings |
|
|
| import torch |
|
|
| from fla.ops.attn.parallel import parallel_attn |
|
|
|
|
| def parallel_forgetting_attn( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| g: torch.Tensor, |
| scale: float | None = None, |
| cu_seqlens: torch.LongTensor | None = None, |
| head_first: bool = False, |
| ) -> torch.Tensor: |
| r""" |
| Args: |
| q (torch.Tensor): |
| queries of shape `[B, T, HQ, K]`. |
| k (torch.Tensor): |
| keys of shape `[B, T, H, K]`. |
| GQA will be applied if HQ is divisible by H. |
| v (torch.Tensor): |
| values of shape `[B, T, H, V]`. |
| g (torch.Tensor): |
| Log decay at rach time step (in **log space**) of shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. |
| scale (Optional[float]): |
| Scale factor for attention scores. |
| If not provided, it will default to `1 / sqrt(K)`. Default: `None`. |
| cu_seqlens (torch.LongTensor): |
| Cumulative sequence lengths of shape `[N+1]` used for variable-length training, |
| consistent with the FlashAttention API. |
| head_first (Optional[bool]): |
| Whether the inputs are in the head-first format. Default: `False`. |
| This argument has been deprecated. |
| |
| Returns: |
| o (torch.Tensor): |
| Outputs of shape `[B, T, HQ, V]`. |
| """ |
|
|
| if scale is None: |
| scale = k.shape[-1] ** -0.5 |
| if cu_seqlens is not None: |
| assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" |
| if head_first: |
| raise DeprecationWarning( |
| "head_first is deprecated and will be removed in a future version. " |
| "Please use head_first=False for now instead.", |
| ) |
| if not head_first and q.shape[1] < q.shape[2]: |
| warnings.warn( |
| f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " |
| "This may indicate the inputs were passed in head-first format [B, H, T, ...] " |
| "when head_first=False was specified. " |
| "Please verify your input tensor format matches the expected shape [B, T, H, ...].", |
| ) |
| o = parallel_attn(q, k, v, g, scale, cu_seqlens) |
| return o |
|
|
