From 5f3c51b394e1765a1b4bca45232df0ea6a3e1151 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 1 Feb 2024 23:41:11 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index b14a4bb793..93d91cad25 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -265,8 +265,8 @@ def LlamaAttention_fast_forward_inference( self.paged_attention = torch.empty((2, bsz, n_kv_heads, 2048, head_dim), dtype = dtype, device = "cuda") self.paged_attention_K = self.paged_attention[0] self.paged_attention_V = self.paged_attention[1] - self.paged_attention_K[:,:,:kv_seq_len,:] = K1 - self.paged_attention_V[:,:,:kv_seq_len,:] = V1 + self.paged_attention_K[:,:,:kv_seq_len,:] = Kn + self.paged_attention_V[:,:,:kv_seq_len,:] = Vn pass Kn = self.paged_attention_K[:,:,seq_len,:] = K1 Vn = self.paged_attention_V[:,:,seq_len,:] = V1 @@ -346,7 +346,7 @@ def LlamaAttention_fast_forward( position_ids, ) return A, None, past_key_value - elif hasattr(self, "paged_attention_K"): + elif hasattr(self, "paged_attention"): del self.paged_attention_K del self.paged_attention_V del self.paged_attention