Enhance SamplingParams (#96)

2023-05-11 15:45:30 -07:00
parent 55f8b0a5de
commit 42f1042e1c
7 changed files with 36 additions and 54 deletions
--- a/cacheflow/core/scheduler.py
+++ b/cacheflow/core/scheduler.py
@@ -316,7 +316,7 @@ class Scheduler:
                    continue

                # Check if the sequence has reached the maximum number of steps.
-                max_num_steps = self.sampling_params[group_id].max_num_steps
+                max_num_steps = self.sampling_params[group_id].max_tokens
                if self.num_steps[group_id] == max_num_steps:
                    self._free_seq(seq)
                    continue