patch fused_recurrent_gated_delta_rule_fwd_kernel

Signed-off-by: Icey <1790571317@qq.com>

Author: wxsIcey

vllm_ascend/ops/sigmoid_gating.py

@@ -10,10 +10,8 @@
 # mypy: ignore-errors
 
 import os
-from typing import Optional
 
-import torch
-import vllm
+import vllm.model_executor.layers.fla.ops.fused_recurrent
 from vllm.triton_utils import tl, tldevice, triton
 
 if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
@@ -179,227 +177,4 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
         # p_beta += HV * (V if IS_BETA_HEADWISE else 1)
 
 
-def fused_recurrent_gated_delta_rule_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    g: torch.Tensor,
-    beta: torch.Tensor,
-    scale: float,
-    initial_state: torch.Tensor,
-    inplace_final_state: bool = True,
-    cu_seqlens: Optional[torch.LongTensor] = None,
-    ssm_state_indices: Optional[torch.Tensor] = None,
-    num_accepted_tokens: Optional[torch.Tensor] = None,
-    use_qk_l2norm_in_kernel: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    B, T, H, K, V = *k.shape, v.shape[-1]
-    HV = v.shape[2]
-    N = B if cu_seqlens is None else len(cu_seqlens) - 1
-    BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8)
-    NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
-    assert NK == 1, "NK > 1 is not supported yet"
-    num_stages = 3
-    num_warps = 1
-
-    o = q.new_empty(NK, *v.shape)
-    if inplace_final_state:
-        final_state = initial_state
-    else:
-        final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)
-
-    stride_init_state_token = initial_state.stride(0)
-    stride_final_state_token = final_state.stride(0)
-
-    if ssm_state_indices is None:
-        stride_indices_seq, stride_indices_tok = 1, 1
-    elif ssm_state_indices.ndim == 1:
-        stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
-    else:
-        stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
-
-    # print("N: ", N)
-    # print("T: ", T)
-    # print("B: ", B)
-    # print("H: ", H)
-    # print("HV: ", HV)
-    # print("K: ", K)
-    # print("V: ", V)
-    # print("BK: ", BK)
-    # print("BV: ", BV)
-
-    grid = (NK, NV, N * HV)
-    fused_recurrent_gated_delta_rule_fwd_kernel[grid](
-        q=q,
-        k=k,
-        v=v,
-        g=g,
-        beta=beta,
-        o=o,
-        h0=initial_state,
-        ht=final_state,
-        cu_seqlens=cu_seqlens,
-        ssm_state_indices=ssm_state_indices,
-        num_accepted_tokens=num_accepted_tokens,
-        scale=scale,
-        N=N,
-        T=T,
-        B=B,
-        H=H,
-        HV=HV,
-        K=K,
-        V=V,
-        BK=BK,
-        BV=BV,
-        stride_init_state_token=stride_init_state_token,
-        stride_final_state_token=stride_final_state_token,
-        stride_indices_seq=stride_indices_seq,
-        stride_indices_tok=stride_indices_tok,
-        IS_BETA_HEADWISE=beta.ndim == v.ndim,
-        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
-        INPLACE_FINAL_STATE=inplace_final_state,
-        num_warps=num_warps,
-        num_stages=num_stages,
-    )
-    o = o.squeeze(0)
-    return o, final_state
-
-
-class FusedRecurrentFunction(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx,
-                q: torch.Tensor,
-                k: torch.Tensor,
-                v: torch.Tensor,
-                g: torch.Tensor,
-                beta: torch.Tensor,
-                scale: float,
-                initial_state: torch.Tensor,
-                inplace_final_state: bool = True,
-                cu_seqlens: Optional[torch.LongTensor] = None,
-                ssm_state_indices: Optional[torch.Tensor] = None,
-                num_accepted_tokens: Optional[torch.Tensor] = None,
-                use_qk_l2norm_in_kernel: bool = False):
-        o, final_state = fused_recurrent_gated_delta_rule_fwd(
-            q=q.contiguous(),
-            k=k.contiguous(),
-            v=v.contiguous(),
-            g=g.contiguous(),
-            beta=beta.contiguous(),
-            scale=scale,
-            initial_state=initial_state,
-            inplace_final_state=inplace_final_state,
-            cu_seqlens=cu_seqlens,
-            ssm_state_indices=ssm_state_indices,
-            num_accepted_tokens=num_accepted_tokens,
-            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
-        )
-
-        return o, final_state
-
-
-def fused_recurrent_gated_delta_rule(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    g: torch.Tensor,
-    beta: torch.Tensor = None,
-    scale: float = None,
-    initial_state: torch.Tensor = None,
-    inplace_final_state: bool = True,
-    cu_seqlens: Optional[torch.LongTensor] = None,
-    ssm_state_indices: Optional[torch.Tensor] = None,
-    num_accepted_tokens: Optional[torch.Tensor] = None,
-    use_qk_l2norm_in_kernel: bool = False,
-) -> 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[int]):
-            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`.
-        inplace_final_state: bool:
-            Whether to store the final state in-place to save memory.
-            Default: `True`.
-        cu_seqlens (torch.LongTensor):
-            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
-            consistent with the FlashAttention API.
-        ssm_state_indices (Optional[torch.Tensor]):
-            Indices to map the input sequences to the initial/final states.
-        num_accepted_tokens (Optional[torch.Tensor]):
-            Number of accepted tokens for each sequence during decoding.
-    Returns:
-        o (torch.Tensor):
-            Outputs of shape `[B, T, HV, V]`.
-        final_state (torch.Tensor):
-            Final state of shape `[N, HV, K, V]`.
-    Examples::
-        >>> import torch
-        >>> import torch.nn.functional as F
-        >>> from einops import rearrange
-        >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
-        # 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, device='cuda'))
-        >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
-        >>> h0 = torch.randn(B, HV, K, V, device='cuda')
-        >>> o, ht = fused_gated_recurrent_delta_rule(
-            q, k, v, g, beta,
-            initial_state=h0,
-        )
-        # 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_gated_recurrent_delta_rule(
-            q, k, v, g, beta,
-            initial_state=h0,
-            cu_seqlens=cu_seqlens
-        )
-    """
-    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.")
-    if scale is None:
-        scale = k.shape[-1]**-0.5
-    else:
-        assert scale > 0, "scale must be positive"
-    if beta is None:
-        beta = torch.ones_like(q[..., 0])
-    o, final_state = FusedRecurrentFunction.apply(
-        q,
-        k,
-        v,
-        g,
-        beta,
-        scale,
-        initial_state,
-        inplace_final_state,
-        cu_seqlens,
-        ssm_state_indices,
-        num_accepted_tokens,
-        use_qk_l2norm_in_kernel,
-    )
-    return o, final_state
-
-
-vllm.model_executor.layers.fla.ops.fused_recurrent.fused_recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule
+vllm.model_executor.layers.fla.ops.fused_recurrent.fused_recurrent_gated_delta_rule_fwd_kernel = fused_recurrent_gated_delta_rule_fwd_kernel

vllm_ascend/patch/worker/patch_common/__init__.py

@@ -17,6 +17,6 @@
 
 import vllm_ascend.patch.worker.patch_common.patch_distributed  # noqa
 import vllm_ascend.patch.worker.patch_common.patch_logits  # noqa
-
+import vllm_ascend.ops.sigmoid_gating
 # TODO: revert me when triton import is fixed
 # import vllm_ascend.patch.worker.patch_common.patch_minicpm  # noqa