Refactor InfiniAttention class to return past_key_value in infiniAtte…

…ntion.py and modeling_qwen_transformers.py
jlamprou · Apr 22, 2024 · 7581c10 · 7581c10
1 parent c18eb55
commit 7581c10
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/infiniAttention.py b/infiniAttention.py
@@ -123,7 +123,7 @@ def forward(
         combined_output = combined_output.transpose(1, 2).contiguous()
         combined_output = combined_output.view(bsz, q_len, self.hidden_size)
         final_output = self.o_proj(combined_output)
-        return final_output, None, None, (M, z)
+        return final_output, None, past_key_value, (M, z)
 
     def _retrieve_from_memory(self, Q, M, z):
         # Retrieve context from compressive memory using linear attention (Eq. 3)

diff --git a/modeling_qwen_transformers.py b/modeling_qwen_transformers.py
@@ -576,7 +576,7 @@ def forward(
         combined_output = combined_output.transpose(1, 2).contiguous()
         combined_output = combined_output.view(bsz, q_len, self.hidden_size)
         final_output = self.o_proj(combined_output)
-        return final_output, None, None, (M, z)
+        return final_output, None, past_key_value, (M, z)
 
     def _retrieve_from_memory(self, Q, M, z):
         # Retrieve context from compressive memory using linear attention (Eq. 3)