ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

use fused qkv forward in qwen2 (#10185)

Browse source 
* use fused qkv forward in qwen2

* support both

* fix style

* fix rope

* remove pring

* fix style

* clean up
This commit is contained in:
Yang Wang 2024-03-01 00:46:35 -08:00 committed by GitHub
parent 509e206de0
commit f4d7dbcde2
3 changed files with 94 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
python/llm/src/bigdl/llm/transformers/models/llama.py
View file

@ -323,7 +323,8 @@ def llama_attention_forward_4_31_quantized(
self.q_proj.weight.qtype, self.q_proj.weight.qtype,
self.v_proj.weight.qtype, self.v_proj.weight.qtype,
0, 0,
self.head_dim) self.head_dim,
self.rotary_emb.base,)
else: else:
query_states = self.q_proj(hidden_states) query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states) key_states = self.k_proj(hidden_states)
@ -511,7 +512,8 @@ def llama_attention_forward_4_31_original(
self.q_proj.weight.qtype, self.q_proj.weight.qtype,
self.v_proj.weight.qtype, self.v_proj.weight.qtype,
kv_seq_len, kv_seq_len,
self.head_dim) self.head_dim,
self.rotary_emb.base,)
kv_seq_len += 1 kv_seq_len += 1
else: else:
@ -762,7 +764,9 @@ def llama_attention_selective_batching_forward_4_31(
self.q_proj.weight.qtype, self.q_proj.weight.qtype,
self.v_proj.weight.qtype, self.v_proj.weight.qtype,
kv_seq_len, kv_seq_len,
self.head_dim) self.head_dim,
self.rotary_emb.base,
)
kv_seq_len += 1 kv_seq_len += 1
else: else:
if self.config.pretraining_tp > 1: if self.config.pretraining_tp > 1:
@ -942,7 +946,8 @@ def llama_attention_forward_4_36(
self.q_proj.weight.qtype, self.q_proj.weight.qtype,
self.v_proj.weight.qtype, self.v_proj.weight.qtype,
kv_seq_len, kv_seq_len,
self.head_dim) self.head_dim,
self.rotary_emb.base,)
kv_seq_len += 1 kv_seq_len += 1
# update past_key_value's seem_tokens and kv caches. # update past_key_value's seem_tokens and kv caches.
if self.layer_idx == 0: if self.layer_idx == 0:

3
python/llm/src/bigdl/llm/transformers/models/mixtral.py
View file

@ -171,7 +171,8 @@ def mixtral_attention_forward(
self.q_proj.weight.qtype, self.q_proj.weight.qtype,
self.v_proj.weight.qtype, self.v_proj.weight.qtype,
kv_seq_len, kv_seq_len,
self.head_dim) self.head_dim,
self.rotary_emb.base,)
kv_seq_len += 1 kv_seq_len += 1
# update past_key_value's seem_tokens and kv caches. # update past_key_value's seem_tokens and kv caches.
if self.layer_idx == 0: if self.layer_idx == 0:

24
python/llm/src/bigdl/llm/transformers/models/qwen2.py
View file

@ -223,6 +223,9 @@ def qwen2_attention_forward_quantized(
attn_weights = None attn_weights = None
return attn_output, attn_weights, past_key_value return attn_output, attn_weights, past_key_value
from bigdl.llm.ggml.quantize import ggml_tensor_qtype
SYM_INT4 = ggml_tensor_qtype["sym_int4"]
FP8E5 = ggml_tensor_qtype["fp8_e5m2"]
def qwen2_attention_forward_origin( def qwen2_attention_forward_origin(
@ -247,6 +250,27 @@ def qwen2_attention_forward_origin(
device = hidden_states.device device = hidden_states.device
enough_kv_room = is_enough_kv_cache_room_4_36(past_key_value, self.layer_idx) enough_kv_room = is_enough_kv_cache_room_4_36(past_key_value, self.layer_idx)
qtype = getattr(self.q_proj, "qtype", None)
qtype_check = qtype in [SYM_INT4, FP8E5]
decoding_fast_path = (qtype_check and use_fuse_rope
and enough_kv_room and bsz * q_len == 1)
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
args = [hidden_states, self.q_proj.weight, self.k_proj.weight, self.v_proj.weight,
self.q_proj.bias, self.k_proj.bias, self.v_proj.bias, position_ids, cache_k,
cache_v, self.q_proj.weight.qtype, kv_seq_len, self.head_dim, self.rotary_emb.base]
query_states, key_states, value_states = linear_q4_0.forward_qkv_bias(*args)
kv_seq_len += 1
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) query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states) key_states = self.k_proj(hidden_states)