refactor and remove pos_id in apply_rotary_emb

Author: wanfengcxz

lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py

@@ -5,7 +5,6 @@
 from lmdeploy.pytorch.kernels.dlinfer import apply_rotary_pos_emb
 
 from ..apply_rotary_emb import ApplyRotaryEmbBuilder, ApplyRotaryEmbImpl
-from .attention import DlinferAttentionMetadata
 
 class DlinferApplyRotaryEmbImpl(ApplyRotaryEmbImpl):
     """Apply rotary embedding implementation."""
@@ -15,19 +14,16 @@ def forward(self,
                 key: Tensor,
                 cos: Tensor,
                 sin: Tensor,
-                attn_metadata: DlinferAttentionMetadata,
+                cu_seqlens: Tensor,
                 inplace: bool = True):
         """forward."""
-        cos_sin_ids = attn_metadata.cos_sin_ids
-        cu_seqlens = attn_metadata.cu_seqlens
-
         if inplace:
             q_embed = None
             k_embed = None
         else:
             q_embed = torch.empty_like(query)
             k_embed = torch.empty_like(key)
-        return apply_rotary_pos_emb(query, key, cos, sin, q_embed, k_embed, cos_sin_ids, cu_seqlens)
+        return apply_rotary_pos_emb(query, key, cos, sin, q_embed, k_embed, cu_seqlens)
 
 
 class DlinferApplyRotaryEmbBuilder(ApplyRotaryEmbBuilder):

lmdeploy/pytorch/backends/dlinfer/attention.py

@@ -10,14 +10,13 @@
 @dataclass
 class DlinferAttentionMetadata(AttentionMetadata):
     kv_start_indices: Optional[Tensor] = None
-    block_size: int = 16
+    block_size: int = 64
     attention_mask: Sequence[Tensor] = tuple()
     is_unpaged_prefill: Optional[bool] = None
     max_q_seq_len: int = 1
     max_kv_seq_len: int = 1
     cu_seqlens: Optional[Tensor] = None
-    cos_sin_ids: Optional[Tensor] = None
-
+    is_flash_attn_support_inplace: bool = True
 
 class DlinferAttentionImpl(AttentionImpl[DlinferAttentionMetadata]):
     """dlinfer attention implementation."""
@@ -82,6 +81,10 @@ def forward(
         k_cache, v_cache = self.fill_kv_cache(key, value, k_cache, v_cache,
                                               kv_start_indices)
 
+        if is_unpaged_prefill:
+            inplace = inplace if attn_metadata.is_flash_attn_support_inplace \
+                    else False
+
         if inplace:
             attn_output = query[..., :self.v_head_size]
         else:

lmdeploy/pytorch/backends/dlinfer/camb/op_backend.py

@@ -62,12 +62,6 @@ def update_step_context(cls, step_context):
         cu_seqlens = torch.zeros(batch_size+1, dtype=torch.int32, device=device)
         cu_seqlens[:-1] = step_context.q_start_loc
         cu_seqlens[-1] = step_context.q_seqlens.sum()
-        cu_seqlens_list = cu_seqlens.tolist()
-
-        if not step_context.is_decoding:
-            cos_sin_ids = step_context.position_ids[0].to(torch.int32)
-        else:
-            cos_sin_ids = torch.zeros(batch_size, dtype=torch.int32, device=device)
 
         if not step_context.is_decoding:
             is_unpaged_prefill = \
@@ -104,7 +98,7 @@ def update_step_context(cls, step_context):
             max_q_seq_len=max_q_seq_len,
             max_kv_seq_len=max_kv_seq_len,
             cu_seqlens=cu_seqlens,
-            cos_sin_ids=cos_sin_ids,
+            is_flash_attn_support_inplace=False,
         )
 
         step_context.attn_metadata = attn_metadata

lmdeploy/pytorch/kernels/dlinfer/apply_rotary_pos_emb.py

@@ -10,12 +10,11 @@ def apply_rotary_pos_emb(
     sin: Tensor,
     q_embed: Tensor = None,
     k_embed: Tensor = None,
-    cos_sin_ids=None,
     cu_seqlens=None,
 ):
     query_states = query_states.contiguous()
     key_states = key_states.contiguous()
-    query_states, key_states = ext_ops.apply_rotary_pos_emb(query_states, key_states, cos, sin, None, cos_sin_ids, cu_seqlens)
+    query_states, key_states = ext_ops.apply_rotary_pos_emb(query_states, key_states, cos, sin, None, cu_seqlens)
 
     if q_embed is None:
         q_embed = query_states

lmdeploy/pytorch/kernels/dlinfer/pagedattention.py

@@ -25,8 +25,7 @@ def prefill_attention(
     num_kv_heads = value_states.shape[1]
 
     if is_unpaged_prefill:
-        output = torch.empty_like(query_states)
-        ext_ops.prefill_attention(
+        return ext_ops.prefill_attention(
             query_states,
             key_states,
             value_states,
@@ -36,10 +35,8 @@ def prefill_attention(
             num_q_heads,
             num_kv_heads,
             attn_mask,
-            attn_output=output,
+            attn_output=attn_output,
         )
-        attn_output.copy_(output)
-        return attn_output
     else:
         return ext_ops.paged_prefill_attention(
             query_states,

lmdeploy/pytorch/models/internlm2.py

@@ -75,14 +75,15 @@ def forward(
         query_states, key_states, value_states = self.wqkv.split_qkv(
             qkv_states)
 
+        cu_seqlens = attn_metadata.cu_seqlens
         # apply rotary embedding
         cos, sin = rotary_pos_emb
         query_states, key_states = self.apply_rotary_pos_emb(
             query_states,
             key_states,
             cos,
             sin,
-            attn_metadata,
+            cu_seqlens,
             inplace=True,
         )
 

lmdeploy/pytorch/nn/rotary_embedding.py

@@ -2,7 +2,6 @@
 from torch import Tensor, nn
 
 from ..backends import OpType, get_backend
-from ..backends.attention import AttentionMetadata
 from ..backends.rotary_embedding import (Llama3Parameters,
                                          LongRoPEScalingParameters, RopeType,
                                          YarnParameters)
@@ -44,7 +43,7 @@ def forward(self,
                 key: Tensor,
                 cos: Tensor,
                 sin: Tensor,
-                attn_metadata: AttentionMetadata,
+                cu_seqlens: Tensor,
                 inplace: bool = True):
         """forward."""
-        return self.impl.forward(query, key, cos, sin, attn_metadata, inplace)
+        return self.impl.forward(query, key, cos, sin, cu_seqlens, inplace)