From b6b70d1ba0a1931ca46d5953ecc40548f8cad836 Mon Sep 17 00:00:00 2001 From: Yina Chen <33650826+cyita@users.noreply.github.com> Date: Tue, 28 May 2024 12:00:18 +0800 Subject: [PATCH] Divide core-xe packages (#11131) * temp * add batch * fix style * update package name * fix style * add workflow * use temp version to run uts * trigger performance test * trigger win igpu perf * revert workflow & setup --- .github/actions/llm/setup-llm-env/action.yml | 8 + .github/workflows/llm_performance_tests.yml | 4 +- python/llm/setup.py | 2 + python/llm/src/ipex_llm/transformers/bmm.py | 4 +- .../src/ipex_llm/transformers/embedding.py | 6 +- .../transformers/layers/rope_embedding.py | 16 +- .../ipex_llm/transformers/low_bit_linear.py | 58 +++++-- .../ipex_llm/transformers/models/baichuan.py | 20 ++- .../ipex_llm/transformers/models/baichuan2.py | 38 ++-- .../src/ipex_llm/transformers/models/bloom.py | 12 +- .../ipex_llm/transformers/models/chatglm2.py | 12 +- .../ipex_llm/transformers/models/cohere.py | 33 ++-- .../src/ipex_llm/transformers/models/gemma.py | 30 ++-- .../ipex_llm/transformers/models/internlm.py | 14 +- .../src/ipex_llm/transformers/models/llama.py | 163 +++++++++--------- .../ipex_llm/transformers/models/mistral.py | 123 +++++++------ .../ipex_llm/transformers/models/mixtral.py | 40 ++--- .../src/ipex_llm/transformers/models/phi3.py | 37 ++-- .../src/ipex_llm/transformers/models/qwen.py | 49 +++--- .../src/ipex_llm/transformers/models/qwen2.py | 21 +-- .../ipex_llm/transformers/models/qwen2_moe.py | 20 +-- .../ipex_llm/transformers/models/qwen_vl.py | 4 +- .../src/ipex_llm/transformers/models/rwkv4.py | 12 +- .../src/ipex_llm/transformers/models/rwkv5.py | 12 +- .../ipex_llm/transformers/models/stablelm.py | 13 +- .../transformers/models/starcoder2.py | 5 +- .../src/ipex_llm/transformers/models/utils.py | 32 ++-- .../src/ipex_llm/transformers/models/yuan.py | 12 +- 28 files changed, 427 insertions(+), 373 deletions(-) diff --git a/.github/actions/llm/setup-llm-env/action.yml b/.github/actions/llm/setup-llm-env/action.yml index 87d9ff20..31115de6 100644 --- a/.github/actions/llm/setup-llm-env/action.yml +++ b/.github/actions/llm/setup-llm-env/action.yml @@ -13,12 +13,20 @@ runs: run: | # make sure we install the latest version for bigdl-core-xe related packages pip uninstall bigdl-core-xe -y || true + pip uninstall bigdl-core-xe-batch -y || true + pip uninstall bigdl-core-xe-addons -y || true pip uninstall bigdl-core-xe-esimd -y || true pip uninstall bigdl-core-xe-21 -y || true + pip uninstall bigdl-core-xe-batch-21 -y || true + pip uninstall bigdl-core-xe-addons-21 -y || true pip uninstall bigdl-core-xe-esimd-21 -y || true sed -i 's/"bigdl-core-xe==" + CORE_XE_VERSION + "/"bigdl-core-xe/g' python/llm/setup.py + sed -i 's/"bigdl-core-xe-batch==" + CORE_XE_VERSION + "/"bigdl-core-xe-batch/g' python/llm/setup.py + sed -i 's/"bigdl-core-xe-addons==" + CORE_XE_VERSION + "/"bigdl-core-xe-addons/g' python/llm/setup.py sed -i 's/"bigdl-core-xe-esimd==" + CORE_XE_VERSION + "/"bigdl-core-xe-esimd/g' python/llm/setup.py sed -i 's/"bigdl-core-xe-21==" + CORE_XE_VERSION/"bigdl-core-xe-21"/g' python/llm/setup.py + sed -i 's/"bigdl-core-xe-batch-21==" + CORE_XE_VERSION/"bigdl-core-xe-batch-21"/g' python/llm/setup.py + sed -i 's/"bigdl-core-xe-addons-21==" + CORE_XE_VERSION/"bigdl-core-xe-addons-21"/g' python/llm/setup.py sed -i 's/"bigdl-core-xe-esimd-21==" + CORE_XE_VERSION/"bigdl-core-xe-esimd-21"/g' python/llm/setup.py pip install requests diff --git a/.github/workflows/llm_performance_tests.yml b/.github/workflows/llm_performance_tests.yml index 73098d4d..9f49a997 100644 --- a/.github/workflows/llm_performance_tests.yml +++ b/.github/workflows/llm_performance_tests.yml @@ -312,6 +312,8 @@ jobs: # shell: bash # run: | # sed -i 's/"bigdl-core-xe-21==" + CORE_XE_VERSION/"bigdl-core-xe-21"/g' python/llm/setup.py + # sed -i 's/"bigdl-core-xe-batch-21==" + CORE_XE_VERSION/"bigdl-core-xe-batch-21"/g' python/llm/setup.py + # sed -i 's/"bigdl-core-xe-addons-21==" + CORE_XE_VERSION/"bigdl-core-xe-addons-21"/g' python/llm/setup.py # sed -i 's/"bigdl-core-xe-esimd-21==" + CORE_XE_VERSION/"bigdl-core-xe-esimd-21"/g' python/llm/setup.py # - name: Install ipex-llm and other related packages (install from source) @@ -740,4 +742,4 @@ jobs: # if: ${{ always() }} # shell: cmd # run: | - # call conda env remove -n igpu-perf -y + # call conda env remove -n igpu-perf -y \ No newline at end of file diff --git a/python/llm/setup.py b/python/llm/setup.py index d3b2f2f3..fa5af247 100644 --- a/python/llm/setup.py +++ b/python/llm/setup.py @@ -298,6 +298,8 @@ def setup_package(): "torchvision==0.16.0a0", "intel_extension_for_pytorch==2.1.10+xpu", "bigdl-core-xe-21==" + CORE_XE_VERSION, + "bigdl-core-xe-batch-21==" + CORE_XE_VERSION, + "bigdl-core-xe-addons-21==" + CORE_XE_VERSION, "bigdl-core-xe-esimd-21==" + CORE_XE_VERSION] xpu_21_requires += oneapi_2024_0_requires # default to ipex 2.1 for linux and windows diff --git a/python/llm/src/ipex_llm/transformers/bmm.py b/python/llm/src/ipex_llm/transformers/bmm.py index 56a12eef..dbde3a2b 100644 --- a/python/llm/src/ipex_llm/transformers/bmm.py +++ b/python/llm/src/ipex_llm/transformers/bmm.py @@ -16,7 +16,7 @@ import torch -import linear_q4_0 +import xe_linear torch_bmm_old_ = torch.bmm @@ -30,7 +30,7 @@ def torch_bmm(a, b): if a.size(1) == 1: torch_bmm_old_(a, b, out=C) else: - linear_q4_0.bmm(a.contiguous(), b.contiguous(), C) + xe_linear.bmm(a.contiguous(), b.contiguous(), C) return C diff --git a/python/llm/src/ipex_llm/transformers/embedding.py b/python/llm/src/ipex_llm/transformers/embedding.py index 8031e020..1046c9e2 100644 --- a/python/llm/src/ipex_llm/transformers/embedding.py +++ b/python/llm/src/ipex_llm/transformers/embedding.py @@ -104,11 +104,11 @@ class LowBitEmbedding(torch.nn.Embedding): "`LowBitEmbedding` only supports GPU now.") try: import intel_extension_for_pytorch - import linear_q4_0 + import xe_linear except ModuleNotFoundError: invalidInputError(False, "Please `pip install bigdl_core_xe` first.") - result = linear_q4_0.dequantize_rows(x.contiguous(), self.weight.data, - self.weight.qtype, self.embedding_dim) + result = xe_linear.dequantize_rows(x.contiguous(), self.weight.data, + self.weight.qtype, self.embedding_dim) return result.to(self.torch_dtype) diff --git a/python/llm/src/ipex_llm/transformers/layers/rope_embedding.py b/python/llm/src/ipex_llm/transformers/layers/rope_embedding.py index 0c6c3714..be03c5ab 100644 --- a/python/llm/src/ipex_llm/transformers/layers/rope_embedding.py +++ b/python/llm/src/ipex_llm/transformers/layers/rope_embedding.py @@ -42,26 +42,26 @@ class FastRopeEmbedding(torch.autograd.Function): @staticmethod @custom_fwd def forward(ctx, x, position_ids): - import linear_q4_0 + import xe_addons x_embed = torch.empty(x.shape, dtype=x.dtype, device=x.device) - linear_q4_0.apply_rotary_embedding_half_q_or_k(x, position_ids, - x_embed, False) + xe_addons.apply_rotary_embedding_half_q_or_k(x, position_ids, + x_embed, False) ctx.save_for_backward(position_ids) return x_embed @staticmethod @custom_bwd def backward(ctx, grad_output): - import linear_q4_0 + import xe_addons # LOG.info(f"backward, grad_output: {grad_output}") position_ids, = ctx.saved_tensors dx = torch.empty(grad_output.shape, dtype=grad_output.dtype, device=grad_output.device) - linear_q4_0.apply_rotary_embedding_half_q_or_k(grad_output, - position_ids, - dx, - True) + xe_addons.apply_rotary_embedding_half_q_or_k(grad_output, + position_ids, + dx, + True) # LOG.info(f"backward, dx: {dx}, position_ids: {position_ids}, # requires_grad: {ctx.needs_input_grad}") return dx, None diff --git a/python/llm/src/ipex_llm/transformers/low_bit_linear.py b/python/llm/src/ipex_llm/transformers/low_bit_linear.py index f129093f..41e8ea71 100644 --- a/python/llm/src/ipex_llm/transformers/low_bit_linear.py +++ b/python/llm/src/ipex_llm/transformers/low_bit_linear.py @@ -320,6 +320,26 @@ def reshape_lm_head_input(x): return x +def use_batch_forward(x: torch.Tensor, qtype: int, output_len: int): + device = get_xpu_device_type(x) + batch_size = x.shape[0] + hard_condition = ( + x.dtype in [torch.float, torch.half] + and x.shape[1] % 256 == 0 + and output_len % 32 == 0 + and device in ["arc", "flex", "pvc", "mtl"] + and qtype in [SYM_INT4, ASYM_INT4, SYM_INT8, FP4, + FP8E5, FP6] + and batch_size <= 64 + ) + if hard_condition: + return ( + batch_size > 1 + or (device in ["arc", "flex"] and qtype in [SYM_INT8, FP4]) + ) + return False + + # Rename to FP4Params to trigger initializing # the params layer with all parameters on the CPU # https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L333 @@ -524,8 +544,8 @@ class MatMulLowBit(torch.autograd.Function): @custom_fwd def forward(ctx, A, weight, input_seq_size): ctx.is_empty = False - import linear_q4_0 - result = linear_q4_0.forward_new(A, weight.data, weight.qtype, input_seq_size) + import xe_linear + result = xe_linear.forward_new(A, weight.data, weight.qtype, input_seq_size) if any(ctx.needs_input_grad[:2]): ctx.tensors = (A, weight) else: @@ -535,7 +555,7 @@ class MatMulLowBit(torch.autograd.Function): @staticmethod @custom_bwd def backward(ctx, grad_output): - import linear_q4_0 + import xe_linear if ctx.is_empty: bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias) return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None @@ -545,7 +565,7 @@ class MatMulLowBit(torch.autograd.Function): if req_gradA: if torch.xpu.is_autocast_xpu_enabled(): grad_output = grad_output.to(torch.xpu.get_autocast_xpu_dtype()) - dequant_weight = linear_q4_0.dequant(A, weight.data, weight.qtype) + dequant_weight = xe_linear.dequant(A, weight.data, weight.qtype) grad_A = torch.matmul(grad_output, dequant_weight.reshape(weight._shape)) return grad_A, grad_weight, None @@ -659,7 +679,7 @@ class LowBitLinear(nn.Linear): # GPU logic try: import intel_extension_for_pytorch - import linear_q4_0 + import xe_linear from ipex_llm.transformers.models.utils import use_xmx except ModuleNotFoundError: invalidInputError(False, @@ -678,12 +698,12 @@ class LowBitLinear(nn.Linear): if x_2d.requires_grad: result = MatMulLowBit.apply(x_2d, self.weight, input_seq_size) else: - result = linear_q4_0.forward_new(x_2d, self.weight.data, - self.weight.qtype, - input_seq_size) + result = xe_linear.forward_new(x_2d, self.weight.data, + self.weight.qtype, + input_seq_size) elif self.enable_xetla: x_2d = x_2d.half() - result = linear_q4_0.mm_xetla(x_2d, self.weight.data, self.qtype) + result = xe_linear.mm_xetla(x_2d, self.weight.data, self.qtype) else: # inference path # current workaround to reduce first token latency of fp32 input @@ -696,12 +716,24 @@ class LowBitLinear(nn.Linear): if self.conver_to_half and x_2d.shape[0] > 1 and x_2d.dtype == torch.float32 and \ not use_xmx(x_2d, self.weight.qtype): x_2d = x_2d.half() - result = linear_q4_0.forward_new(x_2d, self.weight.data, self.weight.qtype, - input_seq_size) + if use_batch_forward(x_2d, self.weight.qtype, self.out_len): + import xe_batch + result = xe_batch.batch_forward(x_2d, self.weight.data, + self.weight.qtype, + input_seq_size) + else: + result = xe_linear.forward_new(x_2d, self.weight.data, self.weight.qtype, + input_seq_size) result = result.to(x.dtype) else: - result = linear_q4_0.forward_new(x_2d, self.weight.data, self.weight.qtype, - input_seq_size) + if use_batch_forward(x_2d, self.weight.qtype, self.out_len): + import xe_batch + result = xe_batch.batch_forward(x_2d, self.weight.data, + self.weight.qtype, + input_seq_size) + else: + result = xe_linear.forward_new(x_2d, self.weight.data, self.weight.qtype, + input_seq_size) if do_empty_cache: torch.xpu.empty_cache() result = result.view(new_shape) diff --git a/python/llm/src/ipex_llm/transformers/models/baichuan.py b/python/llm/src/ipex_llm/transformers/models/baichuan.py index 26509b92..8bcdb637 100644 --- a/python/llm/src/ipex_llm/transformers/models/baichuan.py +++ b/python/llm/src/ipex_llm/transformers/models/baichuan.py @@ -168,9 +168,9 @@ def baichuan_attention_forward_7b_quantized( dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) attn_weights = None invalidInputError( @@ -277,8 +277,9 @@ def baichuan_attention_forward_7b_origin( attn_weights = None elif not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: @@ -400,8 +401,8 @@ def baichuan_attention_forward_13b_quantized( query_states.dtype) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) else: - import linear_q4_0 - attn_weights = linear_q4_0.query_key_fp8_matmul(query_states, key_states) + import xe_addons + attn_weights = xe_addons.query_key_fp8_matmul(query_states, key_states) attn_weights = attn_weights / math.sqrt(self.head_dim) @@ -419,8 +420,9 @@ def baichuan_attention_forward_13b_quantized( if query_states.size(2) != 1 or query_states.device.type != 'xpu': attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.attn_value_fp8_matmul(attn_weights, value_states) + import xe_addons + attn_output = xe_addons.attn_value_fp8_matmul(attn_weights, + value_states) attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) diff --git a/python/llm/src/ipex_llm/transformers/models/baichuan2.py b/python/llm/src/ipex_llm/transformers/models/baichuan2.py index e496e68b..47d046ef 100644 --- a/python/llm/src/ipex_llm/transformers/models/baichuan2.py +++ b/python/llm/src/ipex_llm/transformers/models/baichuan2.py @@ -41,9 +41,9 @@ def pre_compute_inv_freq(module: torch.nn.Module): def baichuan_13b_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training or hidden_states.requires_grad): - import linear_q4_0 + import xe_addons x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous() - output = linear_q4_0.rms_norm(self.weight, x_2d, self.epsilon) + output = xe_addons.rms_norm(self.weight, x_2d, self.epsilon) return output.reshape(hidden_states.shape) input_dtype = hidden_states.dtype @@ -60,10 +60,10 @@ def baichuan_mlp_forward( x_2d = x.view(-1, x.shape[-1]) qtype = getattr(self.gate_proj, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training) and not self.down_proj.enable_xetla: - import linear_q4_0 + import xe_linear if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - return self.down_proj(linear_q4_0.mlp_forward_xpu( + return self.down_proj(xe_linear.mlp_forward_xpu( x_2d, self.gate_proj.weight.data, self.up_proj.weight.data, x_2d.shape[0], x_2d.shape[1], self.gate_proj.out_len, SILU, qtype @@ -96,9 +96,9 @@ def baichuan_attention_forward_7b( # IPEX-LLM OPT: fuse rope if should_use_fuse_rope(hidden_states, position_ids, self.training): - import linear_q4_0 - linear_q4_0.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, - query_states, key_states) + import xe_addons + xe_addons.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, + query_states, key_states) else: cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, @@ -126,18 +126,20 @@ def baichuan_attention_forward_7b( value_states.to(dtype=torch.float16), is_causal=True).to(hidden_states.dtype) elif use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_addons if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_addons if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_fp8_causal(query_states, key_states, + value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_causal(query_states, key_states, value_states) else: if use_quantize_kv: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, @@ -202,8 +204,8 @@ def baichuan_attention_forward_13b( attention_mask = attention_mask[:, None, -q_len:, :] if use_quantize_kv and q_len == 1: - import linear_q4_0 - attn_weights = linear_q4_0.query_key_fp8_matmul(query_states, key_states) + import xe_addons + attn_weights = xe_addons.query_key_fp8_matmul(query_states, key_states) else: if use_quantize_kv: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, @@ -216,8 +218,8 @@ def baichuan_attention_forward_13b( attn_weights = attn_weights.to(query_states.dtype) attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) if use_quantize_kv and q_len == 1: - import linear_q4_0 - attn_output = linear_q4_0.attn_value_fp8_matmul(attn_weights, value_states) + import xe_addons + attn_output = xe_addons.attn_value_fp8_matmul(attn_weights, value_states) else: attn_output = torch.matmul(attn_weights.to(dtype=value_states.dtype), value_states) attn_output = attn_output.transpose(1, 2) diff --git a/python/llm/src/ipex_llm/transformers/models/bloom.py b/python/llm/src/ipex_llm/transformers/models/bloom.py index 10010b72..54a6d052 100644 --- a/python/llm/src/ipex_llm/transformers/models/bloom.py +++ b/python/llm/src/ipex_llm/transformers/models/bloom.py @@ -66,12 +66,12 @@ def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: def bloom_layer_norm_forward(self, hidden_states): if use_fused_layer_norm(hidden_states, self.training): - import linear_q4_0 - result = linear_q4_0.fused_layer_norm(hidden_states, - [self.weight.size(0)], - self.weight, - self.bias, - self.eps) + import xe_addons + result = xe_addons.fused_layer_norm(hidden_states, + [self.weight.size(0)], + self.weight, + self.bias, + self.eps) # if nelement == 0, means fused norm failed, go back to python implement. if result.nelement != 0: return result diff --git a/python/llm/src/ipex_llm/transformers/models/chatglm2.py b/python/llm/src/ipex_llm/transformers/models/chatglm2.py index 8c11c3aa..2078087b 100644 --- a/python/llm/src/ipex_llm/transformers/models/chatglm2.py +++ b/python/llm/src/ipex_llm/transformers/models/chatglm2.py @@ -111,9 +111,9 @@ def should_split_qkv_tensor(query_layer, bsz, n_head, seq_len): def chatglm_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad): - import linear_q4_0 + import xe_addons x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous() - output = linear_q4_0.rms_norm(self.weight, x_2d, self.eps) + output = xe_addons.rms_norm(self.weight, x_2d, self.eps) return output.reshape(hidden_states.shape) input_dtype = hidden_states.dtype @@ -322,8 +322,8 @@ def chatglm2_quantized_attention_forward_8eb45c( context_layer = torch.matmul(attn.to(value.dtype), value) else: key, value = k_cache, v_cache - import linear_q4_0 - context_layer = linear_q4_0.sdp_fp8(query_layer, key, value, attn_bias) + import xe_addons + context_layer = xe_addons.sdp_fp8(query_layer, key, value, attn_bias) # context_layer's shape: [bs, n_head, seq_len, head_dim] -> [seq_len, bs, n_head * head_dim] context_layer = context_layer.permute(2, 0, 1, 3).contiguous().view(seq_len, batch_size, -1) @@ -572,8 +572,8 @@ def core_attn_forward_8eb45c(query_layer, key_layer, value_layer, attention_mask if use_sdp(query_layer.shape[2], key_layer.shape[2], query_layer.shape[-1], query_layer): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_layer, key_layer, value_layer, attn_bias) + import xe_addons + attn_output = xe_addons.sdp(query_layer, key_layer, value_layer, attn_bias) context_layer = attn_output.view(query_layer.shape) else: head_dim = query_layer.size(-1) diff --git a/python/llm/src/ipex_llm/transformers/models/cohere.py b/python/llm/src/ipex_llm/transformers/models/cohere.py index 1ff8b53b..9ee4f142 100644 --- a/python/llm/src/ipex_llm/transformers/models/cohere.py +++ b/python/llm/src/ipex_llm/transformers/models/cohere.py @@ -261,9 +261,8 @@ def cohere_attention_forward_quantized( cache_kwargs, new_layout=True) if q_len == 1 and query_states.device.type == 'xpu' and not self.training \ and not hidden_states.requires_grad: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, attention_mask) attn_weights = None else: key_states, value_states = restore_fp8_kv_cache(key_states, @@ -325,18 +324,18 @@ def cohere_attention_forward_origin( 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim, + self.rotary_emb.base,) kv_seq_len += 1 # update past_key_value's seem_tokens and kv caches. if self.layer_idx == 0: @@ -421,12 +420,12 @@ def cohere_attention_forward_origin( attn_weights = None elif not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 + import xe_addons if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] else: causal_mask = None - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, causal_mask) + attn_output = xe_addons.sdp(query_states, key_states, value_states, causal_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: diff --git a/python/llm/src/ipex_llm/transformers/models/gemma.py b/python/llm/src/ipex_llm/transformers/models/gemma.py index d6ac66d2..71c8dfc5 100644 --- a/python/llm/src/ipex_llm/transformers/models/gemma.py +++ b/python/llm/src/ipex_llm/transformers/models/gemma.py @@ -79,9 +79,9 @@ def should_use_fuse_rope(self, hidden_states, position_ids): def gemma_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad): - import linear_q4_0 + import xe_addons x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous() - output = linear_q4_0.rms_norm(self.weight + 1, x_2d, self.eps) + output = xe_addons.rms_norm(self.weight + 1, x_2d, self.eps) return output.reshape(hidden_states.shape) input_dtype = hidden_states.dtype @@ -100,10 +100,10 @@ def gemma_mlp_forward( bsz, hidden_size = x_2d.shape qtype = getattr(self.gate_proj, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training) and not self.down_proj.enable_xetla: - import linear_q4_0 + import xe_linear if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - out = self.down_proj(linear_q4_0.mlp_forward_xpu( + out = self.down_proj(xe_linear.mlp_forward_xpu( x_2d, self.gate_proj.weight.data, self.up_proj.weight.data, x_2d.shape[0], x_2d.shape[1], self.gate_proj.out_len, GELU, qtype @@ -146,17 +146,17 @@ def gemma_attention_forward( 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim) kv_seq_len += 1 # update past_key_value's seem_tokens and kv caches. diff --git a/python/llm/src/ipex_llm/transformers/models/internlm.py b/python/llm/src/ipex_llm/transformers/models/internlm.py index 3df5264f..aa5567df 100644 --- a/python/llm/src/ipex_llm/transformers/models/internlm.py +++ b/python/llm/src/ipex_llm/transformers/models/internlm.py @@ -398,18 +398,18 @@ def internlm_xcomposser2_attention_forward( # IPEX-LLM OPT: sdp if use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_linear if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_linear.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + attn_output = xe_linear.sdp(query_states, key_states, value_states, attention_mask) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_linear if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_linear.sdp_fp8_causal(query_states, key_states, value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_linear.sdp_causal(query_states, key_states, value_states) else: if use_quantize_kv: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, diff --git a/python/llm/src/ipex_llm/transformers/models/llama.py b/python/llm/src/ipex_llm/transformers/models/llama.py index cadb908b..cb18bc94 100644 --- a/python/llm/src/ipex_llm/transformers/models/llama.py +++ b/python/llm/src/ipex_llm/transformers/models/llama.py @@ -169,9 +169,9 @@ def llama_model_forward_4_38( def llama_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad): - import linear_q4_0 + import xe_addons x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous() - output = linear_q4_0.rms_norm(self.weight, x_2d, self.variance_epsilon) + output = xe_addons.rms_norm(self.weight, x_2d, self.variance_epsilon) return output.reshape(hidden_states.shape) input_dtype = hidden_states.dtype @@ -190,10 +190,10 @@ def llama_mlp_forward( bsz, hidden_size = x_2d.shape qtype = getattr(self.gate_proj, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training) and not self.down_proj.enable_xetla: - import linear_q4_0 + import xe_linear if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - out = self.down_proj(linear_q4_0.mlp_forward_xpu( + out = self.down_proj(xe_linear.mlp_forward_xpu( x_2d, self.gate_proj.weight.data, self.up_proj.weight.data, x_2d.shape[0], x_2d.shape[1], self.gate_proj.out_len, SILU, qtype @@ -429,18 +429,18 @@ def llama_attention_forward_4_31_quantized( dtype=hidden_states.dtype, device=device ) - 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, - tmp_cache_k, tmp_cache_v, - self.q_proj.weight.qtype, - self.v_proj.weight.qtype, - 0, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.forward_qkv(hidden_states, + self.q_proj.weight, + self.k_proj.weight, + self.v_proj.weight, + position_ids, + tmp_cache_k, tmp_cache_v, + self.q_proj.weight.qtype, + self.v_proj.weight.qtype, + 0, + self.head_dim, + self.rotary_emb.base,) else: query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) @@ -504,9 +504,8 @@ def llama_attention_forward_4_31_quantized( bsz, q_len, kv_seq_len, self.head_dim, self.num_heads, output_attentions) else: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, attention_mask) attn_weights = None attn_output = attn_output.transpose(1, 2).contiguous() @@ -562,18 +561,18 @@ def llama_attention_forward_4_31_original( kv_seq_len = past_key_value[0].shape[-2] cache_k = past_key_value[0] cache_v = past_key_value[1] - 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim, + self.rotary_emb.base,) kv_seq_len += 1 else: @@ -625,12 +624,12 @@ def llama_attention_forward_4_31_original( self.k_proj, self.v_proj, self.q_proj.weight.qtype,) - import linear_q4_0 + import xe_linear q_out_len = self.q_proj.out_len k_out_len = self.k_proj.out_len v_out_len = self.v_proj.out_len - qkv_states = linear_q4_0.mm_xetla(hidden_states, self.qkv_proj_qweight, - self.q_proj.weight.qtype) + qkv_states = xe_linear.mm_xetla(hidden_states, self.qkv_proj_qweight, + self.q_proj.weight.qtype) query_states = qkv_states[:, :, :q_out_len] key_states = qkv_states[:, :, q_out_len:q_out_len + k_out_len] value_states = qkv_states[:, :, q_out_len + k_out_len:] @@ -712,8 +711,8 @@ def llama_attention_forward_4_31_original( attn_weights = None elif not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: @@ -811,19 +810,19 @@ def llama_attention_selective_batching_forward_4_31( past_k = new_cache_k past_v = new_cache_v hidden_states = hidden_states.view(1, -1) - 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, - past_k, past_v, - self.q_proj.weight.qtype, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim, - self.rotary_emb.base, - ) + import xe_linear + query_states, key_states, value_states = xe_linear.forward_qkv(hidden_states, + self.q_proj.weight, + self.k_proj.weight, + self.v_proj.weight, + position_ids, + past_k, past_v, + self.q_proj.weight.qtype, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim, + self.rotary_emb.base, + ) kv_seq_len += 1 else: if self.config.pretraining_tp > 1: @@ -1028,18 +1027,18 @@ def llama_attention_forward_4_38_quantized( dtype=hidden_states.dtype, device=device ) - 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, - tmp_cache_k, tmp_cache_v, - self.q_proj.weight.qtype, - self.v_proj.weight.qtype, - 0, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.forward_qkv(hidden_states, + self.q_proj.weight, + self.k_proj.weight, + self.v_proj.weight, + position_ids, + tmp_cache_k, tmp_cache_v, + self.q_proj.weight.qtype, + self.v_proj.weight.qtype, + 0, + self.head_dim, + self.rotary_emb.base,) else: query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) @@ -1176,13 +1175,12 @@ def llama_attention_forward_4_38_quantized( dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 + import xe_addons if cache_position is not None: new_attn_mask = attention_mask[:, :, kv_seq_len-q_len:kv_seq_len, 0:kv_seq_len] else: new_attn_mask = attention_mask - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - new_attn_mask) + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, new_attn_mask) attn_weights = None if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): @@ -1251,18 +1249,18 @@ def llama_attention_forward_4_38_original( 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim, + self.rotary_emb.base,) kv_seq_len += 1 # update past_key_value's seem_tokens and kv caches. if self.layer_idx == 0: @@ -1319,13 +1317,13 @@ def llama_attention_forward_4_38_original( self.k_proj, self.v_proj, self.q_proj.weight.qtype,) - import linear_q4_0 + import xe_linear q_out_len = self.q_proj.out_len k_out_len = self.k_proj.out_len v_out_len = self.v_proj.out_len - qkv_states = linear_q4_0.mm_xetla(hidden_states, - self.qkv_proj_qweight, - self.q_proj.weight.qtype) + qkv_states = xe_linear.mm_xetla(hidden_states, + self.qkv_proj_qweight, + self.q_proj.weight.qtype) query_states = qkv_states[:, :, :q_out_len] key_states = qkv_states[:, :, q_out_len:q_out_len + k_out_len] value_states = qkv_states[:, :, q_out_len + k_out_len:] @@ -1425,8 +1423,9 @@ def llama_attention_forward_4_38_original( attn_weights = None elif not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, new_attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, + new_attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: diff --git a/python/llm/src/ipex_llm/transformers/models/mistral.py b/python/llm/src/ipex_llm/transformers/models/mistral.py index fcbb6961..d9307cb3 100644 --- a/python/llm/src/ipex_llm/transformers/models/mistral.py +++ b/python/llm/src/ipex_llm/transformers/models/mistral.py @@ -278,17 +278,17 @@ def mistral_attention_forward_quantized( dtype=hidden_states.dtype, device=device ) - 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, - tmp_cache_k, tmp_cache_v, - self.q_proj.weight.qtype, - self.v_proj.weight.qtype, - 0, - self.head_dim) + import xe_linear + query_states, key_states, value_states = xe_linear.forward_qkv(hidden_states, + self.q_proj.weight, + self.k_proj.weight, + self.v_proj.weight, + position_ids, + tmp_cache_k, tmp_cache_v, + self.q_proj.weight.qtype, + self.v_proj.weight.qtype, + 0, + self.head_dim) else: query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) @@ -427,9 +427,9 @@ def mistral_attention_forward_quantized( attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) attn_weights = None attn_output_size = (bsz, self.num_heads, q_len, self.head_dim) @@ -476,17 +476,17 @@ def mistral_attention_forward_original( kv_seq_len = past_key_value[0].shape[-2] cache_k = past_key_value[0] cache_v = past_key_value[1] - 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim) kv_seq_len += 1 else: @@ -496,13 +496,13 @@ def mistral_attention_forward_original( self.k_proj, self.v_proj, self.q_proj.qtype) - import linear_q4_0 + import xe_linear q_out_len = self.q_proj.out_len k_out_len = self.k_proj.out_len v_out_len = self.v_proj.out_len - qkv_states = linear_q4_0.mm_xetla(hidden_states, - self.qkv_proj_qweight, - self.q_proj.qtype) + qkv_states = xe_linear.mm_xetla(hidden_states, + self.qkv_proj_qweight, + self.q_proj.qtype) query_states = qkv_states[:, :, :q_out_len] key_states = qkv_states[:, :, q_out_len:q_out_len + k_out_len] value_states = qkv_states[:, :, q_out_len + k_out_len:] @@ -592,8 +592,8 @@ def mistral_attention_forward_original( attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) elif use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): # new fp16 sdp doesn't require repeat_kv - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None attn_output = attn_output.transpose(1, 2).contiguous() @@ -695,17 +695,17 @@ def mistral_attention_forward_4_36_quantized( dtype=hidden_states.dtype, device=device ) - 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, - tmp_cache_k, tmp_cache_v, - self.q_proj.weight.qtype, - self.v_proj.weight.qtype, - 0, - self.head_dim) + import xe_linear + query_states, key_states, value_states = xe_linear.forward_qkv(hidden_states, + self.q_proj.weight, + self.k_proj.weight, + self.v_proj.weight, + position_ids, + tmp_cache_k, tmp_cache_v, + self.q_proj.weight.qtype, + self.v_proj.weight.qtype, + 0, + self.head_dim) else: query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) @@ -852,9 +852,8 @@ def mistral_attention_forward_4_36_quantized( attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, attention_mask) attn_weights = None attn_output_size = (bsz, self.num_heads, q_len, self.head_dim) @@ -905,17 +904,17 @@ def mistral_attention_forward_4_36_original( 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim) kv_seq_len += 1 # update past_key_value's seem_tokens and kv caches. @@ -931,13 +930,13 @@ def mistral_attention_forward_4_36_original( self.k_proj, self.v_proj, self.q_proj.qtype) - import linear_q4_0 + import xe_linear q_out_len = self.q_proj.out_len k_out_len = self.k_proj.out_len v_out_len = self.v_proj.out_len - qkv_states = linear_q4_0.mm_xetla(hidden_states, - self.qkv_proj_qweight, - self.q_proj.qtype) + qkv_states = xe_linear.mm_xetla(hidden_states, + self.qkv_proj_qweight, + self.q_proj.qtype) query_states = qkv_states[:, :, :q_out_len] key_states = qkv_states[:, :, q_out_len:q_out_len + k_out_len] value_states = qkv_states[:, :, q_out_len + k_out_len:] @@ -1033,8 +1032,8 @@ def mistral_attention_forward_4_36_original( attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) elif use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): # new fp16 sdp doesn't require repeat_kv - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None attn_output = attn_output.transpose(1, 2).contiguous() diff --git a/python/llm/src/ipex_llm/transformers/models/mixtral.py b/python/llm/src/ipex_llm/transformers/models/mixtral.py index 8223ad61..9069c49c 100644 --- a/python/llm/src/ipex_llm/transformers/models/mixtral.py +++ b/python/llm/src/ipex_llm/transformers/models/mixtral.py @@ -105,8 +105,8 @@ def mixtral_moeblock_forward(self, elif bs < 256 and hidden_states.device.type == 'xpu': final_hidden_states = torch.zeros((batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device) - import linear_q4_0 - indexes = linear_q4_0.get_moe_indexes(selected_experts.to(torch.int32).cpu(), 8) + import xe_linear + indexes = xe_linear.get_moe_indexes(selected_experts.to(torch.int32).cpu(), 8) for expert_idx in range(self.num_experts): expert_layer = self.experts[expert_idx] idx_list = indexes[0][expert_idx] @@ -184,18 +184,18 @@ def mixtral_attention_forward( 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, - self.v_proj.weight.qtype, - kv_seq_len, - self.head_dim, - self.rotary_emb.base,) + import xe_linear + query_states, key_states, value_states = xe_linear.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, + self.v_proj.weight.qtype, + kv_seq_len, + self.head_dim, + self.rotary_emb.base,) kv_seq_len += 1 # update past_key_value's seem_tokens and kv caches. if self.layer_idx == 0: @@ -209,8 +209,8 @@ def mixtral_attention_forward( # 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 = linear_q4_0.forward_qk(hidden_states, + # import xe_linear + # query_states, key_states = xe_linear.forward_qk(hidden_states, # self.q_proj.weight, # self.k_proj.weight, # position_ids, @@ -333,8 +333,8 @@ def mixtral_attention_forward( is_causal=True) attn_weights = None elif use_sdp(query_states.shape[2], key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: @@ -389,8 +389,8 @@ def mixtral_mlp_forward( ) -> torch.Tensor: qtype = getattr(self.w1, "qtype", None) if mlp_fusion_check(x, qtype, self.training) and not self.w1.enable_xetla: - import linear_q4_0 - return self.w2(linear_q4_0.mlp_forward_xpu( + import xe_linear + return self.w2(xe_linear.mlp_forward_xpu( x, self.w1.weight.data, self.w3.weight.data, x.shape[0], x.shape[1], self.w1.out_len, SILU, qtype, diff --git a/python/llm/src/ipex_llm/transformers/models/phi3.py b/python/llm/src/ipex_llm/transformers/models/phi3.py index c1ea3ed2..94c44e9b 100644 --- a/python/llm/src/ipex_llm/transformers/models/phi3.py +++ b/python/llm/src/ipex_llm/transformers/models/phi3.py @@ -108,17 +108,17 @@ def attention_forward( # IPEX-LLM OPT: fuse rope if should_use_fuse_rope(hidden_states, position_ids, self.training): - import linear_q4_0 + import xe_addons if self.rotary_emb.__class__.__name__ == "Phi3RotaryEmbedding": # 4k - linear_q4_0.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, - query_states, key_states) + xe_addons.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, + query_states, key_states) else: # 128k if kv_seq_len > self.rotary_emb.original_max_position_embeddings: - linear_q4_0.rotary_half_inplaced(self.rotary_emb.long_inv_freq, position_ids, - query_states, key_states) + xe_addons.rotary_half_inplaced(self.rotary_emb.long_inv_freq, + position_ids, query_states, key_states) else: - linear_q4_0.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, - query_states, key_states) + xe_addons.rotary_half_inplaced(self.rotary_emb.inv_freq, + position_ids, query_states, key_states) # todo: fuse scaling_factor query_states *= self.rotary_emb.scaling_factor key_states *= self.rotary_emb.scaling_factor @@ -132,18 +132,19 @@ def attention_forward( self.layer_idx, None) if use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_addons if isinstance(past_key_value, DynamicFp8Cache): - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_addons if isinstance(past_key_value, DynamicFp8Cache): - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_fp8_causal(query_states, key_states, value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_causal(query_states, key_states, value_states) else: if isinstance(past_key_value, DynamicFp8Cache): key_states, value_states = restore_fp8_kv_cache(key_states, value_states, @@ -204,8 +205,8 @@ def mlp_forward( qtype = getattr(self.gate_proj, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training): x_2d = x_2d.contiguous() - import linear_q4_0 - return self.down_proj(linear_q4_0.mlp_forward_xpu( + import xe_linear + return self.down_proj(xe_linear.mlp_forward_xpu( x_2d, self.gate_proj.weight.data, self.up_proj.weight.data, x_2d.shape[0], x_2d.shape[1], self.gate_proj.out_features, SILU, qtype @@ -293,9 +294,9 @@ def phi3v_model_forward_wrapper(origin_model_forward): def phi3_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad): - import linear_q4_0 + import xe_addons x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous() - output = linear_q4_0.rms_norm(self.weight, x_2d, self.variance_epsilon) + output = xe_addons.rms_norm(self.weight, x_2d, self.variance_epsilon) return output.reshape(hidden_states.shape) input_dtype = hidden_states.dtype diff --git a/python/llm/src/ipex_llm/transformers/models/qwen.py b/python/llm/src/ipex_llm/transformers/models/qwen.py index 2856eb1c..99f1d726 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen.py @@ -92,9 +92,10 @@ def qwen_attention_forward( rotary_pos_emb = rotary_pos_emb_list[0] if use_fuse_rope: rot_dim = rotary_pos_emb[0].size(-1) - import linear_q4_0 - linear_q4_0.rotary_half_inplaced(inv_freq, position_ids, - query_states[..., :rot_dim], key_states[..., :rot_dim]) + import xe_addons + xe_addons.rotary_half_inplaced(inv_freq, position_ids, + query_states[..., :rot_dim], + key_states[..., :rot_dim]) else: rotary_pos_emb = [i[:, -q_len:, :, :].transpose(1, 2) for i in rotary_pos_emb] query_states = apply_rotary_pos_emb(query_states, rotary_pos_emb) @@ -124,11 +125,11 @@ def qwen_attention_forward( value_states.to(dtype=torch.float16), is_causal=True).to(hidden_states.dtype) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_addons if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_fp8_causal(query_states, key_states, value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_causal(query_states, key_states, value_states) else: if q_len > 1: causal_mask = torch.tril( @@ -146,13 +147,13 @@ def qwen_attention_forward( attention_mask = None if use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_addons if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, - attention_mask) + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) else: if use_quantize_kv: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, @@ -221,9 +222,9 @@ def qwen_attention_forward_registered( rotary_pos_emb = rotary_pos_emb_list[0] if use_fuse_rope: rot_dim = rotary_pos_emb[0].size(-1) - import linear_q4_0 - linear_q4_0.rotary_half_inplaced(inv_freq, position_ids, - query_states[..., :rot_dim], key_states[..., :rot_dim]) + import xe_linear + xe_linear.rotary_half_inplaced(inv_freq, position_ids, + query_states[..., :rot_dim], key_states[..., :rot_dim]) else: rotary_pos_emb = [i[:, -q_len:, :, :].transpose(1, 2) for i in rotary_pos_emb] query_states = apply_rotary_pos_emb(query_states, rotary_pos_emb) @@ -253,11 +254,11 @@ def qwen_attention_forward_registered( value_states.to(dtype=torch.float16), is_causal=True).to(hidden_states.dtype) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_linear if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_linear.sdp_fp8_causal(query_states, key_states, value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_linear.sdp_causal(query_states, key_states, value_states) else: if q_len > 1: causal_mask = registered_causal_mask[ @@ -272,13 +273,13 @@ def qwen_attention_forward_registered( attention_mask = None if use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_linear if use_quantize_kv: - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_linear.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, - attention_mask) + attn_output = xe_linear.sdp(query_states, key_states, value_states, + attention_mask) else: if use_quantize_kv: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, @@ -310,10 +311,10 @@ def qwen_mlp_forward(self, x: torch.Tensor) -> torch.Tensor: x_2d = x.view(-1, x.shape[-1]) qtype = getattr(self.w1, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training) and not self.w1.enable_xetla: - import linear_q4_0 + import xe_linear if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - return self.c_proj(linear_q4_0.mlp_forward_xpu( + return self.c_proj(xe_linear.mlp_forward_xpu( x_2d, self.w2.weight.data, self.w1.weight.data, x_2d.shape[0], x_2d.shape[1], self.w2.out_len, SILU, qtype diff --git a/python/llm/src/ipex_llm/transformers/models/qwen2.py b/python/llm/src/ipex_llm/transformers/models/qwen2.py index c02d28f9..002f6589 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen2.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen2.py @@ -310,9 +310,9 @@ def qwen2_attention_forward( kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) if should_use_fuse_rope(hidden_states, position_ids, self.training): - import linear_q4_0 - linear_q4_0.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, - query_states, key_states) + import xe_addons + xe_addons.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, + query_states, key_states) else: cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, @@ -337,18 +337,19 @@ def qwen2_attention_forward( value_states.to(device, dtype=torch.float16), is_causal=True).to(hidden_states.dtype) elif use_sdp(q_len, kv_seq_len, self.head_dim, query_states): - import linear_q4_0 + import xe_addons if isinstance(past_key_value, DynamicFp8Cache): - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, - attention_mask) + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) else: - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): - import linear_q4_0 + import xe_addons if isinstance(past_key_value, DynamicFp8Cache): - attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_fp8_causal(query_states, key_states, value_states) else: - attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states) + attn_output = xe_addons.sdp_causal(query_states, key_states, value_states) else: if isinstance(past_key_value, DynamicFp8Cache): key_states, value_states = restore_fp8_kv_cache(key_states, value_states, diff --git a/python/llm/src/ipex_llm/transformers/models/qwen2_moe.py b/python/llm/src/ipex_llm/transformers/models/qwen2_moe.py index c136c8f7..9f14ca08 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen2_moe.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen2_moe.py @@ -372,8 +372,8 @@ def qwen2moe_attention_forward_quantized( self.layer_idx, cache_kwargs) if q_len == 1 and query_states.device.type == 'xpu' and not self.training \ and not hidden_states.requires_grad: - import linear_q4_0 - attn_weights = linear_q4_0.query_key_fp8_matmul(query_states, key_states) + import xe_addons + attn_weights = xe_addons.query_key_fp8_matmul(query_states, key_states) else: key_states, value_states = restore_fp8_kv_cache(key_states, value_states, query_states.dtype) @@ -404,8 +404,8 @@ def qwen2moe_attention_forward_quantized( p=self.attention_dropout, training=self.training) if q_len == 1 and query_states.device.type == 'xpu' and not self.training \ and not hidden_states.requires_grad: - import linear_q4_0 - attn_output = linear_q4_0.attn_value_fp8_matmul(attn_weights, value_states) + import xe_addons + attn_output = xe_addons.attn_value_fp8_matmul(attn_weights, value_states) else: attn_output = torch.matmul(attn_weights, value_states) @@ -456,12 +456,12 @@ def qwen2moe_attention_forward_origin( 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 + import xe_linear 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, self.v_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) + query_states, key_states, value_states = xe_linear.forward_qkv_bias(*args) kv_seq_len += 1 if self.layer_idx == 0: past_key_value._seen_tokens = kv_seq_len @@ -613,12 +613,12 @@ def qwen2moe_attention_forward_sdpa( 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 + import xe_linear 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, self.v_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) + query_states, key_states, value_states = xe_linear.forward_qkv_bias(*args) kv_seq_len += 1 if self.layer_idx == 0: past_key_value._seen_tokens = kv_seq_len @@ -765,8 +765,8 @@ def qwen2moe_moeblock_forward(self, hidden_states: torch.Tensor): elif bs < 256 and hidden_states.device.type == 'xpu': final_hidden_states = torch.zeros((batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device) - import linear_q4_0 - indexes = linear_q4_0.get_moe_indexes(selected_experts.to(torch.int32).cpu(), 60) + import xe_linear + indexes = xe_linear.get_moe_indexes(selected_experts.to(torch.int32).cpu(), 60) for expert_idx in range(self.num_experts): expert_layer = self.experts[expert_idx] idx_list = indexes[0][expert_idx] diff --git a/python/llm/src/ipex_llm/transformers/models/qwen_vl.py b/python/llm/src/ipex_llm/transformers/models/qwen_vl.py index c518e1a6..3f494de9 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen_vl.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen_vl.py @@ -162,8 +162,8 @@ def qwen_attention_forward_vl( if not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key.shape[2], self.head_dim, query): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query, key, value, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query, key, value, attention_mask) attn_output = attn_output.view(query.shape) attn_output = attn_output.transpose(1, 2) attn_weight = None diff --git a/python/llm/src/ipex_llm/transformers/models/rwkv4.py b/python/llm/src/ipex_llm/transformers/models/rwkv4.py index 2c0dc7ae..0deef7cd 100644 --- a/python/llm/src/ipex_llm/transformers/models/rwkv4.py +++ b/python/llm/src/ipex_llm/transformers/models/rwkv4.py @@ -56,8 +56,8 @@ def extract_key_value(self, hidden, state=None): self.time_mix_receptance.data, ]).to(dtype=hidden.dtype) - import linear_q4_0 - mixed_result = linear_q4_0.rwkv_time_shift(hidden, shifted, self.mixed_mix) + import xe_linear + mixed_result = xe_linear.rwkv_time_shift(hidden, shifted, self.mixed_mix) key, value, receptance = mixed_result key = self.key(key) @@ -92,8 +92,8 @@ def rwkv_linear_attention_xpu( time_decay = -torch.exp(time_decay) # `num_state`, `den_state`, `max_state` will be modified during this call - import linear_q4_0 - output = linear_q4_0.rwkv_linear_attention_v4( + import xe_linear + output = xe_linear.rwkv_linear_attention_v4( time_decay, time_first, key, @@ -167,8 +167,8 @@ def rwkv_ffn_forward( self.mixed_mix = torch.cat([self.time_mix_key.data, self.time_mix_receptance.data]).to(dtype=hidden.dtype) - import linear_q4_0 - mixed_result = linear_q4_0.rwkv_time_shift(hidden, shifted, self.mixed_mix) + import xe_linear + mixed_result = xe_linear.rwkv_time_shift(hidden, shifted, self.mixed_mix) key, receptance = mixed_result key = torch.square(torch.relu(self.key(key))) diff --git a/python/llm/src/ipex_llm/transformers/models/rwkv5.py b/python/llm/src/ipex_llm/transformers/models/rwkv5.py index 5619c16f..58889e5a 100644 --- a/python/llm/src/ipex_llm/transformers/models/rwkv5.py +++ b/python/llm/src/ipex_llm/transformers/models/rwkv5.py @@ -58,8 +58,8 @@ def extract_key_value(self, hidden, state=None): self.time_mix_gate.data, ]).to(dtype=hidden.dtype) - import linear_q4_0 - mixed_result = linear_q4_0.rwkv_time_shift(hidden, shifted, self.mixed_mix) + import xe_linear + mixed_result = xe_linear.rwkv_time_shift(hidden, shifted, self.mixed_mix) key, value, receptance, gate = mixed_result key = self.key(key) @@ -98,8 +98,8 @@ def rwkv_linear_attention_xpu( time_first = time_first.float() # `state` will be updated inplaced during this call - import linear_q4_0 - out = linear_q4_0.rwkv_linear_attention_v5( + import xe_linear + out = xe_linear.rwkv_linear_attention_v5( time_decay, time_first, receptance, @@ -236,8 +236,8 @@ def rwkv_ffn_forward_wrapper(origin_rwkv_ffn_forward): self.mixed_mix = torch.cat([self.time_mix_key.data, self.time_mix_receptance.data]).to(dtype=hidden.dtype) - import linear_q4_0 - mixed_result = linear_q4_0.rwkv_time_shift(hidden, shifted, self.mixed_mix) + import xe_linear + mixed_result = xe_linear.rwkv_time_shift(hidden, shifted, self.mixed_mix) key, receptance = mixed_result key = torch.square(torch.relu(self.key(key))) diff --git a/python/llm/src/ipex_llm/transformers/models/stablelm.py b/python/llm/src/ipex_llm/transformers/models/stablelm.py index 4e8b685f..c8a84557 100644 --- a/python/llm/src/ipex_llm/transformers/models/stablelm.py +++ b/python/llm/src/ipex_llm/transformers/models/stablelm.py @@ -267,8 +267,9 @@ def stablelm_attention_forward_original( attn_weights = None elif not self.training and not hidden_states.requires_grad and \ use_sdp(q_len, key_states.shape[2], self.head_dim, query_states): - import linear_q4_0 - attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp(query_states, key_states, value_states, + attention_mask) attn_output = attn_output.view(query_states.shape) attn_weights = None else: @@ -420,8 +421,8 @@ def stablelm_attention_forward_quantized( value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) else: - import linear_q4_0 - attn_weights = linear_q4_0.query_key_fp8_matmul(query_states, key_states) + import xe_addons + attn_weights = xe_addons.query_key_fp8_matmul(query_states, key_states) attn_weights = attn_weights / math.sqrt(self.head_dim) @@ -444,8 +445,8 @@ def stablelm_attention_forward_quantized( if query_states.size(2) != 1 or query_states.device.type != 'xpu': attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.attn_value_fp8_matmul(attn_weights, value_states) + import xe_addons + attn_output = xe_addons.attn_value_fp8_matmul(attn_weights, value_states) attn_output_size = (bsz, self.num_heads, q_len, self.head_dim) invalidInputError(attn_output.size() == attn_output_size, diff --git a/python/llm/src/ipex_llm/transformers/models/starcoder2.py b/python/llm/src/ipex_llm/transformers/models/starcoder2.py index 014bd51d..b0e83f48 100644 --- a/python/llm/src/ipex_llm/transformers/models/starcoder2.py +++ b/python/llm/src/ipex_llm/transformers/models/starcoder2.py @@ -135,8 +135,9 @@ def attention_forward( self.layer_idx, None) if use_quantize_kv and q_len == 1: - import linear_q4_0 - attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, attention_mask) + import xe_addons + attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, + attention_mask) attn_weights = None else: if use_quantize_kv: diff --git a/python/llm/src/ipex_llm/transformers/models/utils.py b/python/llm/src/ipex_llm/transformers/models/utils.py index f902c059..3e2878b9 100644 --- a/python/llm/src/ipex_llm/transformers/models/utils.py +++ b/python/llm/src/ipex_llm/transformers/models/utils.py @@ -20,7 +20,8 @@ import warnings from ipex_llm.utils.common import invalidInputError from ipex_llm.ggml.quantize import ggml_tensor_qtype from ipex_llm.transformers.utils import get_ipex_version, get_xpu_device_type -from ipex_llm.transformers.low_bit_linear import SYM_INT4, SYM_INT8, FP8E5, IQ2_XXS, FP4, FP8E4, FP6 +from ipex_llm.transformers.low_bit_linear import SYM_INT4, SYM_INT8, FP8E5, IQ2_XXS, FP4, FP8E4,\ + FP6, ASYM_INT4 from ipex_llm.transformers.convert import is_deepspeed_available FP8_KV_ALLOC_LENGTH = 512 @@ -128,10 +129,10 @@ def append_fp8_kv_cache(k_cache, v_cache, key, value): new_k_cache = k_cache.as_strided(new_size, k_cache.stride(), storage_offset=0) new_v_cache = v_cache.as_strided(new_size, v_cache.stride(), storage_offset=0) - import linear_q4_0 - linear_q4_0.quantize_key_value(key, value, - new_k_cache[:, :, cur_length:new_length, :], - new_v_cache[:, :, cur_length:new_length, :]) + import xe_addons + xe_addons.quantize_key_value(key, value, + new_k_cache[:, :, cur_length:new_length, :], + new_v_cache[:, :, cur_length:new_length, :]) return new_k_cache, new_v_cache @@ -140,8 +141,8 @@ def restore_fp8_kv_cache(k_cache, v_cache, dtype): key_states = torch.empty(k_cache.shape, device=k_cache.device, dtype=dtype) value_states = torch.empty(v_cache.shape, device=v_cache.device, dtype=dtype) - import linear_q4_0 - linear_q4_0.dequantize_key_value(k_cache, v_cache, key_states, value_states) + import xe_addons + xe_addons.dequantize_key_value(k_cache, v_cache, key_states, value_states) return key_states, value_states @@ -211,13 +212,13 @@ def apply_rotary_pos_emb_no_cache_xpu(q, k, position_ids, model_family, rope_the if q.device.type != "xpu": invalidInputError(False, f"only xpu is supported in this function") - import linear_q4_0 + import xe_addons q_embed = torch.empty(q.shape, dtype=q.dtype, device=q.device) k_embed = torch.empty(k.shape, dtype=k.dtype, device=k.device) if model_family in ["llama", "baichuan", "internlm", "aquila", "gpt_neox", "mistral", "mixtral"]: - linear_q4_0.apply_rotary_embedding_half_q_and_k(q, k, position_ids, - q_embed, k_embed, rope_theta) + xe_addons.apply_rotary_embedding_half_q_and_k(q, k, position_ids, + q_embed, k_embed, rope_theta) return q_embed, k_embed else: invalidInputError(False, @@ -228,11 +229,12 @@ def apply_rotary_pos_emb_cache_freq_xpu(q, k, sin, cos, model_family, position_i if q.device.type != "xpu": invalidInputError(False, f"only xpu is supported in this function") - import linear_q4_0 + import xe_addons q_embed = torch.empty(q.shape, dtype=q.dtype, device=q.device) k_embed = torch.empty(k.shape, dtype=k.dtype, device=k.device) if model_family in ["qwen", "mixtral"]: - linear_q4_0.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, q_embed, k_embed) + xe_addons.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, + q_embed, k_embed) elif model_family in ["qwen2", "yuan", "stablelm", "qwen2_moe", "internlm"]: cos = cos.to(q.dtype) sin = sin.to(q.dtype) @@ -240,11 +242,13 @@ def apply_rotary_pos_emb_cache_freq_xpu(q, k, sin, cos, model_family, position_i sin = sin.squeeze(1).squeeze(0) # [seq_len, dim] cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] - linear_q4_0.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, q_embed, k_embed) + xe_addons.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, + q_embed, k_embed) elif model_family in ["gemma", "phi3"]: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) - linear_q4_0.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, q_embed, k_embed) + xe_addons.apply_rotary_embedding_half_q_and_k_cache_freq(q, k, sin, cos, + q_embed, k_embed) else: invalidInputError(False, f"{model_family} is not supported.") diff --git a/python/llm/src/ipex_llm/transformers/models/yuan.py b/python/llm/src/ipex_llm/transformers/models/yuan.py index 6e0674ef..a2d48bb5 100644 --- a/python/llm/src/ipex_llm/transformers/models/yuan.py +++ b/python/llm/src/ipex_llm/transformers/models/yuan.py @@ -97,10 +97,10 @@ def yuan_mlp_forward( bsz, hidden_size = x_2d.shape qtype = getattr(self.up_proj, "qtype", None) if mlp_fusion_check(x_2d, qtype, self.training): - import linear_q4_0 + import xe_linear if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - out = self.down_proj(linear_q4_0.mlp_forward_xpu( + out = self.down_proj(xe_linear.mlp_forward_xpu( x_2d, self.up_proj.weight.data, self.gate_proj.weight.data, x_2d.shape[0], x_2d.shape[1], self.up_proj.out_len, SILU, qtype @@ -268,8 +268,8 @@ def yuan_attention_forward_quantized( query_states.dtype) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) else: - import linear_q4_0 - attn_weights = linear_q4_0.query_key_fp8_matmul(query_states, key_states) + import xe_addons + attn_weights = xe_addons.query_key_fp8_matmul(query_states, key_states) attn_weights = attn_weights / math.sqrt(self.head_dim) @@ -292,8 +292,8 @@ def yuan_attention_forward_quantized( if query_states.size(2) != 1 or device.type != 'xpu': attn_output = torch.matmul(attn_weights, value_states) else: - import linear_q4_0 - attn_output = linear_q4_0.attn_value_fp8_matmul(attn_weights, value_states) + import xe_addons + attn_output = xe_addons.attn_value_fp8_matmul(attn_weights, value_states) invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), "`attn_output` should be of size "