feat: Wx/modify ops for npu

deeplink_ext/internevo_ops/_flash_attention_npu.py

@@ -25,24 +25,25 @@ def flash_attn_func(
     deterministic=False,
     return_attn_probs=False,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
 
     seqlen_q = q.shape[1]
     seqlen_k = k.shape[1]
     head_num = q.shape[-2]
 
-    if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
-        sparse_mode = 0
-    else:
-        sparse_mode = 2
-
-    seqlen_q = min(seqlen_q, 2048)
-    seqlen_k = min(seqlen_k, 2048)
+    assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
+    sparse_mode = 2 if seqlen_q > 2048 else 0
+    seqlen = min(seqlen_q, 2048)
 
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
+            torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
@@ -81,25 +82,28 @@ def flash_attn_varlen_func(
     alibi_slopes=None,
     deterministic=False,
     return_attn_probs=False,
-    block_table=None,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
     head_num = q.shape[-2]
 
     cu_seqlens_q = cu_seqlens_q[1:].tolist()
     cu_seqlens_k = cu_seqlens_k[1:].tolist()
-    seqlen_q = min(max_seqlen_q, 2048)
-    seqlen_k = min(max_seqlen_k, 2048)
 
-    if max_seqlen_q < 2048:
-        sparse_mode = 0
-    else:
-        sparse_mode = 2
+    assert (
+        max_seqlen_q == max_seqlen_k
+    ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
+    sparse_mode = 2 if max_seqlen_q > 2048 else 0
+    max_seqlen = min(max_seqlen_q, 2048)
 
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
+            torch.ones([max_seqlen, max_seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
@@ -114,8 +118,8 @@ def flash_attn_varlen_func(
         "TND",
         atten_mask=attention_mask,
         scale=softmax_scale,
-        pre_tockens=q.shape[0],  # seq_len
-        next_tockens=0,  # 0
+        pre_tockens=q.shape[0],
+        next_tockens=0,
         keep_prob=1 - dropout_p,
         sparse_mode=sparse_mode,
         actual_seq_qlen=cu_seqlens_q,
@@ -134,6 +138,11 @@ def flash_attn_qkvpacked_func(
     deterministic=False,
     return_attn_probs=False,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = qkv.shape[-1] ** (-0.5)
     q = qkv[:, :, 0]
@@ -143,16 +152,12 @@ def flash_attn_qkvpacked_func(
     seqlen_qkv = qkv.shape[1]
     head_num = q.shape[-2]
 
-    if seqlen_qkv < 2048:
-        sparse_mode = 0
-    else:
-        sparse_mode = 2
-
-    seqlen_qkv = min(qkv.shape[1], 2048)
+    sparse_mode = 2 if seqlen_qkv > 2048 else 0
+    seqlen = min(seqlen_qkv, 2048)
 
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
+            torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
@@ -187,26 +192,27 @@ def flash_attn_kvpacked_func(
     deterministic=False,
     return_attn_probs=False,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
     k = kv[:, :, 0]
     v = kv[:, :, 1]
 
-    s0 = q.shape[1]
-    s1 = kv.shape[1]
+    seqlen_q = q.shape[1]
+    seqlen_kv = kv.shape[1]
     head_num = q.shape[-2]
 
-    if s0 == s1 and s0 < 2048 and s1 < 2048:
-        sparse_mode = 0
-    else:
-        sparse_mode = 2
-
-    seqlen_q = min(s0, 2048)
-    seqlen_k = min(s1, 2048)
+    assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
+    sparse_mode = 2 if seqlen_q > 2048 else 0
+    seqlen = min(seqlen_q, 2048)
 
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
+            torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
@@ -222,7 +228,7 @@ def flash_attn_kvpacked_func(
         atten_mask=attention_mask,
         scale=softmax_scale,
         keep_prob=1 - dropout_p,
-        pre_tockens=seqlen_k,
+        pre_tockens=seqlen_q,
         next_tockens=0,
         sparse_mode=sparse_mode,
     )[0]
@@ -242,37 +248,42 @@ def flash_attn_varlen_qkvpacked_func(
     deterministic=False,
     return_attn_probs=False,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = qkv.shape[-1] ** (-0.5)
     q = qkv[:, 0]
     k = qkv[:, 1]
     v = qkv[:, 2]
-    n = q.shape[1]
-    if max_seqlen > 2048:
-        sparse_mode = 2
-    else:
-        sparse_mode = 0
+    head_num = q.shape[1]
+
     cu_seqlens_q = cu_seqlens[1:].tolist()
     cu_seqlens_k = cu_seqlens[1:].tolist()
-    seqlen = min(max_seqlen, 2048)
+
+    sparse_mode = 2 if max_seqlen > 2048 else 0
+    max_seqlen = min(max_seqlen, 2048)
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
+            torch.ones([max_seqlen, max_seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
         else None
     )
+
     out = torch_npu.npu_fusion_attention(
         q,
         k,
         v,
-        n,
+        head_num,
         "TND",
         atten_mask=attention_mask,
         scale=softmax_scale,
-        pre_tockens=q.shape[0],  # seq_len
-        next_tockens=0,  # 0
+        pre_tockens=q.shape[0],
+        next_tockens=0,
         keep_prob=1 - dropout_p,
         sparse_mode=sparse_mode,
         actual_seq_qlen=cu_seqlens_q,
@@ -296,39 +307,44 @@ def flash_attn_varlen_kvpacked_func(
     deterministic=False,
     return_attn_probs=False,
 ):
+    assert window_size == (
+        -1,
+        -1,
+    ), "Npu currently does not support sliding window attention"
+    assert alibi_slopes is None, "Npu currently does not support ALiBi."
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
     k = kv[:, 0]
     v = kv[:, 1]
-    n = q.shape[1]
+    head_num = q.shape[1]
     cu_seqlens_q = cu_seqlens_q[1:].tolist()
     cu_seqlens_k = cu_seqlens_k[1:].tolist()
-    seqlen_q = min(max_seqlen_q, 2048)
-    seqlen_k = min(max_seqlen_k, 2048)
 
-    if max_seqlen_q > 2048:
-        sparse_mode = 2
-    else:
-        sparse_mode = 0
+    assert (
+        max_seqlen_q == max_seqlen_k
+    ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
+    sparse_mode = 2 if max_seqlen_q > 2048 else 0
+    max_seqlen = min(max_seqlen_q, 2048)
 
     attention_mask = (
         torch.triu(
-            torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
+            torch.ones([max_seqlen, max_seqlen], dtype=torch.bool, device=q.device),
             diagonal=1,
         )
         if causal
         else None
     )
+
     out = torch_npu.npu_fusion_attention(
         q,
         k,
         v,
-        n,
+        head_num,
         "TND",
         atten_mask=attention_mask,
         scale=softmax_scale,
-        pre_tockens=q.shape[0],  # seq_len
-        next_tockens=0,  # 0
+        pre_tockens=q.shape[0],
+        next_tockens=0,
         keep_prob=1 - dropout_p,
         sparse_mode=sparse_mode,
         actual_seq_qlen=cu_seqlens_q,

deeplink_ext/internevo_ops/_rotary_embedding_npu.py

@@ -1,8 +1,8 @@
 # Copyright (c) 2024, DeepLink.
 
 import torch
-import torch_npu
-from einops import rearrange
+from einops import repeat
+from mindspeed.ops.npu_rotary_position_embedding import npu_rotary_position_embedding
 
 __all__ = ["ApplyRotaryEmb"]
 
@@ -38,38 +38,73 @@ def forward(
         assert seqlen <= rotary_seqlen
         assert sin.shape == (rotary_seqlen, rotary_dim // 2)
 
-        re_cos = rearrange(cos[:seqlen], "s d -> s 1 d")
-        re_sin = rearrange(sin[:seqlen], "s d -> s 1 d")
-
-        cat_cos = torch.cat([re_cos, re_cos], -1)
-        cat_sin = torch.cat([re_sin, re_sin], -1)
+        if interleaved:
+            cos = repeat(cos[:seqlen], "... d -> 1 ... 1 (d 2)")
+            sin = repeat(sin[:seqlen], "... d -> 1 ... 1 (d 2)")
+        else:
+            cos = repeat(cos[:seqlen], "... d -> 1 ... 1 (2 d)")
+            sin = repeat(sin[:seqlen], "... d -> 1 ... 1 (2 d)")
 
-        rot = torch_npu.npu_rotary_mul(x[..., :rotary_dim], cat_cos, cat_sin)
-        ctx.save_for_backward(cat_cos, cat_sin)
+        ctx.save_for_backward(cos, sin)
         ctx.interleaved = interleaved
         ctx.in_place = in_place
-        if in_place:
-            x[..., :rotary_dim].copy_(rot)
-            return x
+
+        if interleaved:
+            x_ro = x[..., :rotary_dim]
+            out_ro = npu_rotary_position_embedding(x_ro, cos, sin, 1)
+            if in_place:
+                x[..., :rotary_dim].copy_(out_ro)
+                return x
+            if rotary_dim < head_dim:
+                out = torch.empty_like(x)
+                out[..., :rotary_dim].copy_(out_ro)
+                out[..., rotary_dim:].copy_(x[..., rotary_dim:])
+                return out
+            return out_ro
         else:
-            out = x.detach().clone()
-            if rotary_dim < head_dim and not in_place:
+            x_ro = x[..., :rotary_dim]
+            out_ro = npu_rotary_position_embedding(x_ro, cos, sin, 0)
+            if in_place:
+                x[..., :rotary_dim].copy_(out_ro)
+                return x
+            if rotary_dim < head_dim:
+                out = torch.empty_like(x)
+                out[..., :rotary_dim].copy_(out_ro)
                 out[..., rotary_dim:].copy_(x[..., rotary_dim:])
-            return out
+                return out
+            return out_ro
 
     @staticmethod
-    def backward(ctx, do):
-        cat_cos, cat_sin = ctx.saved_tensors
-        *_, seqlen, _, head_dim = do.shape
-        rotary_dim = cat_cos.shape[-1]
+    def backward(ctx, grad_out):
+        cos, sin = ctx.saved_tensors
+        rotary_dim = cos.shape[-1]
+        head_dim = grad_out.shape[-1]
 
-        dx_out = torch_npu.npu_rotary_mul(
-            do[..., :rotary_dim], cat_cos, torch.neg(cat_sin)
-        )
-        if ctx.in_place:
-            do[..., :rotary_dim].copy_(dx_out)
-            return do, None, None, None, None
+        if ctx.interleaved:
+            grad_out_ro = grad_out[..., :rotary_dim]
+            grad_input_ro = npu_rotary_position_embedding(
+                grad_out_ro, cos, torch.neg(sin), 1
+            )
+            if ctx.in_place:
+                grad_out[..., :rotary_dim].copy_(grad_input_ro)
+                return grad_out, None, None, None, None
+            if rotary_dim < head_dim:
+                grad_input = torch.empty_like(grad_out)
+                grad_input[..., :rotary_dim].copy_(grad_input_ro)
+                grad_input[..., rotary_dim:].copy_(grad_out[..., rotary_dim:])
+                return grad_input, None, None, None, None
+            return grad_input_ro, None, None, None, None
         else:
-            dx = do.detach().clone()
-            dx[..., :rotary_dim].copy_(dx_out)
-            return dx, None, None, None, None
+            grad_out_ro = grad_out[..., :rotary_dim]
+            grad_input_ro = npu_rotary_position_embedding(
+                grad_out_ro, cos, torch.neg(sin), 0
+            )
+            if ctx.in_place:
+                grad_out[..., :rotary_dim].copy_(grad_input_ro)
+                return grad_out, None, None, None, None
+            if rotary_dim < head_dim:
+                grad_input = torch.empty_like(grad_out)
+                grad_input[..., :rotary_dim].copy_(grad_input_ro)
+                grad_input[..., rotary_dim:].copy_(grad_out[..., rotary_dim:])
+                return grad_input, None, None, None, None
+            return grad_input_ro, None, None, None, None

deeplink_ext/internevo_ops/rotary_embedding.py

@@ -4,8 +4,7 @@
 
 platform_type = deeplink_ext_get_platform_type()
 if platform_type == PlatformType.TORCH_NPU:
-    # from ._rotary_embedding_npu import ApplyRotaryEmb
-    from .rotary_embedding_fallback import ApplyRotaryEmbTorch as ApplyRotaryEmb
+    from ._rotary_embedding_npu import ApplyRotaryEmb
 elif platform_type == PlatformType.TORCH_DIPU:
     from ._rotary_embedding_dipu import ApplyRotaryEmb
 else: