# # Copyright 2016 The BigDL Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os import torch 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, ASYM_INT4 FP8_KV_ALLOC_LENGTH = 512 KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256)) # used in fused mlp forward SILU = 0 GELU = 1 def decoding_fast_path_qtype_check(proj): qtype = getattr(proj, "qtype", None) return qtype in [SYM_INT4, FP8E5, FP4] def init_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device): key_cache_storage = torch.empty(batch_size, num_heads, max_length, head_dim, dtype=dtype, device=device) value_cache_storage = torch.empty(batch_size, num_heads, max_length, head_dim, dtype=dtype, device=device) key_cache = key_cache_storage.as_strided((batch_size, num_heads, current_length, head_dim), key_cache_storage.stride(), storage_offset=0) value_cache = value_cache_storage.as_strided((batch_size, num_heads, current_length, head_dim), value_cache_storage.stride(), storage_offset=0) return key_cache, value_cache def extend_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device): # empty cache to reduce gpu memory if device.type == 'xpu': torch.xpu.empty_cache() return init_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device) def append_kv_cache(cache_k, cache_v, key_states, value_states): new_size = (cache_k.size(0), cache_k.size(1), cache_k.size(2) + key_states.size(2), cache_k.size(3)) new_cache_k = cache_k.as_strided(new_size, cache_k.stride(), storage_offset=0) new_cache_k[:, :, cache_k.size(2):cache_k.size(2) + key_states.size(2), :] = key_states new_cache_v = cache_v.as_strided(new_size, cache_v.stride(), storage_offset=0) new_cache_v[:, :, cache_v.size(2):cache_v.size(2) + key_states.size(2), :] = value_states return new_cache_k, new_cache_v def use_quantize_kv_cache(linear: torch.nn.Module, x: torch.Tensor, kv_group: int = 1) -> bool: if os.environ.get("BIGDL_QUANTIZE_KV_CACHE", None) is not None: warnings.warn( "`BIGDL_QUANTIZE_KV_CACHE` is deprecated and will be removed in future releases. " "Please use `IPEX_LLM_QUANTIZE_KV_CACHE` instead." ) return os.environ["BIGDL_QUANTIZE_KV_CACHE"] == "1" elif os.environ.get("IPEX_LLM_QUANTIZE_KV_CACHE", None) is not None: return os.environ["IPEX_LLM_QUANTIZE_KV_CACHE"] == "1" elif os.environ.get("IPEX_LLM_LOW_MEM", None) is not None: return os.environ["IPEX_LLM_LOW_MEM"] == "1" else: return x.device.type == 'xpu' and kv_cache_device_check(x, kv_group) \ and hasattr(linear, "qtype") and \ linear.qtype != ggml_tensor_qtype["fp16"] and linear.qtype != ggml_tensor_qtype["bf16"] def kv_cache_device_check(x: torch.Tensor, kv_group: int) -> bool: return (get_xpu_device_type(x) in ["mtl", "lnl"] and kv_group <= 1) or \ ((get_xpu_device_type(x) == "arc" or get_xpu_device_type(x) == "flex") and 1 < x.size(0) and x.size(0) <= 8) def init_fp8_kv_cache(batch_size, num_heads, current_length, head_dim, device): max_length = current_length + FP8_KV_ALLOC_LENGTH k_cache_storage = torch.empty(batch_size, num_heads, max_length, head_dim, dtype=torch.uint8, device=device) k_cache = k_cache_storage.as_strided((batch_size, num_heads, 0, head_dim), k_cache_storage.stride(), storage_offset=0) v_cache_storage = torch.empty(batch_size, num_heads, max_length, head_dim, dtype=torch.uint8, device=device) v_cache = v_cache_storage.as_strided((batch_size, num_heads, 0, head_dim), v_cache_storage.stride(), storage_offset=0) return k_cache, v_cache def append_fp8_kv_cache(k_cache, v_cache, key, value): batch_size, num_heads, cur_length, head_dim = k_cache.shape new_length = cur_length + key.size(2) new_size = (batch_size, num_heads, new_length, head_dim) if k_cache.stride(1) < new_length * k_cache.size(3): new_k_cache, new_v_cache = init_fp8_kv_cache(batch_size, num_heads, new_length, head_dim, key.device) new_k_cache = new_k_cache.as_strided(new_size, new_k_cache.stride(), storage_offset=0) new_v_cache = new_v_cache.as_strided(new_size, new_v_cache.stride(), storage_offset=0) new_k_cache[:, :, :cur_length, :] = k_cache new_v_cache[:, :, :cur_length, :] = v_cache else: 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 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 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 xe_addons xe_addons.dequantize_key_value(k_cache, v_cache, key_states, value_states) return key_states, value_states def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., :x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def rotate_every_two(x): x1 = x[:, :, :, ::2] x2 = x[:, :, :, 1::2] x = torch.stack((-x2, x1), dim=-1) return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)') def should_use_fuse_rope(hidden_states, position_ids, training): return ( hidden_states.device.type == "xpu" and not training and not hidden_states.requires_grad and position_ids is not None ) def apply_rotary_pos_emb(q, k, cos, sin, position_ids, model_family): if model_family in ["llama", "baichuan", "internlm", "aquila", "gpt_neox", "mistral", "mixtral", "qwen2", "yuan", "stablelm", "qwen2_moe"]: # The first two dimensions of cos and sin are always 1, so we can `squeeze` them. cos = cos.squeeze(1).squeeze(0) # [seq_len, dim] 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] q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed elif model_family == "llama2": cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed elif model_family in ["gptj", "chatglm"]: q_embed = (q * cos) + (rotate_every_two(q) * sin) k_embed = (k * cos) + (rotate_every_two(k) * sin) return q_embed, k_embed else: invalidInputError(False, f"{model_family} is not supported.") def apply_ipex_rotate_every_two(q, k, cos, sin): # ipex's apply_rotary_embedding_two_qk can change the origin storage, # so q/k will get the result directly. from ipex_llm.transformers.utils import get_ipex_version if get_ipex_version() >= "2.1.10+xpu": torch.ops.torch_ipex.apply_rotary_embedding_two_qk( q, k, sin, cos, q, k ) else: torch.ops.torch_ipex.apply_rotary_embedding(q, sin, cos, q) torch.ops.torch_ipex.apply_rotary_embedding(k, sin, cos, k) def is_enough_kv_cache_room_4_36(past_key_value, idx, seq_len=1): # to determinate if is enough kv cache room in transformers==4.36 # seq_len for current seq len # For llama like kv cache, i.e., [bs, n_head, seq_len, head_dim] 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) + seq_len) * \ past_key_value.key_cache[idx].size(3) def is_enough_kv_cache_room_4_31(past_key_value, seq_len=1): # to determinate if is enough kv cache room in transformers between 4.31 and 4.35 # seq_len for current seq len # For llama like kv cache, i.e., [bs, n_head, seq_len, head_dim] return past_key_value is not None and \ past_key_value[0].stride()[1] >= \ (past_key_value[0].size(2) + seq_len) * past_key_value[0].size(3) def use_flash_attention(query, key, attention_mask=None): # here we support query's shape is always [batch_size, head_num, q_len, head_dim], # key's shape is always [batch_size, head_num, k_len, head_dim] invalidInputError(query.dim() == 4, "Here query input of use_flash_attention should be [batch_size, " "head_num, q_len, head_dim]") invalidInputError(key.dim() == 4, "Here key input of use_flash_attention should be [batch_size, " "head_num, k_len, head_dim]") bsz, _, q_len, _ = query.size() k_len = key.size()[2] # check whether ipex flash attention can be used if q_len != k_len: # now only use flash attention for first token # as it seems have no performance benifit for rest token now return False if query.device.type != "xpu": # ipex flash attention only support for xpu return False ipex_version = get_ipex_version() if ipex_version <= "2.0.110+xpu": # ipex flash attention is supported from ipex 2.1 return False if not torch.xpu.has_xetla(): # ipex flash attention is only supported for xetla # may update this later return False elif get_xpu_device_type(query) != "pvc": return False if query.dtype not in [torch.float32, torch.float16]: # only use flash attention for fp32/fp16 input return False if bsz > 1: # as flash attention doesn't support attn_mask in ipex 2.1, # so it will cause output error for padded batch input if attention_mask is None: return True else: # TODO: below logic may change for different model # attention mask shape : [bsz, 1, q_len, k_len] if attention_mask[0].squeeze()[0, 0].item() != 0: # first batch contains padding # otherwise we suppose it should be a upper triangular matrix # at the same time, the diagonal is also 0 return False elif not attention_mask.equal(attention_mask[0].repeat(bsz, 1, 1, 1)): # check whether mask of every batch is the same return False return True def use_sdp(q_len, kv_len, head_dim, query_states): return ( query_states.device.type == "xpu" and query_states.dtype in [torch.float, torch.half] # fp32/fp16 and head_dim in [-1, 64, 80, 96, 128] and q_len != kv_len # next token and q_len <= 32 # lookup ) def use_sdp_causal(q_len, kv_len, head_dim, query_states, training): return ( q_len == kv_len # first token and head_dim in [-1, 64, 80, 96, 128] # for now and query_states.device.type == "xpu" # GPU and query_states.dtype in [torch.float, torch.half] # fp32/fp16 and not query_states.requires_grad and not training # not training ) def use_sdp_non_causal(head_dim, device, dtype): return ( head_dim in [64, 80, 128] and device.type == "xpu" # GPU and dtype in [torch.float, torch.half] # fp32/fp16 ) def mlp_fusion_check(x, qtype, training): if x.numel() // x.size(-1) != 1: return False if x.device.type != 'xpu': return False if qtype not in [SYM_INT4, FP8E5, FP4, IQ2_XXS, FP6]: return False if training or x.requires_grad: return False if qtype == FP6: device = get_xpu_device_type(x) if device in ["mtl", "lnl"]: return False return True def use_decoding_fast_path(proj, use_fuse_rope, enough_kv_room, bs, qtype_check=decoding_fast_path_qtype_check): if proj is None: return False device = get_xpu_device_type(proj.weight) if not qtype_check(proj): return False if not use_fuse_rope: return False if not enough_kv_room: return False if bs != 1: return False if proj.enable_xetla: return False if device in ["uhd"]: return False return True def use_xmx(x: torch.Tensor, qtype: int): device = get_xpu_device_type(x) return ( os.environ.get("BIGDL_LLM_XMX_DISABLED", "0") != "1" and device in ["arc", "flex", "pvc"] and qtype in [SYM_INT4, SYM_INT8, FP8E4, FP8E5] and ( (device == "pvc" and 1 < x.size(0) <= 16) or (device != "pvc" and 1 < x.size(0) <= 64) ) ) def fp16_fusion_check(proj, x, training): # only use fp16 fusion on PVC inference if proj is None: return False if not hasattr(proj, "qtype"): return False if proj.qtype != ggml_tensor_qtype["fp16"]: return False if proj.weight_type != 2: return False if training: return False if x.requires_grad: return False device_type = get_xpu_device_type(x) if device_type != "pvc": return False return True def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def update_past_key_value(past_key_value, key_states, value_states, kv_seq_len, use_quantize_kv, device): bsz, num_heads, _, head_dim = key_states.shape if use_quantize_kv: if past_key_value is None: k_cache, v_cache = init_fp8_kv_cache( bsz, num_heads, kv_seq_len, head_dim, device=device ) else: k_cache, v_cache = past_key_value key_states, value_states = append_fp8_kv_cache(k_cache, v_cache, key_states, value_states) else: if past_key_value is None: max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH k_cache, v_cache = init_kv_cache(bsz, num_heads, head_dim, kv_seq_len, max_cache_length, dtype=key_states.dtype, device=device) k_cache[...] = key_states v_cache[...] = value_states key_states = k_cache value_states = v_cache else: k_cache, v_cache = past_key_value if k_cache.stride(1) < kv_seq_len * k_cache.size(3): max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH new_k_cache, new_v_cache = extend_kv_cache(bsz, num_heads, head_dim, k_cache.size(2), max_cache_length, dtype=k_cache.dtype, device=device) new_k_cache[...] = k_cache new_v_cache[...] = v_cache k_cache = new_k_cache v_cache = new_v_cache key_states, value_states = append_kv_cache(k_cache, v_cache, key_states, value_states) return key_states, value_states def should_use_compresskv(x: torch.Tensor, prompt_len: int): use_compress_kv = os.environ.get("IPEX_LLM_COMPRESS_KV_CACHE", None) perf_mode = os.environ.get("IPEX_LLM_PERFORMANCE_MODE", None) if perf_mode == "1": return False else: if use_compress_kv is None: return ( get_xpu_device_type(x) in ["mtl", "lnl"] and prompt_len >= 1800 and prompt_len <= 4500 ) else: return x.device.type == 'xpu' and use_compress_kv == "1" def get_compresskv_attn_mask(key_states: torch.Tensor, attention_mask: torch.Tensor): if attention_mask is not None: attention_mask = attention_mask[:, :, :, -key_states.size(2):] return attention_mask def get_q_proj_or_qkv_proj(self): if hasattr(self, "q_proj"): proj = self.q_proj elif hasattr(self, "qkv_proj"): proj = self.qkv_proj return proj def make_cache_contiguous_inplaced(cos: torch.Tensor, sin: torch.Tensor): if not cos.is_contiguous(): new_cos = cos.contiguous() new_sin = sin.contiguous() cos.set_(new_cos) sin.set_(new_sin)