fix

Author: shihaobai

lightllm/server/router/model_infer/infer_batch.py

@@ -395,14 +395,15 @@ def diverse_copy(self, req_manager, is_prefill):
             prefix_len = prev_req.shared_kv_node.node_prefix_total_len
         else:
             prefix_len = 0
-        pre_input_token_ids = prev_req.get_input_token_ids()
+        prefix_len = max(prefix_len, prev_req.cur_kv_len)
+        pre_input_token_ids = prev_req.get_chuncked_input_token_ids()
         cache_token_id = req_manager.req_to_token_indexs[prev_req.req_idx][prefix_len : len(pre_input_token_ids)]
         # update the InferReq status and mem_manager status for cache sharing
         for req_id in self.req_ids_group[:]:
             if req_id == convert_sub_id_to_group_id(req_id):
                 continue
             req = g_infer_context.requests_mapping[req_id]
             req.finish_status.set_status(FinishStatus.NO_FINISH)
-            input_token_ids = req.get_input_token_ids()
+            input_token_ids = req.get_chuncked_input_token_ids()
             req_manager.req_to_token_indexs[req.req_idx][prefix_len : len(input_token_ids)] = cache_token_id
             assert len(input_token_ids) == len(pre_input_token_ids)

lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py

@@ -38,8 +38,7 @@ def diverse_copy(self, groups: List[InferReqGroup]):
         for i in range(len(groups)):
             req_group = groups[i]
             best_of = req_group.best_of()
-            _0_req_obj = req_group.get_req(0)
-            if best_of > 1 and _0_req_obj.get_chuncked_input_token_len() == _0_req_obj.get_cur_total_len():
+            if best_of > 1:
                 req_group.diverse_copy(g_infer_context.req_manager, is_prefill=True)
                 batch_idx.extend([i for _ in range(best_of)])
             else:
@@ -58,7 +57,6 @@ def decode(self):
 
         if aborted_reqs:
             g_infer_context.filter_reqs(aborted_reqs)
-
         if prefill_reqs:
             group_reqs = [
                 g_infer_context.requests_mapping[req.req_id]

lightllm/server/router/req_queue/__init__.py

@@ -1,8 +1,8 @@
 from .continues_batch.impl import ContinuesBatchQueue
-from .continues_batch.beam_impl import BeamContinuesBatchQueue
 from .continues_batch.impl_for_pd_decode import QueueForPDDecode
 from .chunked_prefill.impl_for_pd_prefill import QueueForPDChunkedPrefill
 from .chunked_prefill.impl import ChunkedPrefillQueue
+from .chunked_prefill.beam_impl import BeamContinuesBatchQueue
 from .dp_base_queue import DpQueue