From 7993288caab18c28e8483c07869a381d19cd6d24 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 20 Aug 2024 15:21:18 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 42f01b7eb5..376b4b4ebc 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -187,8 +187,9 @@ def LlamaAttention_fast_forward_inference( # 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[position_ids].unsqueeze(1) - sin = self.rotary_emb.sin_cached[position_ids].unsqueeze(1) + cos, sin = self.rotary_emb.get_cached(kv_seq_len) + cos = cos[position_ids].unsqueeze(1) + sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim RH_Q = self.RH_Q