Deepseek MTP for dp backend (#923)

Author: sufubao

lightllm/common/basemodel/basemodel.py

@@ -311,127 +311,82 @@ def _decode(
         return self._token_forward(model_input.input_ids, infer_state)
 
     @torch.no_grad()
-    def microbatch_overlap_decode(self, batch: DecodeMicroBatch, batch1: DecodeMicroBatch):
-        assert batch.batch_size == batch1.batch_size
-        assert batch.mem_indexes.is_cuda
-        assert batch1.mem_indexes.is_cuda
-        input_ids, input_ids1 = batch.input_ids, batch1.input_ids
-
-        def create_inferstate(cur_batch: DecodeMicroBatch, batch_index):
-            infer_state = self.infer_state_class()
-            infer_state.is_prefill = False
-            infer_state.batch_size = cur_batch.batch_size
-            infer_state.total_token_num = cur_batch.total_token_num
-            infer_state.max_len_in_batch = cur_batch.max_len_in_batch
-            infer_state.use_dynamic_prompt_cache = self.use_dynamic_prompt_cache
-            assert cur_batch.b_req_idx.shape[0] == cur_batch.b_seq_len.shape[0]
-            infer_state.b_req_idx = cur_batch.b_req_idx
-            infer_state.b_seq_len = cur_batch.b_seq_len
-            infer_state.multimodal_params = None
-            infer_state.microbatch_index = batch_index
-
-            infer_state.mem_manager = self.mem_manager
-            infer_state.req_manager = self.req_manager
-
-            infer_state.mem_index = cur_batch.mem_indexes
-            infer_state.kv_buffer_shapedtype = (
-                (cur_batch.batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-                self.data_type,
-            )
-            infer_state.dist_group = dist_group_manager.get_group(batch_index)
-            copy_kv_index_to_req(
-                self.req_manager.req_to_token_indexs, cur_batch.b_req_idx, cur_batch.b_seq_len, infer_state.mem_index
-            )
-            return infer_state
+    def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: ModelInput):
+        assert model_input0.batch_size == model_input1.batch_size
+        assert model_input0.mem_indexes.is_cuda
+        assert model_input1.mem_indexes.is_cuda
+        input_ids0, input_ids1 = model_input0.input_ids, model_input1.input_ids
 
-        infer_state = create_inferstate(batch, 0)
-        infer_state1 = create_inferstate(batch1, 1)
+        infer_state0 = self._create_inferstate(model_input0, 0)
+        copy_kv_index_to_req(
+            self.req_manager.req_to_token_indexs, model_input0.b_req_idx, model_input0.b_seq_len, infer_state0.mem_index
+        )
+        infer_state0.init_some_extra_state(self, input_ids0)
 
-        infer_state.init_some_extra_state(self, input_ids)
+        infer_state1 = self._create_inferstate(model_input1, 1)
+        copy_kv_index_to_req(
+            self.req_manager.req_to_token_indexs, model_input1.b_req_idx, model_input1.b_seq_len, infer_state1.mem_index
+        )
         infer_state1.init_some_extra_state(self, input_ids1)
 
-        batch_size = batch.batch_size
-        max_len_in_batch = max(batch.max_len_in_batch, batch1.max_len_in_batch)
+        batch_size = model_input0.batch_size
+        max_len_in_batch = max(model_input0.max_len_in_batch, model_input1.max_len_in_batch)
 
         if self.graph is not None and self.graph.can_run(batch_size, max_len_in_batch):
             if self.graph.need_capture(batch_size):
-                infer_state.is_cuda_graph = True
+                infer_state0.is_cuda_graph = True
                 infer_state1.is_cuda_graph = True
 
-                predict_logits, predict_logits1 = self.graph.capture_decode(
+                model_output0, model_output1 = self.graph.capture_decode(
                     self._overlap_tpsp_token_forward,
-                    input_ids,
-                    infer_state,
+                    input_ids0,
+                    infer_state0,
                     input_ids1=input_ids1,
                     infer_state1=infer_state1,
                 )
             else:
-                predict_logits, predict_logits1 = self.graph.replay(
-                    input_ids, infer_state, input_ids1=input_ids1, infer_state1=infer_state1
+                model_output0, model_output1 = self.graph.replay(
+                    input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
                 )
         else:
-            predict_logits, predict_logits1 = self._overlap_tpsp_token_forward(
-                input_ids, infer_state, input_ids1=input_ids1, infer_state1=infer_state1
+            model_output0, model_output1 = self._overlap_tpsp_token_forward(
+                input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
             )
-        return predict_logits, predict_logits1
+        return model_output0, model_output1
 
     @torch.no_grad()
-    def microbatch_overlap_prefill(self, batch: PrefillMicroBatch, batch1: PrefillMicroBatch):
-        assert batch.mem_indexes.is_cuda
-        assert batch1.mem_indexes.is_cuda
-        input_ids, input_ids1 = batch.input_ids, batch1.input_ids
-
-        def create_inferstate(cur_batch: PrefillMicroBatch, batch_index):
-            infer_state = self.infer_state_class()
-            infer_state.is_prefill = True
-            infer_state.is_token_healing = self.is_token_healing
-            infer_state.return_all_prompt_logics = self.return_all_prompt_logics
-            infer_state.use_dynamic_prompt_cache = self.use_dynamic_prompt_cache
-            infer_state.batch_size = cur_batch.batch_size
-            infer_state.total_token_num = cur_batch.total_token_num
-            infer_state.max_len_in_batch = cur_batch.max_len_in_batch
-            assert cur_batch.b_req_idx.shape[0] == cur_batch.b_seq_len.shape[0]
-            infer_state.b_req_idx = cur_batch.b_req_idx
-            infer_state.b_seq_len = cur_batch.b_seq_len
-            if cur_batch.b_ready_cache_len is not None:
-                infer_state.b_ready_cache_len = cur_batch.b_ready_cache_len
-            else:
-                infer_state.b_ready_cache_len = torch.zeros_like(
-                    cur_batch.b_seq_len, dtype=cur_batch.b_seq_len.dtype, device=cur_batch.b_seq_len.device
-                )
-            infer_state.multimodal_params = cur_batch.multimodal_params
-            infer_state.microbatch_index = batch_index
+    def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: ModelInput):
+        assert model_input0.mem_indexes.is_cuda
+        assert model_input1.mem_indexes.is_cuda
+        input_ids0, input_ids1 = model_input0.input_ids, model_input1.input_ids
 
-            infer_state.mem_manager = self.mem_manager
-            infer_state.req_manager = self.req_manager
-
-            infer_state.mem_index = cur_batch.mem_indexes
-            infer_state.kv_buffer_shapedtype = (
-                (cur_batch.input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-                self.data_type,
-            )
-            infer_state.dist_group = dist_group_manager.get_group(batch_index)
-            init_req_to_token_indexes(
-                self.req_manager.req_to_token_indexs,
-                cur_batch.b_req_idx,
-                cur_batch.b_seq_len,
-                infer_state.b_ready_cache_len,
-                cur_batch.max_len_in_batch,
-                infer_state.mem_index,
-            )
-            return infer_state
-
-        infer_state = create_inferstate(batch, 0)
-        infer_state1 = create_inferstate(batch1, 1)
-
-        infer_state.init_some_extra_state(self, input_ids)
+        infer_state0 = self._create_inferstate(model_input0, 0)
+        init_req_to_token_indexes(
+            self.req_manager.req_to_token_indexs,
+            model_input0.b_req_idx,
+            model_input0.b_seq_len,
+            infer_state0.b_ready_cache_len,
+            model_input0.max_len_in_batch,
+            infer_state0.mem_index,
+        )
+        infer_state0.init_some_extra_state(self, input_ids0)
+        
+        infer_state1 = self._create_inferstate(model_input1, 1)
+        init_req_to_token_indexes(
+            self.req_manager.req_to_token_indexs,
+            model_input1.b_req_idx,
+            model_input1.b_seq_len,
+            infer_state1.b_ready_cache_len,
+            model_input1.max_len_in_batch,
+            infer_state1.mem_index,
+        )
         infer_state1.init_some_extra_state(self, input_ids1)
 
-        predict_logits, predict_logits1 = self._overlap_tpsp_context_forward(
-            input_ids, infer_state, input_ids1=input_ids1, infer_state1=infer_state1
+        model_output0, model_output1 = self._overlap_tpsp_context_forward(
+            input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
         )
         dist_group_manager.clear_deepep_buffer()
-        return predict_logits, predict_logits1
+        return model_output0, model_output1
 
     @final
     def _context_forward(self, input_ids, infer_state: InferStateInfo):
@@ -508,9 +463,21 @@ def _overlap_tpsp_token_forward(
         predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward(
             input_embs, input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
-
+        
         g_cache_manager.cache_env_out()
-        return predict_logits, predict_logits1
+        is_return_hidden_states = self.spec_algo.is_mtp() or (
+            self.spec_algo.is_mtp_module() and not self.last_mtp_module
+        )
+        model_output = ModelOutput(
+            logits=predict_logits,
+            hidden_states=input_embs if is_return_hidden_states else None,
+        )
+        
+        model_output1 = ModelOutput(
+            logits=predict_logits1,
+            hidden_states=input_embs1 if is_return_hidden_states else None,
+        )
+        return model_output, model_output1
 
     @final
     def _overlap_tpsp_context_forward(
@@ -528,7 +495,21 @@ def _overlap_tpsp_context_forward(
             input_embs, input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
         g_cache_manager.cache_env_out()
-        return predict_logits, predict_logits1
+        
+        is_return_hidden_states = self.spec_algo.is_mtp() or (
+            self.spec_algo.is_mtp_module() and not self.last_mtp_module
+        )
+        model_output = ModelOutput(
+            logits=predict_logits,
+            hidden_states=input_embs if is_return_hidden_states else None,
+        )
+        
+        model_output1 = ModelOutput(
+            logits=predict_logits1,
+            hidden_states=input_embs1 if is_return_hidden_states else None,
+        )
+        
+        return model_output, model_output1
 
     @final
     @torch.no_grad()

lightllm/models/deepseek_mtp/deepseek3_mtp_mem_manager.py

@@ -26,11 +26,6 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
 
         self.can_use_mem_size = self.size
 
-        rank_in_node = get_current_rank_in_node()
-        self.shared_can_use_token_num = SharedInt(f"MTP_mem_manger_can_use_token_num_{rank_in_node}")
-
-        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
-
         self._init_buffers(
             self.size,
             dtype,

lightllm/server/router/model_infer/infer_batch.py

@@ -325,7 +325,7 @@ def get_chunked_input_token_ids_shift(self, shift=1):
         shift_input_ids = np.roll(input_ids, -1 * shift)
         chunked_start = self.cur_kv_len
         chunked_end = min(self.get_cur_total_len(), chunked_start + self.shm_req.chunked_prefill_size)
-        is_last_chunked = chunked_end == self.get_cur_total_len() + shift
+        is_last_chunked = chunked_end == self.get_cur_total_len() - shift
         return shift_input_ids[0:chunked_end], is_last_chunked
 
     def get_chuncked_input_token_len(self):

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

@@ -7,6 +7,7 @@
 from .chunked_prefill.impl_for_outlines_constraint_mode import OutlinesConstraintBackend
 from .chunked_prefill.impl_for_first_token_constraint_mode import FirstTokenConstraintBackend
 from .dp_backend.impl import DPChunkedPrefillBackend
+from .dp_backend.impl_mtp import DPChunkedPrefillWithMTPBackend
 from .continues_batch.pd_mode.prefill_node_impl.prefill_impl import ChunckedPrefillForPrefillNode
 from .continues_batch.pd_mode.decode_node_impl.decode_impl import ContinuesBatchBackendForDecodeNode
 from .chunked_prefill.impl_for_xgrammar_mode import XgrammarBackend

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

@@ -11,6 +11,11 @@
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from lightllm.utils.envs_utils import get_env_start_args
 
+from lightllm.server.router.model_infer.mode_backend.generic_pre_process import (
+    prepare_prefill_inputs,
+    prepare_decode_inputs
+)
+from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
 
 class DPChunkedPrefillBackend(ModeBackend):
     def __init__(self) -> None:
@@ -24,16 +29,14 @@ def __init__(self) -> None:
         pass
 
     def init_custom(self):
-        self.reduce_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
         # 这个地方预先进行一次 prefill 推理，主要是为了填充后续fake请求的第一个token位置，因为填充的decode请求
         # 在推理的时候至少是两个token，1个是已经有kv的token，一个是等待计算kv的token，然后生成第三个token，这几个
         # token 实际引用的都是 g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX，但是需要初始化排除
         # nan 值，避免后续构建的fake请求在计算的过程中出现计算错误。
-        from .pre_process import padded_prepare_prefill_inputs
-
-        kwargs, run_reqs, padded_req_num = padded_prepare_prefill_inputs([], 1, is_multimodal=self.is_multimodal)
-        self.model.forward(**kwargs)
-        assert len(run_reqs) == 0 and padded_req_num == 1
+        self.reduce_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
+        model_input, run_reqs = prepare_prefill_inputs([], is_chuncked_mode=True, is_multimodal=self.is_multimodal, pad_for_empty_batch=True)
+        self.model.forward(model_input)
+        assert len(run_reqs) == 0 and model_input.batch_size == 1
         return
 
     def prefill(self, reqs: List[Tuple]):
@@ -71,15 +74,14 @@ def decode(self):
         return
 
     def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs):
-        from .pre_process import padded_prepare_prefill_inputs
-
-        kwargs, run_reqs, padded_req_num = padded_prepare_prefill_inputs(
-            prefill_reqs, max_prefill_num, is_multimodal=self.is_multimodal
+        model_input, run_reqs = prepare_prefill_inputs(
+            prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal, pad_for_empty_batch=True
         )
-        logits = self.model.forward(**kwargs)
+        model_output: ModelOutput = self.model.forward(model_input)
+        
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         if len(run_reqs) != 0:
-            logits = logits[0 : len(run_reqs), :]
+            logits = model_output.logits[0 : len(run_reqs), :]
             next_token_ids, next_token_probs = sample(logits, run_reqs, self.eos_id)
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
@@ -89,43 +91,31 @@ def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int
         return
 
     def normal_decode(self, decode_reqs: List[InferReq], max_decode_num: int, uninit_reqs, ok_finished_reqs):
-        from .pre_process import padded_prepare_decode_inputs
-
-        kwargs, run_reqs, padded_req_num = padded_prepare_decode_inputs(
-            decode_reqs, max_decode_num, is_multimodal=self.is_multimodal
-        )
-        logits = self.model.forward(**kwargs)
+        model_input, run_reqs = prepare_decode_inputs(decode_reqs, pad_for_empty_batch=True)
+        model_output: ModelOutput = self.model.forward(model_input)
 
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
 
         if len(run_reqs) != 0:
-            logits = logits[0 : len(run_reqs), :]
+            logits = model_output.logits[0 : len(run_reqs), :]
             next_token_ids, next_token_probs = sample(logits, run_reqs, self.eos_id)
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
             self._post_handle(
                 run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=False, do_filter_finished_reqs=False
             )
-        logits = None
 
     def overlap_decode(self, decode_reqs: List[InferReq], max_decode_num: int, uninit_reqs, ok_finished_reqs):
         from .pre_process import padded_overlap_prepare_decode_inputs
 
-        (
-            micro_batch,
-            run_reqs,
-            padded_req_num,
-            micro_batch1,
-            run_reqs1,
-            padded_req_num1,
-        ) = padded_overlap_prepare_decode_inputs(decode_reqs, max_decode_num, is_multimodal=self.is_multimodal)
-        logits, logits1 = self.model.microbatch_overlap_decode(micro_batch, micro_batch1)
+        micro_input, run_reqs, micro_input1, run_reqs1 = padded_overlap_prepare_decode_inputs(decode_reqs, max_decode_num, is_multimodal=self.is_multimodal)
+        micro_output, micro_output1 = self.model.microbatch_overlap_decode(micro_input, micro_input1)
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         req_num, req_num1 = len(run_reqs), len(run_reqs1)
-        all_logits = torch.empty((req_num + req_num1, logits.shape[1]), dtype=logits.dtype, device=logits.device)
+        all_logits = torch.empty((req_num + req_num1, micro_output.logits.shape[1]), dtype=micro_output.logits.dtype, device=micro_output.logits.device)
 
-        all_logits[0:req_num, :].copy_(logits[0:req_num, :], non_blocking=True)
-        all_logits[req_num : (req_num + req_num1), :].copy_(logits1[0:req_num1, :], non_blocking=True)
+        all_logits[0:req_num, :].copy_(micro_output.logits[0:req_num, :], non_blocking=True)
+        all_logits[req_num : (req_num + req_num1), :].copy_(micro_output1.logits[0:req_num1, :], non_blocking=True)
 
         all_run_reqs = run_reqs + run_reqs1
         if all_run_reqs:
@@ -140,21 +130,14 @@ def overlap_decode(self, decode_reqs: List[InferReq], max_decode_num: int, unini
     def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs):
         from .pre_process import padded_overlap_prepare_prefill_inputs
 
-        (
-            micro_batch,
-            run_reqs,
-            padded_req_num,
-            micro_batch1,
-            run_reqs1,
-            padded_req_num1,
-        ) = padded_overlap_prepare_prefill_inputs(prefill_reqs, max_prefill_num, is_multimodal=self.is_multimodal)
-        logits, logits1 = self.model.microbatch_overlap_prefill(micro_batch, micro_batch1)
+        micro_input, run_reqs, micro_input1, run_reqs1 = padded_overlap_prepare_prefill_inputs(prefill_reqs, max_prefill_num, is_multimodal=self.is_multimodal)
+        micro_output, micro_output1 = self.model.microbatch_overlap_prefill(micro_input, micro_input1)
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         req_num, req_num1 = len(run_reqs), len(run_reqs1)
-        all_logits = torch.empty((req_num + req_num1, logits.shape[1]), dtype=logits.dtype, device=logits.device)
+        all_logits = torch.empty((req_num + req_num1, micro_output.logits.shape[1]), dtype=micro_output.logits.dtype, device=micro_output.logits.device)
 
-        all_logits[0:req_num, :].copy_(logits[0:req_num, :], non_blocking=True)
-        all_logits[req_num : (req_num + req_num1), :].copy_(logits1[0:req_num1, :], non_blocking=True)
+        all_logits[0:req_num, :].copy_(micro_output.logits[0:req_num, :], non_blocking=True)
+        all_logits[req_num : (req_num + req_num1), :].copy_(micro_output1.logits[0:req_num1, :], non_blocking=True)
 
         all_run_reqs = run_reqs + run_reqs1
         if all_run_reqs: