Update llama.py

This commit is contained in:
Daniel Han-Chen 2024-02-02 13:52:12 +11:00
commit 6cc9835021

View file

@ -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)