Added _reorder_cach to gpt_neox

imangohari1 · Mar 7, 2025 · 7e9180b · 7e9180b
1 parent 11140b2
commit 7e9180b
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -530,6 +530,15 @@ def prepare_inputs_for_generation(
 
         return model_inputs
 
+    def _reorder_cache(self, past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past[:2])
+                + layer_past[2:],
+            )
+        return reordered_past
+
 
 def apply_customized_rope(q, k, cos, sin, position_ids, training=True):
     if q.device.type == "hpu" and FusedRoPE is not None: