| |
|
|
| import warnings |
|
|
| import torch |
|
|
| from fla.ops.generalized_delta_rule import chunk_dplr_delta_rule |
|
|
|
|
| def chunk_rwkv7( |
| r: torch.Tensor, |
| w: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| a: torch.Tensor, |
| b: torch.Tensor, |
| scale: float = 1.0, |
| initial_state: torch.Tensor = None, |
| output_final_state: bool = True, |
| cu_seqlens: torch.LongTensor | None = None, |
| head_first: bool = False, |
| ): |
| """ |
| Args: |
| r (torch.Tensor): |
| r of shape `[B, T, H, K]`. |
| w (torch.Tensor): |
| log decay of shape `[B, T, H, K]`. |
| k (torch.Tensor): |
| k of shape `[B, T, H, K]`. |
| v (torch.Tensor): |
| v of shape `[B, T, H, V]`. |
| a (torch.Tensor): |
| a of shape `[B, T, H, K]`. |
| b (torch.Tensor): |
| b of shape `[B, T, H, K]`. |
| scale (float): |
| scale of the attention. |
| initial_state (Optional[torch.Tensor]): |
| Initial state of shape `[N, H, K, V]` for `N` input sequences. |
| For equal-length input sequences, `N` equals the batch size `B`. |
| Default: `None`. |
| output_final_state (Optional[bool]): |
| Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. |
| 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. |
| """ |
| 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 r.shape[1] < r.shape[2]: |
| warnings.warn( |
| f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.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, ...].", |
| ) |
| return chunk_dplr_delta_rule( |
| q=r, |
| k=k, |
| v=v, |
| a=a, |
| b=b, |
| gk=w, |
| scale=scale, |
| initial_state=initial_state, |
| output_final_state=output_final_state, |
| cu_seqlens=cu_seqlens, |
| head_first=head_first, |
| ) |
|
|
