diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 25ad680008..60fa43bc8e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -266,7 +266,7 @@ def LlamaAttention_fast_forward_inference( 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] = K1.permute(2, 0, 1, 3) + self.paged_attention_K[: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)