Repository: unslothai/unsloth
File: unsloth/models/llama.py (Python)

Code before:
from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model
from peft import PeftModelForCausalLM, PeftModelForSequenceClassification
from ..save import patch_saving_functions
import re, os, inspect, math, sys
import types
try:
    from huggingface_hub.utils import get_token
except:
    # Old HF Hub versions <= 0.0.25
    from huggingface_hub.utils._token import get_token
pass
from triton import __version__ as triton_version
HAS_XFORMERS = xformers is not None
BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None

def clean_gpu_cache():
    if DEVICE_TYPE == "xpu":
        torch.xpu.empty_cache()
    else:
        torch.cuda.empty_cache()

def get_current_device():
    if DEVICE_TYPE == "xpu":
        return torch.xpu.current_device()
    else:
        return torch.cuda.current_device()

def original_apply_qkv(self, X):
    Q = self.q_proj(X)
    K = self.k_proj(X)
    V = self.v_proj(X)
    return Q, K, V
pass


def original_apply_o(self, X):
    O = self.o_proj(X)
    return O
pass

from math import sqrt as math_sqrt
KV_CACHE_INCREMENT = 512 # KV Cache update size
torch_nn_functional_softmax = torch.nn.functional.softmax
# SDPA has GQA internally
SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__

# Fix new HF's inference code
def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs,):
    past_key_values = kwargs.get("past_key_values", None)
    if past_key_values is not None:
        # Check for uninitialized DynamicCache

Review comment (suggestion):
```suggestion
if DEVICE_TYPE == "xpu":
    get_current_device = torch.xpu.current_device
else:
    get_current_device = torch.cuda.current_device
```

Code after:
from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model
from peft import PeftModelForCausalLM, PeftModelForSequenceClassification
from ..save import patch_saving_functions
import re, os, inspect, math, sys
import types
try:
    from huggingface_hub.utils import get_token
except:
    # Old HF Hub versions <= 0.0.25
    from huggingface_hub.utils._token import get_token
pass
from triton import __version__ as triton_version
HAS_XFORMERS = xformers is not None
BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None

def clean_gpu_cache():
    if DEVICE_TYPE == "xpu":
        torch.xpu.empty_cache()
    else:
        torch.cuda.empty_cache()

if DEVICE_TYPE == "xpu":
    get_current_device = torch.xpu.current_device
else:
    get_current_device = torch.cuda.current_device

def original_apply_qkv(self, X):
    Q = self.q_proj(X)
    K = self.k_proj(X)
    V = self.v_proj(X)
    return Q, K, V
pass


def original_apply_o(self, X):
    O = self.o_proj(X)
    return O
pass

from math import sqrt as math_sqrt
KV_CACHE_INCREMENT = 512 # KV Cache update size
torch_nn_functional_softmax = torch.nn.functional.softmax
# SDPA has GQA internally
SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__

# Fix new HF's inference code
def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs,):
    past_key_values = kwargs.get("past_key_values", None)
    if past_key_values is not None:
        # Check for uninitialized DynamicCache
        if len(past_key_values) == 0: