fix stopwords kv cache (#3494)

Author: grimoire

lmdeploy/pytorch/engine/engine.py

@@ -50,8 +50,8 @@ class InferOutput:
 def _tensorlize_block_offsets(block_offsets, dtype=torch.int32):
     """tensorlize block_offsets."""
     from torch.nn.utils.rnn import pad_sequence
-    block_offsets = [torch.from_numpy(off).to(dtype) for off in block_offsets]
-    block_offsets = pad_sequence(block_offsets, batch_first=True)
+    block_offsets = [torch.from_numpy(off) for off in block_offsets]
+    block_offsets = pad_sequence(block_offsets, batch_first=True).to(dtype)
     return block_offsets
 
 
@@ -563,6 +563,7 @@ def __update_max_new_tokens(msg):
                     req.data['token_ids'],
                     multimodals=req.data.get('input_multimodals'),
                     embeddings=req.data.get('input_embeddings'),
+                    append_tokens=True,
                 )
                 msg.num_new_tokens = 0
                 msg.sampling_param = sampling_param
@@ -721,8 +722,6 @@ def update_running(self, running: SeqList, next_token_ids: torch.Tensor, stopped
             msg.update_token_ids(update_token, model_meta=model_meta)
             msg.num_new_tokens += 1
             if stop:
-                update_token = _EMPTY_TOKEN
-                msg.update_token_ids(update_token, model_meta=model_meta)
                 msg.status = MessageStatus.STOPPED
 
     def update_running_migration(self, running: SeqList, next_token_ids: np.ndarray, stopped: torch.Tensor,

lmdeploy/pytorch/messages.py

@@ -592,11 +592,15 @@ def update_token_ids(self,
                          token_ids: Tensor,
                          multimodals: MultiModalInputs = None,
                          embeddings: List[InputEmbeddings] = None,
-                         model_meta: Dict[str, Any] = None):
+                         model_meta: Dict[str, Any] = None,
+                         append_tokens: bool = False):
         """Update token ids, old token ids will be added to history."""
         old_num_history_ids = self._num_history_ids
 
-        self._num_history_ids += self._num_token_ids
+        # update history
+        if not append_tokens:
+            self._num_history_ids += self._num_token_ids
+
         # update history image nums
         self._num_history_images += self._num_images
         self._num_images = 0
@@ -626,7 +630,10 @@ def update_token_ids(self,
             token_ids = np.array(token_ids)
         if token_ids.ndim == 0:
             token_ids = token_ids[None]
-        self._num_token_ids = len(token_ids)
+        if append_tokens:
+            self._num_token_ids += len(token_ids)
+        else:
+            self._num_token_ids = len(token_ids)
         self.history_cache.append(token_ids)
         self.random_offsets += 1
         self.arrive_time = time.time()