From 9a5062e6c4a4dc05bf930288113e131126377f68 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 23 Jan 2024 01:44:10 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index cac4943a6c..9f60adc743 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -144,8 +144,10 @@ 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) + Kn = torch.cat([K1, Kn], dim = 1) + Vn = torch.cat([V1, Vn], dim = 1) # Grouped query attention if n_groups != 1: