diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index a05ecc10aa..87aca586b1 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -426,7 +426,7 @@ def LlamaDecoderLayer_fast_forward(
         past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
     """
     bsz, q_len, hd = hidden_states.size()
-    if False:#(past_key_value is not None and q_len == 1):
+    if (past_key_value is not None and q_len == 1):
         # Self Attention
         residual = hidden_states
         hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
@@ -597,7 +597,7 @@ def LlamaModel_fast_forward(
 
     hidden_states = inputs_embeds
 
-    if self.gradient_checkpointing and self.training:
+    if past_key_values is None and self.gradient_checkpointing and self.training:
         if use_cache:
             logger.warning_once(
                 "Unsloth: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`"
@@ -656,7 +656,7 @@ def LlamaModel_fast_forward(
     pass
 
     bsz, q_len, hd = hidden_states.size()
-    if False:#(past_key_values is not None and q_len == 1):
+    if (past_key_values is not None and q_len == 1):
         hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states)
     else:
         hidden_states = fast_rms_layernorm(self.norm, hidden_states)
@@ -719,12 +719,12 @@ def LlamaForCausalLM_fast_forward(
     )
 
     hidden_states = outputs[0]
-    # bsz, q_len, hd = hidden_states.shape
-    # if bsz == 1 and q_len == 1:
-    #     logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
-    #     logits = logits.unsqueeze(0).unsqueeze(0)
-    # else:
-    logits = self.lm_head(hidden_states)
+    bsz, q_len, hd = hidden_states.shape
+    if bsz == 1 and q_len == 1:
+        logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
+        logits = logits.unsqueeze(0).unsqueeze(0)
+    else:
+        logits = self.lm_head(hidden_states)
     pass
 
     loss = None
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 26081da023..9944367d21 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -210,12 +210,12 @@ def MistralForCausalLM_fast_forward(
     )
 
     hidden_states = outputs[0]
-    # bsz, q_len, hd = hidden_states.shape
-    # if bsz == 1 and q_len == 1:
-    #     logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
-    #     logits = logits.unsqueeze(0).unsqueeze(0)
-    # else:
-    logits = self.lm_head(hidden_states)
+    bsz, q_len, hd = hidden_states.shape
+    if bsz == 1 and q_len == 1:
+        logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
+        logits = logits.unsqueeze(0).unsqueeze(0)
+    else:
+        logits = self.lm_head(hidden_states)
     pass
 
     loss = None