| |
|
|
|
|
| import torch |
|
|
| from fla.ops.gated_delta_rule.fused_recurrent import fused_recurrent_gated_delta_rule |
|
|
|
|
| def fused_recurrent_kda( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| g: torch.Tensor, |
| beta: torch.Tensor = None, |
| scale: float = None, |
| initial_state: torch.Tensor = None, |
| output_final_state: bool = False, |
| use_qk_l2norm_in_kernel: 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, HV, V]`. |
| GVA is applied if `HV > H`. |
| g (torch.Tensor): |
| g (decays) of shape `[B, T, HV]`. |
| beta (torch.Tensor): |
| betas of shape `[B, T, HV]`. |
| scale (Optional[float]): |
| Scale factor for the RetNet attention scores. |
| If not provided, it will default to `1 / sqrt(K)`. Default: `None`. |
| initial_state (Optional[torch.Tensor]): |
| Initial state of shape `[N, HV, 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, HV, K, V]`. Default: `False`. |
| use_qk_l2norm_in_kernel (Optional[bool]): |
| Whether to use L2 normalization in the kernel. 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, HV, V]`. |
| final_state (torch.Tensor): |
| Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. |
| |
| Examples:: |
| >>> import torch |
| >>> import torch.nn.functional as F |
| >>> from einops import rearrange |
| >>> from fla.ops.kda import fused_recurrent_kda |
| # inputs with equal lengths |
| >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 |
| >>> q = torch.randn(B, T, H, K, device='cuda') |
| >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) |
| >>> v = torch.randn(B, T, HV, V, device='cuda') |
| >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda')) |
| >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() |
| >>> h0 = torch.randn(B, HV, K, V, device='cuda') |
| >>> o, ht = fused_recurrent_kda( |
| q, k, v, g, beta, |
| initial_state=h0, |
| output_final_state=True |
| ) |
| # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required |
| >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) |
| # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected |
| >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) |
| >>> o_var, ht_var = fused_recurrent_kda( |
| q, k, v, g, beta, |
| initial_state=h0, |
| output_final_state=True, |
| cu_seqlens=cu_seqlens |
| ) |
| """ |
|
|
| if cu_seqlens is not None: |
| if 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.", |
| ) |
| if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: |
| raise ValueError( |
| f"The number of initial states is expected to be equal to the number of input sequences, " |
| f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", |
| ) |
| if scale is None: |
| scale = k.shape[-1] ** -0.5 |
|
|
| o, final_state = fused_recurrent_gated_delta_rule( |
| q=q, |
| k=k, |
| v=v, |
| gk=g, |
| beta=beta, |
| scale=scale, |
| initial_state=initial_state, |
| output_final_state=output_final_state, |
| use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, |
| cu_seqlens=cu_seqlens, |
| ) |
| return o, final_state |
|
|
