From f231c4f395be0735a6b1947fe084a2d81f4190fd Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 1 Feb 2024 23:41:23 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d91cad25..9d267bde6c 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,:] = Kn - self.paged_attention_V[:,:,:kv_seq_len,:] = Vn + self.paged_attention_K[:,:,:kv_seq_len,:] = K1 + self.paged_attention_V[:,:,:kv_seq_len,:] = V1 pass Kn = self.paged_attention_K[:,:,seq_len,:] = K1 Vn = self.paged_attention_V[:,:,seq_len,:] = V1