Fix the reprefill of evicted seqs with invalid draft tokens (#4564)

Author: RunningLeon

lmdeploy/pytorch/paging/seq_states/states.py

@@ -91,6 +91,8 @@ def activate(self):
         self.to_state(RunningState)
 
     def evict(self):
+        # clean up meta before evict
+        self.seq.cleanup()
         self.to_state(WaitingState)
 
 

lmdeploy/pytorch/strategies/ar/sequence.py

@@ -87,6 +87,10 @@ def set_step(self, step: int):
             if len(self.all_routed_experts) > step:
                 self.all_routed_experts.resize(step)
 
+    def cleanup(self):
+        """Setup history meta after sequence stopped or cancelled."""
+        pass
+
 
 class ARSequenceStrategy(SequenceStrategy):
 

lmdeploy/pytorch/strategies/ar_spec/sequence.py

@@ -166,6 +166,13 @@ def set_step(self, step: int):
             if len(self.all_routed_experts) > step:
                 self.all_routed_experts.resize(step)
 
+    def cleanup(self):
+        """Setup history meta after sequence stopped or cancelled."""
+        # truncate history cache to valid ids
+        self.history_cache.resize(self.num_valid_ids)
+        self._num_history_ids = self.num_valid_ids - 1
+        self._num_token_ids = 1
+
 
 class ARSpecSequenceStrategy(ARSequenceStrategy):