diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index b042689e14..122e6f2869 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -230,7 +230,7 @@ def LlamaAttention_fast_forward_inference(
     n_groups   = self.num_key_value_groups
     n_kv_heads = self.num_key_value_heads
     head_dim   = self.head_dim
-    assert(n_kv_heads * n_groups == n_heads)
+    # ssert(n_kv_heads * n_groups == n_heads)
 
     Qn = fast_linear_forward(self.q_proj, Xn)
     Kn = fast_linear_forward(self.k_proj, Xn)
@@ -241,19 +241,19 @@ def LlamaAttention_fast_forward_inference(
 
     seq_len = K1.shape[-2]
     kv_seq_len = seq_len + 1
-    # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
-    # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
-    cos = self.rotary_emb.cos_cached[seq_len]
-    sin = self.rotary_emb.sin_cached[seq_len]
-    h = head_dim // 2
+    cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
+    Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
+    # cos = self.rotary_emb.cos_cached[seq_len]
+    # sin = self.rotary_emb.sin_cached[seq_len]
+    # h = head_dim // 2
 
-    RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = Xn.dtype, device = "cuda")
-    RH_Q[:,:,:,:h] = Qn[:,:,:,h:]; RH_Q[:,:,:,h:] = Qn[:,:,:,:h]; torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]);
-    Qn *= cos; Qn.addcmul_(RH_Q, sin);
+    # RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = Xn.dtype, device = "cuda")
+    # RH_Q[:,:,:,:h] = Qn[:,:,:,h:]; RH_Q[:,:,:,h:] = Qn[:,:,:,:h]; torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]);
+    # Qn *= cos; Qn.addcmul_(RH_Q, sin);
 
-    RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda")
-    RH_K[:,:,:,:h] = Kn[:,:,:,h:]; RH_K[:,:,:,h:] = Kn[:,:,:,:h]; torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]);
-    Kn *= cos; Kn.addcmul_(RH_K, sin);
+    # RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda")
+    # RH_K[:,:,:,:h] = Kn[:,:,:,h:]; RH_K[:,:,:,h:] = Kn[:,:,:,:h]; torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]);
+    # Kn *= cos; Kn.addcmul_(RH_K, sin);
     
     # New KV cache
     Kn = torch.cat([K1, Kn], dim = 2)
@@ -290,7 +290,7 @@ def fast_mlp_inference(self, X):
     gate *= up
 
     # X = self.down_proj(gate)
-    down = fast_linear_forward(self.down_proj, gate)
+    down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:X.shape[2]])
     return down
 pass