fix

Author: shihaobai

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

@@ -253,7 +253,8 @@ def decode_mtp(
 
         # 处理需要释放的内存索引
         need_free_mem_indexes = model_input.mem_indexes_cpu[verify_info["accepted_index_cpu"] == 0]
-        need_free_mem_indexes = torch.cat([need_free_mem_indexes, additional_mem_indexes_cpu], dim=0)
+        if additional_mem_indexes_cpu is not None:
+            need_free_mem_indexes = torch.cat([need_free_mem_indexes, additional_mem_indexes_cpu], dim=0)
 
         self._update_mtp_accept_ratio(decode_reqs=decode_reqs, mtp_accept_len_cpu=verify_info["mtp_accept_len_cpu"])
         select_mask = torch.tensor(verify_info["accepted_index_cpu"], dtype=torch.bool, device="cpu")

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

@@ -37,7 +37,7 @@ def __init__(self) -> None:
         # 在 mtp 模式下切换绑定的prefill 和 decode 函数
         if get_env_start_args().mtp_mode:
             self.is_mtp_eagle = get_env_start_args().mtp_mode == "deepseekv3_eagle"
-            self.prefill_mtp_step = 1 if self.is_mtp_eagle else get_env_start_args().mtp_step
+            self.num_mtp_models = 1 if self.is_mtp_eagle else get_env_start_args().mtp_step
             if self.enable_prefill_microbatch_overlap:
                 self.prefill = self.prefill_overlap_mtp
             else:
@@ -360,7 +360,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq]
             self._draft_prefill_forward(
                 model_input=model_input,
                 model_output=model_output,
-                mtp_step=self.prefill_mtp_step,
+                mtp_step=self.num_mtp_models,
                 next_token_ids=draft_next_token_ids_gpu,
             )
             sync_event = torch.cuda.Event()
@@ -596,7 +596,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I
 
             draft_model_output0, draft_model_output1 = model_output0, model_output1
 
-            for draft_model_idx in range(self.prefill_mtp_step):
+            for draft_model_idx in range(self.num_mtp_models):
 
                 draft_model_input0 = prepare_mtp_prefill_inputs(
                     model_input=draft_model_input0,

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

@@ -133,6 +133,8 @@ def padded_prepare_decode_inputs(
     b_req_idx = []
     b_mtp_index = []
     b_seq_len = []
+    max_q_seq_len = 0
+    max_kv_seq_len = 0
     for req in req_objs:
         run_reqs.append(req)
         b_req_idx.append(req.req_idx)
@@ -141,6 +143,8 @@ def padded_prepare_decode_inputs(
         b_seq_len.append(seq_len)
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
+        max_q_seq_len = max(max_q_seq_len, req.mtp_step + 1)
+        max_kv_seq_len = max(max_kv_seq_len, seq_len)
         b_mtp_index.append(0)
         # process the draft tokens.
         for step in range(req.mtp_step):
@@ -150,6 +154,7 @@ def padded_prepare_decode_inputs(
             b_seq_len.append(seq_len)
             total_token_num += seq_len
             max_len_in_batch = max(max_len_in_batch, seq_len)
+            max_kv_seq_len = max(max_kv_seq_len, seq_len)
             b_mtp_index.append(step + 1)
 
     if dest_batch_size is None:
@@ -170,6 +175,8 @@ def padded_prepare_decode_inputs(
         b_mtp_index.append(0)
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
+        max_q_seq_len = max(max_q_seq_len, 1)
+        max_kv_seq_len = max(max_kv_seq_len, seq_len)
 
     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")
@@ -194,6 +201,8 @@ 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,
+        max_q_seq_len=max_q_seq_len,
+        max_kv_seq_len=max_kv_seq_len,
         input_ids=None,
         mem_indexes_cpu=mem_indexes,
         b_req_idx=b_req_idx,