fix mtp norm and fix chunked

Author: shihaobai

lightllm/server/router/model_infer/infer_batch.py

@@ -320,12 +320,13 @@ def get_chuncked_input_token_ids(self):
         chunked_end = min(self.get_cur_total_len(), chunked_start + self.shm_req.chunked_prefill_size)
         return self.shm_req.shm_prompt_ids.arr[0:chunked_end]
 
-    def get_chunked_input_token_ids_shift(self, shift=-1):
+    def get_chunked_input_token_ids_shift(self, shift=1):
         input_ids = self.get_input_token_ids()
-        shift_input_ids = np.roll(input_ids, shift)
+        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)
-        return shift_input_ids[shift:chunked_end]
+        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):
         chunked_start = self.cur_kv_len

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

@@ -13,3 +13,4 @@
 from .continues_batch.pd_mode.prefill_node_impl.prefill_impl_for_dp_chuncked import DPChunkedForPrefillNode
 from .continues_batch.pd_mode.decode_node_impl.decode_impl_for_dp import DPForDecodeNode
 from .continues_batch.impl_mtp import ContinuesBatchWithMTPBackend
+from .chunked_prefill.impl_mtp import ChunkedPrefillWithMTPBackend

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

@@ -9,7 +9,7 @@
     prepare_prefill_inputs,
 )
 from lightllm.server.router.model_infer.mode_backend.mtp_pre_process import (
-    prepare_mtp_prefill_inputs,
+    prepare_mtp_chunked_prefill_inputs,
     prepare_draft_main_model_decode_inputs,
 )
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
@@ -38,7 +38,7 @@ def decode(self):
 
         if prefill_reqs:
             model_input, run_reqs = prepare_prefill_inputs(
-                prefill_reqs, is_chuncked_mode=False, is_multimodal=self.is_multimodal
+                prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
             )
             model_output = self.model.forward(model_input)
 
@@ -49,27 +49,37 @@ def decode(self):
             next_token_ids, next_token_probs = sample(model_output.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
-            )
+            prev_step_has_output = [
+                req_obj.get_chuncked_input_token_len() == req_obj.get_cur_total_len() for req_obj in prefill_reqs
+            ]
             # spec prefill: MTP
             last_input_ids_cpu = None
             draft_model_input = model_input
             last_hidden_states = model_output.hidden_states
+            draft_next_token_ids = next_token_ids
             for draft_model_idx in range(self.spec_step):
-                device0_print(f"main {draft_model_input}")
-                draft_model_input, last_input_ids_cpu = prepare_mtp_prefill_inputs(
-                    prefill_reqs, model_input, last_hidden_states, next_token_ids, last_input_ids_cpu
+
+                draft_model_input, last_input_ids_cpu, prev_step_has_output = prepare_mtp_chunked_prefill_inputs(
+                    prefill_reqs,
+                    model_input,
+                    last_hidden_states,
+                    draft_next_token_ids,
+                    draft_model_idx + 1,
+                    prev_step_has_output,
+                    last_input_ids_cpu,
                 )
-                device0_print(f"draft_model_input {draft_model_input}")
+
                 draft_model_output = self.draft_models[draft_model_idx].forward(draft_model_input)
                 draft_next_token_ids, _ = sample(draft_model_output.logits, run_reqs, self.eos_id)
                 draft_next_token_ids = draft_next_token_ids.detach().cpu().numpy()
 
                 last_hidden_states = draft_model_output.hidden_states
-                next_token_ids = draft_next_token_ids
                 self._save_draft_token_ids(draft_next_token_ids, run_reqs, draft_model_idx)
 
+            self._post_handle(
+                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False
+            )
+
         if decode_reqs:
             model_input, run_reqs, mem_indexes_cpu = prepare_draft_main_model_decode_inputs(
                 decode_reqs, self.draft_token_id_map
@@ -93,9 +103,11 @@ def decode(self):
                 accepted_reqs,
                 next_token_ids[accepted_index],
                 next_token_logprobs[accepted_index],
-                is_chuncked_mode=False,
+                is_chuncked_mode=True,
                 do_filter_finished_reqs=False,
             )
+            self.main_step += 1
+
             # share some inference info with the main model
             draft_model_input = model_input
             draft_model_input.input_ids = next_token_ids_cuda
@@ -118,37 +130,3 @@ def decode(self):
 
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         return
-
-    def verify(self, next_token_ids, run_reqs, draft_mem_indexes):
-        accepted_reqs = []
-        accepted_index = []
-        need_free_mem_indexes = []
-        assert next_token_ids.shape[0] % self.spec_stride == 0
-
-        for i, req in enumerate(run_reqs):
-            # main model output
-            if i % self.spec_stride == 0:
-                accepted_reqs.append(req)
-                accepted_index.append(i)
-                continue
-            draft_model_idx = i % self.spec_stride - 1
-            if (
-                self.draft_token_id_map[req.req_idx][draft_model_idx] == next_token_ids[i - 1]
-                and req.cur_accepted_len == draft_model_idx
-            ):
-                accepted_reqs.append(req)
-                accepted_index.append(i)
-                req.cur_accepted_len += 1
-                device0_print(f"req {req.req_idx} accepted, cur_accepted_len {req.cur_accepted_len}")
-            else:
-                need_free_mem_indexes.append(draft_mem_indexes[i])
-        return accepted_reqs, accepted_index, need_free_mem_indexes
-
-    def _save_draft_token_ids(self, draft_next_token_ids, run_reqs, draft_model_idx):
-        batch_size = len(run_reqs) // self.spec_stride
-        for i in range(batch_size):
-            req = run_reqs[self.spec_stride * i]
-            self.draft_token_id_map[req.req_idx][draft_model_idx] = draft_next_token_ids[i + req.cur_accepted_len]
-            #  reset the cur_accepted_len
-            if draft_model_idx == self.spec_step - 1:
-                req.cur_accepted_len = 0

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

@@ -15,18 +15,49 @@ def prepare_mtp_prefill_inputs(
     for i, req in enumerate(req_objs):
         if last_input_ids_cpu is None:
             input_token_ids = req.get_input_token_ids()
+        else:
+            input_token_ids = last_input_ids_cpu[i]
+        input_token_ids = np.roll(input_token_ids, -1)
+        input_token_ids[-1] = tgt_input_ids[i]
+        input_ids.append(input_token_ids[req.cur_kv_len :])
+    input_ids_cpu = input_ids
+    input_ids = np.concatenate(input_ids, dtype=np.int64)
+    input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cuda")
+    model_input.input_ids = input_ids
+    # mtp embedding
+    model_input.hidden_states = last_hidden_states
+    return model_input, input_ids_cpu
+
+
+def prepare_mtp_chunked_prefill_inputs(
+    req_objs: List[InferReq],
+    model_input: ModelInput,
+    last_hidden_states,
+    tgt_input_ids,
+    shift,
+    prev_step_has_output,
+    last_input_ids_cpu=None,
+):
+    input_ids = []
+    for i, req in enumerate(req_objs):
+        if last_input_ids_cpu is None or not prev_step_has_output[i]:
+            input_token_ids, is_last_chunked = req.get_chunked_input_token_ids_shift(shift)
+            if prev_step_has_output[i]:
+                input_token_ids[-1] = tgt_input_ids[i]
+            prev_step_has_output[i] = is_last_chunked
         else:
             input_token_ids = last_input_ids_cpu[i]
             input_token_ids = np.roll(input_token_ids, -1)
             input_token_ids[-1] = tgt_input_ids[i]
+            prev_step_has_output[i] = True
         input_ids.append(input_token_ids[req.cur_kv_len :])
     input_ids_cpu = input_ids
     input_ids = np.concatenate(input_ids, dtype=np.int64)
     input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cuda")
     model_input.input_ids = input_ids
     # mtp embedding
     model_input.hidden_states = last_hidden_states
-    return model_input, input_ids_cpu
+    return model_input, input_ids_cpu, prev_step_has_output
 
 
 def prepare_draft_main_model_decode_inputs(req_objs: List[InferReq], draft_token_id_map):

lightllm/server/router/model_infer/model_rpc.py

@@ -22,6 +22,7 @@
     ChunckedPrefillForPrefillNode,
     DPChunkedForPrefillNode,
     ContinuesBatchWithMTPBackend,
+    ChunkedPrefillWithMTPBackend,
 )
 from lightllm.server.router.model_infer.mode_backend.redundancy_expert_manager import RedundancyExpertManager
 from lightllm.server.core.objs import RpcShmParams, RpcShmResults, ShmSyncStatusArray
@@ -161,7 +162,10 @@ def init_model(self, kvargs):
             else:
                 self.backend = ContinuesBatchBackend()
         else:
-            self.backend = ChunkedPrefillBackend()
+            if kvargs.get("spec_algo", "NONE") == "MTP":
+                self.backend = ChunkedPrefillWithMTPBackend()
+            else:
+                self.backend = ChunkedPrefillBackend()
 
         logger.info(f"use {self.backend.__class__.__name__}")
         self.backend.init_model(kvargs)