diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 53abedf52f..780bc7544c 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -124,7 +124,7 @@ def LlamaAttention_fast_forward_inference(
 
     # Prefill phase
     # if not hasattr(self, "paged_attention"):
-    # if do_prefill:
+    if do_prefill:
     #     self.paged_attention = torch.empty((KV_CACHE_INCREMENT+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda")
     #     self.paged_attention_K = self.paged_attention[:,0]
     #     self.paged_attention_V = self.paged_attention[:,1]
@@ -134,7 +134,7 @@ def LlamaAttention_fast_forward_inference(
     #     self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda")
     #     self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda")
     #     self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT), dtype = dtype, device = "cuda")
-    #     self.scalar = 1.0 / math_sqrt(self.head_dim)
+        self.scalar = 1.0 / math_sqrt(self.head_dim)
     # elif kv_seq_len >= self.paged_attention.shape[0]:
     #     self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim))
     #     self.paged_attention_K = self.paged_attention[:,0]