From 7c2b04254cbae8e58dedf8ea57375f721053233e Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 2 Feb 2024 16:59:02 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 994eedb479..ca85138195 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) - if not hasattr(self, "paged_attention"): - 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_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[:kv_seq_len].permute(1, 2, 0, 3) - Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3) + Kn = torch.cat([K1, Kn], dim = 2) + Vn = torch.cat([V1, Vn], dim = 2) + # if not hasattr(self, "paged_attention"): + # 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_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[: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: