mixstral fused qkv and rope (#9724)

* mixstral fused qkv and rope

* fix and clean

* fix style

* update

* update

* fix

* update

* fix

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

@@ -39,7 +39,7 @@ import math
 import torch.nn.functional as F
 from bigdl.llm.utils.common import invalidInputError
 from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
-from bigdl.llm.transformers.models.utils import rotate_half, apply_rotary_pos_emb
+from bigdl.llm.transformers.models.utils import is_enough_kv_cache_room_4_31, apply_rotary_pos_emb
 from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu
 from bigdl.llm.transformers.low_bit_linear import SYM_INT4
 from bigdl.llm.ggml.quantize import ggml_tensor_qtype
@@ -111,11 +111,6 @@ def llama_mlp_forward(
     return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 
-def is_enough_kv_cache_room(past_key_value):
-    return past_key_value is not None and \
-        past_key_value[0].stride()[1] > past_key_value[0].size(2) * past_key_value[0].size(3)
-
-
 def should_use_fuse_rope(self, query_states, position_ids):
     use_fuse_rope = query_states.device.type == "xpu"
     use_fuse_rope = use_fuse_rope and not (self.training and query_states.requires_grad)
@@ -149,7 +144,7 @@ def llama_attention_forward_4_31(
         attention_dtype = original_dtype
 
     use_fuse_rope = should_use_fuse_rope(self, hidden_states, position_ids)
-    enough_kv_room = is_enough_kv_cache_room(past_key_value)
+    enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value)
     is_q4_0 = self.q_proj.qtype == SYM_INT4
     no_tp = not self.config.pretraining_tp > 1
     decoding_fast_path = (no_tp and is_q4_0 and use_fuse_rope and

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

@@ -44,7 +44,7 @@ from bigdl.llm.utils.common import invalidInputError
 from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
 from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb,\
     apply_rotary_pos_emb_no_cache_xpu
-from bigdl.llm.transformers.models.llama import is_enough_kv_cache_room
+from bigdl.llm.transformers.models.utils import is_enough_kv_cache_room_4_31
 from bigdl.llm.transformers.low_bit_linear import SYM_INT4
 
 KV_CACHE_ALLOC_BLOCK_LENGTH = 256
@@ -89,7 +89,7 @@ def mistral_attention_forward(
     device = hidden_states.device
 
     use_fuse_rope = should_use_fuse_rope(self, hidden_states, position_ids)
-    enough_kv_room = is_enough_kv_cache_room(past_key_value)
+    enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value)
     decoding_fast_path = use_decoding_fast_path(self.q_proj.qtype,
                                                 use_fuse_rope,
                                                 enough_kv_room,

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

@@ -47,7 +47,8 @@ from bigdl.llm.ggml.quantize import ggml_tensor_qtype
 from bigdl.llm.utils.common import invalidInputError
 from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
 from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb,\
-    apply_rotary_pos_emb_no_cache_xpu
+    apply_rotary_pos_emb_no_cache_xpu, is_enough_kv_cache_room_4_36
+from bigdl.llm.transformers.models.mistral import should_use_fuse_rope, use_decoding_fast_path
 
 
 KV_CACHE_ALLOC_BLOCK_LENGTH = 256
@@ -142,6 +143,36 @@ def mixtral_attention_forward(
     bsz, q_len, _ = hidden_states.size()
     device = hidden_states.device
 
+    use_fuse_rope = should_use_fuse_rope(self, hidden_states, position_ids)
+    enough_kv_room = is_enough_kv_cache_room_4_36(past_key_value, self.layer_idx)
+    decoding_fast_path = use_decoding_fast_path(self.q_proj.qtype,
+                                                use_fuse_rope,
+                                                enough_kv_room,
+                                                bsz * q_len)
+
+    if decoding_fast_path:
+        hidden_states = hidden_states.view(1, -1)
+        cache_k = past_key_value.key_cache[self.layer_idx]
+        cache_v = past_key_value.value_cache[self.layer_idx]
+        kv_seq_len = cache_k.shape[-2]
+        import linear_q4_0
+        query_states, key_states, value_states = linear_q4_0.forward_qkv(hidden_states,
+                                                                         self.q_proj.weight,
+                                                                         self.k_proj.weight,
+                                                                         self.v_proj.weight,
+                                                                         position_ids,
+                                                                         cache_k, cache_v,
+                                                                         self.q_proj.weight.qtype,
+                                                                         kv_seq_len,
+                                                                         self.head_dim)
+        kv_seq_len += 1
+        # update past_key_value's seem_tokens and kv caches.
+        if self.layer_idx == 0:
+            past_key_value.seen_tokens = kv_seq_len
+        past_key_value.key_cache[self.layer_idx] = key_states
+        past_key_value.value_cache[self.layer_idx] = value_states
+
+    else:
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
@@ -155,13 +186,14 @@ def mixtral_attention_forward(
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
             if self.layer_idx is None:
-            invalidInputError(False, "The cache structure has changed since version v4.36. "
+                invalidInputError(False,
+                                  "The cache structure has changed since version v4.36. "
                                   f"If you are using {self.__class__.__name__} for "
                                   "auto-regressive decodingwith k/v caching, please make sure "
                                   "to initialize the attention class with a layer index.")
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
 
-    if query_states.device.type == "xpu" and not (self.training and query_states.requires_grad):
+        if use_fuse_rope:
             query_states, key_states = apply_rotary_pos_emb_no_cache_xpu(query_states,
                                                                          key_states,
                                                                          position_ids,
@@ -185,9 +217,9 @@ def mixtral_attention_forward(
                 cache_k = past_key_value.key_cache[self.layer_idx]
                 cache_v = past_key_value.value_cache[self.layer_idx]
 
-            if cache_k.stride()[1] <= cache_k.size(2) * cache_k.size(3):
+                if not enough_kv_room:
                     # allocate new
-                new_cache_k, new_cache_v = extend_kv_cache(bsz,
+                    new_c_k, new_c_v = extend_kv_cache(bsz,
                                                        self.num_key_value_heads,  # Support GQA
                                                        self.head_dim,
                                                        cache_k.size(2),
@@ -195,12 +227,15 @@ def mixtral_attention_forward(
                                                        dtype=cache_k.dtype,
                                                        device=device)
 
-                new_cache_k[:] = cache_k
+                    new_c_k[:] = cache_k
-                new_cache_v[:] = cache_v
+                    new_c_v[:] = cache_v
-                cache_k = new_cache_k
+                    cache_k = new_c_k
-                cache_v = new_cache_v
+                    cache_v = new_c_v
 
-            key_states, value_states = append_kv_cache(cache_k, cache_v, key_states, value_states)
+                key_states, value_states = append_kv_cache(cache_k,
+                                                           cache_v,
+                                                           key_states,
+                                                           value_states)
 
                 # update past_key_value
                 past_key_value.key_cache[self.layer_idx] = key_states

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

@@ -106,3 +106,16 @@ def apply_rotary_pos_emb_no_cache_xpu(q, k, position_ids, model_family):
     else:
         invalidInputError(False,
                           f"{model_family} is not supported.")
+
+
+def is_enough_kv_cache_room_4_36(past_key_value, idx):
+    # to determinate if is enough kv cache room in transformers==4.36
+    return past_key_value is not None and len(past_key_value.key_cache) > idx and \
+        past_key_value.key_cache[idx].stride()[1] > past_key_value.key_cache[idx].size(2) * \
+        past_key_value.key_cache[idx].size(3)
+
+
+def is_enough_kv_cache_room_4_31(past_key_value):
+    # to determinate if is enough kv cache room in transformers between 4.31 and 4.35
+    return past_key_value is not None and \
+        past_key_value[0].stride()[1] > past_key_value[0].size(2) * past_key_value[0].size(3)