Revert "Change rolling batch generations order (#1046)" (#1048)

This reverts commit 7f9b9a8.
deepjavalibrary · Sep 1, 2023 · 70bbe55 · 70bbe55
1 parent 086ec90
commit 70bbe55
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py b/engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py
@@ -132,7 +132,7 @@ def _prefill_and_decode(self, new_batch):
             if self.cache:
                 decode_generations, decode_next_batch = self.model.generate_token(
                     self.cache)
-                decode_generations.extend(generations)
+                generations.extend(decode_generations)
 
                 # concatenate with the existing batch of the model
                 if prefill_next_batch is None:
@@ -141,7 +141,7 @@ def _prefill_and_decode(self, new_batch):
                     self.cache = prefill_next_batch
                 else:
                     self.cache = self.model.batch_type.concatenate(
-                        [decode_next_batch, prefill_next_batch])
+                        [prefill_next_batch, decode_next_batch])
             else:
                 self.cache = prefill_next_batch
         else: