From 40e8848c4e7d4bba7490ce9e1a02f43e5dea1509 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 2 Feb 2024 15:56:56 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 60fa43bc8e..994eedb479 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -259,19 +259,19 @@ def LlamaAttention_fast_forward_inference( Kn *= cos; Kn.addcmul_(RH_K, sin); # New KV cache - Kn = torch.cat([K1, Kn], dim = 2) - Vn = torch.cat([V1, Vn], dim = 2) + # Kn = torch.cat([K1, Kn], dim = 2) + # Vn = torch.cat([V1, Vn], dim = 2) if not hasattr(self, "paged_attention"): - self.paged_attention = torch.empty((2048, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda") + self.paged_attention = torch.zeros((2048, 2, bsz, n_kv_heads, 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[:seq_len] = K1.permute(2, 0, 1, 3) - self.paged_attention_K[:seq_len] = V1.permute(2, 0, 1, 3) + self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) pass self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3) self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3) - Kn = self.paged_attention_K.permute(1, 2, 0, 3) - Vn = self.paged_attention_V.permute(1, 2, 0, 3) + Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3) + Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3) # Grouped query attention if n_groups != 1: