fix

Author: hiworldwzj

lightllm/server/router/model_infer/mode_backend/generic_padded_pre_process.py

@@ -33,13 +33,17 @@ def padded_prepare_prefill_inputs(
     b_seq_len = []
     batch_multimodal_params = []
     b_ready_cache_len = []
+    b_mtp_index = []
+    b_prefill_has_output = []
+
     for req in req_objs:
 
         run_reqs.append(req)
         batch_multimodal_params.append(req.multimodal_params)
         b_req_idx.append(req.req_idx)
 
         input_token_ids = req.get_chuncked_input_token_ids()
+        b_prefill_has_output.append(False if len(input_token_ids) < req.get_cur_total_len() else True)
         seq_len = len(input_token_ids)
         input_token_len = seq_len - req.cur_kv_len
         input_id = input_token_ids[req.cur_kv_len :]
@@ -49,27 +53,32 @@ def padded_prepare_prefill_inputs(
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, input_token_len)
         b_ready_cache_len.append(req.cur_kv_len)
+        b_mtp_index.append(0)
 
     # padding fake req for prefill
     for _ in range(padded_req_num):
         input_ids.append([1])
         b_req_idx.append(g_infer_context.req_manager.HOLD_REQUEST_ID)
         b_seq_len.append(1)
+        b_mtp_index.append(0)
+        b_prefill_has_output.append(False)
         b_ready_cache_len.append(0)
         total_token_num += 1
         max_len_in_batch = max(max_len_in_batch, 1)
 
     input_ids = np.concatenate(input_ids, dtype=np.int64)
-    input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cuda")
-    b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cuda")
-    b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cuda")
-    b_ready_cache_len = torch.tensor(b_ready_cache_len, dtype=torch.int32, device="cuda")
+
+    input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cpu")
+    b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
+    b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
+    b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
+    b_ready_cache_len = torch.tensor(b_ready_cache_len, dtype=torch.int32, device="cpu")
 
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
         g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(input_ids.shape[0] - padded_req_num)
-    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(input_ids.shape[0] - padded_req_num).cuda()
+    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(input_ids.shape[0] - padded_req_num)
     g_infer_state_lock.release()
 
     if padded_req_num > 0:
@@ -85,11 +94,13 @@ def padded_prepare_prefill_inputs(
         total_token_num=total_token_num,
         max_len_in_batch=max_len_in_batch,
         input_ids=input_ids,
-        mem_indexes=mem_indexes,
+        mem_indexes_cpu=mem_indexes,
         b_req_idx=b_req_idx,
+        b_mtp_index=b_mtp_index,
         b_seq_len=b_seq_len,
         b_ready_cache_len=b_ready_cache_len,
         is_prefill=True,
+        b_prefill_has_output_cpu=b_prefill_has_output,
     )
     if is_multimodal:
         model_input.multimodal_params = batch_multimodal_params
@@ -98,64 +109,62 @@ def padded_prepare_prefill_inputs(
 
 
 def padded_prepare_decode_inputs(
-    req_objs: List[InferReq], dest_batch_size: Optional[int] = None, is_multimodal=False
+    req_objs: List[InferReq], dest_batch_size: Optional[int] = None
 ) -> Tuple[ModelInput, List[InferReq], int]:
+    mtp_step_num = get_env_start_args().mtp_step
     run_reqs = []
     total_token_num = 0
     max_len_in_batch = 0
-    input_ids = []
     b_req_idx = []
+    b_mtp_index = []
     b_seq_len = []
-
     for req in req_objs:
         run_reqs.append(req)
         b_req_idx.append(req.req_idx)
-        input_token_ids = req.get_input_token_ids()
-        input_id = input_token_ids[-1]
-        seq_len = len(input_token_ids)
+        seq_len = req.get_cur_total_len()
         assert req.cur_kv_len == seq_len - 1
         b_seq_len.append(seq_len)
-        input_ids.append(input_id)
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
+        b_mtp_index.append(0)
         # process the draft tokens.
-        for step in range(len(req.mtp_gen_token_ids)):
+        for step in range(req.mtp_step):
             run_reqs.append(req)
             b_req_idx.append(req.req_idx)
             seq_len += 1
             b_seq_len.append(seq_len)
-            input_ids.append(req.mtp_gen_token_ids[step])
             total_token_num += seq_len
             max_len_in_batch = max(max_len_in_batch, seq_len)
+            b_mtp_index.append(step + 1)
 
     if dest_batch_size is None:
         if len(run_reqs) == 0:
             dest_batch_size = 1
         else:
-            dest_batch_size = len(run_reqs)
+            dest_batch_size = len(run_reqs) * (1 + mtp_step_num)
     else:
-        assert len(run_reqs) <= dest_batch_size
+        assert len(run_reqs) * (1 + mtp_step_num) <= dest_batch_size
 
-    padded_req_num = dest_batch_size - len(run_reqs)
+    padded_req_num = dest_batch_size - len(run_reqs) * (1 + mtp_step_num)
 
     # padding fake req for decode
     for _ in range(padded_req_num):
-        input_ids.append(1)
         seq_len = 2
         b_req_idx.append(g_infer_context.req_manager.HOLD_REQUEST_ID)
         b_seq_len.append(seq_len)
+        b_mtp_index.append(0)
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
 
-    input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cuda")
-    b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cuda")
-    b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cuda")
+    b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
+    b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
+    b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
 
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
-        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(input_ids.shape[0] - padded_req_num)
-    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(input_ids.shape[0] - padded_req_num).cuda()
+        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(b_seq_len.shape[0] - padded_req_num)
+    mem_indexes = g_infer_context.req_manager.mem_manager.alloc(b_seq_len.shape[0] - padded_req_num)
     g_infer_state_lock.release()
 
     if padded_req_num > 0:
@@ -170,36 +179,35 @@ def padded_prepare_decode_inputs(
         batch_size=b_seq_len.shape[0],
         total_token_num=total_token_num,
         max_len_in_batch=max_len_in_batch,
-        input_ids=input_ids,
-        mem_indexes=mem_indexes,
+        input_ids=None,
+        mem_indexes_cpu=mem_indexes,
         b_req_idx=b_req_idx,
+        b_mtp_index=b_mtp_index,
         b_seq_len=b_seq_len,
         is_prefill=False,
     )
     return model_input, run_reqs, padded_req_num
 
 
-def padded_overlap_prepare_decode_inputs(req_objs: List[InferReq], is_multimodal=False):
+def padded_overlap_prepare_decode_inputs(req_objs: List[InferReq]):
     split_req_bound = triton.cdiv(len(req_objs), 2)
     req_objs_0 = req_objs[0:split_req_bound]
     req_objs_1 = req_objs[split_req_bound:]
 
     enable_mtp = get_env_start_args().mtp_mode is not None
     if enable_mtp:
         micro_batch_size = max(
-            sum([len(req.mtp_gen_token_ids) + 1 for req in req_objs_0]),
-            sum([len(req.mtp_gen_token_ids) + 1 for req in req_objs_1]),
+            sum([req.mtp_step + 1 for req in req_objs_0]),
+            sum([req.mtp_step + 1 for req in req_objs_1]),
         )
     else:
         micro_batch_size = triton.cdiv(len(req_objs), 2)
 
     micro_batch_size = max(1, micro_batch_size)
 
-    micro_input, run_reqs, padded_req_num = padded_prepare_decode_inputs(
-        req_objs_0, dest_batch_size=micro_batch_size, is_multimodal=is_multimodal
-    )
+    micro_input, run_reqs, padded_req_num = padded_prepare_decode_inputs(req_objs_0, dest_batch_size=micro_batch_size)
     micro_input1, run_reqs1, padded_req_num1 = padded_prepare_decode_inputs(
-        req_objs_1, dest_batch_size=micro_batch_size, is_multimodal=is_multimodal
+        req_objs_1, dest_batch_size=micro_batch_size
     )
     return micro_input, run_reqs, padded_req_num, micro_input1, run_reqs1, padded_req_num1