diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 7bbfc017d2..a05ecc10aa 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -310,7 +310,7 @@ def LlamaAttention_fast_forward(
     bsz, q_len, _ = hidden_states.size()
 
     # Check for inference
-    if past_key_value is not None and q_len == 1 and bsz == 1:
+    if past_key_value is not None and q_len == 1:
         A, past_key_value = LlamaAttention_fast_forward_inference(
             self,
             hidden_states,
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 52f4873e35..26081da023 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -49,7 +49,7 @@ def MistralAttention_fast_forward(
     bsz, q_len, _ = hidden_states.size()
 
     # Check for inference
-    if True:
+    if past_key_value is not None and q_len == 1:
         A, past_key_value = LlamaAttention_fast_forward_inference(
             self,
             hidden_states,