From 3bfb0ebf9dbf857dba9098b04c25de29a7834612 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sun, 4 Feb 2024 03:50:31 +1100
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 29 ++++++++++-------------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c53f5600ce..4a5b33f2a6 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -316,21 +316,12 @@ def LlamaAttention_fast_forward(
             K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
             V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
         pass
-
-        attn_weights = torch.matmul(Q, K.transpose(2, 3)) / math_sqrt(self.head_dim)
-        attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
-        attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(Q.dtype)
-        attn_weights = torch.nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        A = torch.matmul(attn_weights, V)
-
         # Must be contiguous or else results are False!
         # https://github.com/pytorch/pytorch/issues/112577
-        # Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
-        # # Needs (batch_size, n_heads, seq_len, head_dim)
-        # # is_casual and attention_mask must not be both set!
-        # A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
+        Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
+        # Needs (batch_size, n_heads, seq_len, head_dim)
+        # is_casual and attention_mask must not be both set!
+        A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
         # Go back to (batch_size, seq_len, n_heads, head_dim)
         A = A.transpose(1, 2).contiguous()
     pass
@@ -367,7 +358,7 @@ def LlamaDecoderLayer_fast_forward(
             (see `past_key_values`).
         past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
     """
-    if False:#past_key_value is not None:
+    if past_key_value is not None:
         do_prefill = not hasattr(self.self_attn, "paged_attention")
 
         # Self Attention
@@ -480,7 +471,7 @@ def LlamaModel_fast_forward(
     pass
 
     # We already handle KV cache position_ids ourselves.
-    if (past_key_values_length != 0):
+    if False:#(past_key_values_length != 0):
         position_ids = torch.arange(
             past_key_values_length, seq_length + past_key_values_length,
             dtype  = torch.int32,
@@ -673,7 +664,7 @@ def LlamaForCausalLM_fast_forward(
     *args, **kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
 
-    if causal_mask is None:# and past_key_values is None:
+    if causal_mask is None and past_key_values is not None:
         causal_mask = xformers.attn_bias.LowerTriangularMask()
 
     output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -685,8 +676,8 @@ def LlamaForCausalLM_fast_forward(
     # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
     self.model._has_no_labels = labels is None
 
-    if False:#past_key_values is not None and \
-        #hasattr(self.model.layers[0].self_attn, "paged_attention"):
+    if past_key_values is not None and \
+        hasattr(self.model.layers[0].self_attn, "paged_attention"):
         outputs = LlamaModel_fast_forward_inference(
             self.model,
             input_ids,
@@ -709,7 +700,7 @@ def LlamaForCausalLM_fast_forward(
 
     hidden_states = outputs[0]
     bsz, q_len, hd = hidden_states.shape
-    if False:#bsz == 1 and q_len == 1:
+    if bsz == 1 and q_len == 1:
         logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
         logits = logits.unsqueeze(0).unsqueeze(0)
     else: