[LLM] Enable kv_cache and forward_qkv optimizations for yuan2 (#10225)

* add init kv_cache support for yuan2

* add forward qkv in yuan

python/llm/src/bigdl/llm/transformers/models/yuan.py

@@ -20,13 +20,32 @@
 # https://huggingface.co/IEITYuan/Yuan2-2B-hf/blob/7ab7b3c18eb8e5232ce2a3f720d4e6f4b53a2806/README.md#%E5%A3%B0%E6%98%8E%E4%B8%8E%E5%8D%8F%E8%AE%AEterms-and-conditions
 #
 
-import torch
 import copy
 import math
+from einops import rearrange
 from typing import Optional, Tuple
 
-from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb
+import torch
+import torch.nn as nn
 from bigdl.llm.utils.common import invalidInputError
+from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb
+from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
+from bigdl.llm.transformers.models.utils import is_enough_kv_cache_room_4_31
+from bigdl.llm.transformers.low_bit_linear import SYM_INT4, FP8E5
+
+KV_CACHE_ALLOC_BLOCK_LENGTH = 256
+
+
+def use_decoding_fast_path(q_type, use_fuse_rope, enough_kv_room, bs):
+    return q_type in [SYM_INT4, FP8E5] and \
+        use_fuse_rope and enough_kv_room and bs == 1
+
+
+def should_use_fuse_rope(self, hidden_states, position_ids):
+    use_fuse_rope = hidden_states.device.type == "xpu"
+    use_fuse_rope = use_fuse_rope and not (self.training and hidden_states.requires_grad)
+    use_fuse_rope = use_fuse_rope and position_ids is not None
+    return use_fuse_rope
 
 
 def yuan_attention_forward(
@@ -39,8 +58,13 @@ def yuan_attention_forward(
     use_cache: bool = False,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     bsz, q_len, _ = hidden_states.size()
+    device = hidden_states.device
     before_hidden_states = None
     is_first_step = False
+    self.use_shareqk = False
+
+    enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value)
+
     if use_cache:
         if past_key_value is None:
             inference_hidden_states_memory = torch.empty(bsz, 2,
@@ -64,7 +88,9 @@ def yuan_attention_forward(
 
     value_states = \
         self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
     if self.use_shareqk:
+        # use_shareqk is disabled for now
         qk_states = self.qk_proj(hidden_states).view(bsz, q_len, self.num_heads*self.head_dim)
         query_key = qk_states.unsqueeze(2) * self.qk_weight + self.qk_bias
         query_states, key_states = torch.unbind(query_key, dim=2)
@@ -95,80 +121,71 @@ def yuan_attention_forward(
 
     if past_key_value is not None:
         # reuse k, v, self_attention
-        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        cache_k = past_key_value[0]
-        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        cache_v = past_key_value[1]
+        if not enough_kv_room:
+            # allocate new
+            new_cache_k, new_cache_v = extend_kv_cache(bsz,
+                                                       self.num_heads,
+                                                       self.head_dim,
+                                                       cache_k.size(2),
+                                                       kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH,
+                                                       dtype=cache_k.d_type,
+                                                       device=device)
+            new_cache_k[:] = cache_k
+            new_cache_v[:] = cache_v
+            cache_k = new_cache_k
+            cache_v = new_cache_v
+
+        key_states, value_states = append_kv_cache(cache_k, cache_v, key_states, value_states)
+
+    elif use_cache:
+        max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH
+        new_key_states, new_value_states = init_kv_cache(bsz,
+                                                         self.num_heads,
+                                                         self.head_dim,
+                                                         kv_seq_len,
+                                                         max_cache_length,
+                                                         dtype=key_states.dtype,
+                                                         device=device)
+        new_key_states[:] = key_states
+        new_value_states[:] = value_states
+        key_states = new_key_states
+        value_states = new_value_states
 
     past_key_value = \
         (key_states, value_states, inference_hidden_states_memory) if use_cache else None
 
-    if self.use_flash_attention:
+    attn_weights = \
-        attn_weights = None
+        torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
 
-        batch_size, seqlen_q = query_states.shape[0], query_states.shape[1]
+    invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len),
-        seqlen_k = key_states.shape[1]
+                      "Attention weights should be of size "
+                      f"{(bsz, self.num_heads, q_len, kv_seq_len)}, "
+                      f"but is {attn_weights.size()}")
 
-        q, k, v = \
+    if attention_mask is not None:
-            [rearrange(x, 'b s ... -> (b s) ...') for x in [query_states, key_states, value_states]]
+        invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len),
+                          f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
+                          f"but is {attention_mask.size()}")
+        attn_weights = attn_weights + attention_mask
+        attn_weights = torch.max(attn_weights,
+                                 torch.tensor(torch.finfo(attn_weights.dtype).min))
 
-        cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q,
+    # upcast attention to fp32
-                                    step=seqlen_q,
+    attn_weights = \
-                                    dtype=torch.int,
+        torch.nn.functional.softmax(attn_weights,
-                                    device=q.device)
+                                    dim=-1,
+                                    dtype=torch.float32).to(query_states.dtype)
+    attn_output = torch.matmul(attn_weights, value_states)
 
-        if self.training:
+    invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim),
-            invalidInputError(seqlen_k == seqlen_q,
+                      "`attn_output` should be of size "
-                              "`seqlen_k` should be equal to `seqlen_q`, but is not")
+                      f"{(bsz, self.num_heads, q_len, self.head_dim)}, "
-            cu_seqlens_k = cu_seqlens_q
+                      f"but is {attn_output.size()}")
-            is_causal = self.causal_mask
-        else:
-            is_causal = seqlen_q == seqlen_k
-            cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k,
-                                        step=seqlen_k,
-                                        dtype=torch.int,
-                                        device=q.device)
-            self.dropout = 0
-
-        output = flash_attn_unpadded_func(
-            q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, self.dropout, causal=is_causal
-        )
-
-        attn_output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
-    else:
-        attn_weights = \
-            torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
-        invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len),
-                          "Attention weights should be of size "
-                          f"{(bsz, self.num_heads, q_len, kv_seq_len)}, "
-                          f"but is {attn_weights.size()}")
-
-        if attention_mask is not None:
-            invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len),
-                              f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
-                              f"but is {attention_mask.size()}")
-            attn_weights = attn_weights + attention_mask
-            attn_weights = torch.max(attn_weights,
-                                     torch.tensor(torch.finfo(attn_weights.dtype).min))
-
-        # upcast attention to fp32
-        attn_weights = \
-            torch.nn.functional.softmax(attn_weights,
-                                        dim=-1,
-                                        dtype=torch.float32).to(query_states.dtype)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim),
-                          "`attn_output` should be of size "
-                          f"{(bsz, self.num_heads, q_len, self.head_dim)}, "
-                          f"but is {attn_output.size()}")
-
-        attn_output = attn_output.transpose(1, 2)
 
+    attn_output = attn_output.transpose(1, 2)
     attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
     attn_output = self.o_proj(attn_output)
 
     if not output_attentions: