fix input/output of the other mode

shihaobai · shihaobai · commit ab79bb3528c6 · 2025-05-27T12:58:04.000+08:00
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
@@ -76,6 +76,7 @@ def __init__(self, kvargs):
 
         # Speculative decoding
         self.spec_algo = SpeculativeDecodeAlgorithm.from_string(kvargs.get("spec_algo", "NONE"))
+        self.spec_step = kvargs.get("spec_step", 1)
 
         self._init_datatype()
         self._init_config()
diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py
@@ -130,8 +130,12 @@ def replay(self, input_ids, infer_state, input_ids1=None, infer_state1=None):
     @torch.no_grad()
     def warmup(self, model):
         logger.info("Begin capture cudagraph, use the --disable_cudagraph to disable it.")
-        decode_len = model.spec_algo.decode_len()
-        for batch_size in range(self.max_batch_size, 0, -1):
+        if model.spec_algo is not None:
+            spec_stride = model.spec_step + 1
+        else:
+            spec_stride = 1
+
+        for batch_size in range(self.max_batch_size * spec_stride, 0, -1 * spec_stride):
             # dummy prefill
             prefill_input_len = 1
             dummy_input_ids = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -111,6 +111,7 @@ def init_model(self, kvargs):
             "quant_type": kvargs.get("quant_type", None),
             "quant_cfg": kvargs.get("quant_cfg", None),
             "spec_algo": kvargs.get("spec_algo", "NONE"),
+            "spec_step": kvargs.get("spec_step", 1),
             "run_mode": self.run_mode,
         }
         self.model, self.is_multimodal = get_model(model_cfg, model_kvargs)
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_first_token_constraint_mode.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_first_token_constraint_mode.py
@@ -41,8 +41,8 @@ def decode(self):
 
         # 先 decode
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
+            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
+            logits = self.model.forward(model_input)
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
@@ -59,10 +59,10 @@ def decode(self):
         if len(decode_reqs) == 0 or (self.forward_step % self.max_wait_step == 0) or (self.need_prefill_count > 0):
             if prefill_reqs:
                 self.need_prefill_count -= 1
-                kwargs, run_reqs = prepare_prefill_inputs(
+                model_input, run_reqs = prepare_prefill_inputs(
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
-                logits = self.model.forward(**kwargs)
+                logits = self.model.forward(model_input)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_outlines_constraint_mode.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_outlines_constraint_mode.py
@@ -54,8 +54,8 @@ def decode(self):
 
         # 先 decode
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
+            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
+            logits = self.model.forward(model_input)
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
@@ -85,10 +85,10 @@ def decode(self):
         if len(decode_reqs) == 0 or (self.forward_step % self.max_wait_step == 0) or (self.need_prefill_count > 0):
             if prefill_reqs:
                 self.need_prefill_count -= 1
-                kwargs, run_reqs = prepare_prefill_inputs(
+                model_input, run_reqs = prepare_prefill_inputs(
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
-                logits = self.model.forward(**kwargs)
+                logits = self.model.forward(model_input)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_token_healing.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_token_healing.py
@@ -50,8 +50,8 @@ def decode(self):
 
         # 先 decode
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
+            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
+            logits = self.model.forward(model_input)
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
@@ -83,10 +83,10 @@ def decode(self):
         if len(decode_reqs) == 0 or (self.forward_step % self.max_wait_step == 0) or (self.need_prefill_count > 0):
             if prefill_reqs:
                 self.need_prefill_count -= 1
-                kwargs, run_reqs = prepare_prefill_inputs(
+                model_input, run_reqs = prepare_prefill_inputs(
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
-                logits = self.model.forward(**kwargs)
+                logits = self.model.forward(model_input)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_xgrammar_mode.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_xgrammar_mode.py
@@ -48,8 +48,8 @@ def decode(self):
 
         # 先 decode
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
+            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
+            logits = self.model.forward(model_input)
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
@@ -79,10 +79,10 @@ def decode(self):
         if len(decode_reqs) == 0 or (self.forward_step % self.max_wait_step == 0) or (self.need_prefill_count > 0):
             if prefill_reqs:
                 self.need_prefill_count -= 1
-                kwargs, run_reqs = prepare_prefill_inputs(
+                model_input, run_reqs = prepare_prefill_inputs(
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
-                logits = self.model.forward(**kwargs)
+                logits = self.model.forward(model_input)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )
diff --git a/lightllm/server/router/model_infer/mode_backend/continues_batch/impl_for_return_all_prompt_logprobs.py b/lightllm/server/router/model_infer/mode_backend/continues_batch/impl_for_return_all_prompt_logprobs.py
@@ -17,12 +17,14 @@ def prefill(self, run_reqs: List[Tuple]):
         req_ids = self._init_reqs(run_reqs, init_req_obj=True)
 
         req_objs = self._trans_req_ids_to_req_objs(req_ids)
-        kwargs, run_reqs = prepare_prefill_inputs(req_objs, is_chuncked_mode=False, is_multimodal=self.is_multimodal)
-
-        prompt_all_logits = self.model.forward(**kwargs)
-        input_ids = kwargs["input_ids"]
-        b_ready_cache_len = kwargs["b_ready_cache_len"]
-        b_seq_len = kwargs["b_seq_len"]
+        model_input, run_reqs = prepare_prefill_inputs(
+            req_objs, is_chuncked_mode=False, is_multimodal=self.is_multimodal
+        )
+
+        prompt_all_logits = self.model.forward(model_input)
+        input_ids = model_input.input_ids
+        b_ready_cache_len = model_input.b_ready_cache_len
+        b_seq_len = model_input.b_seq_len
         last_index = torch.cumsum(b_seq_len, dim=0, dtype=torch.long) - 1
         logits = prompt_all_logits[last_index, :]
 
diff --git a/lightllm/server/router/model_infer/mode_backend/continues_batch/impl_for_reward_model.py b/lightllm/server/router/model_infer/mode_backend/continues_batch/impl_for_reward_model.py
@@ -14,9 +14,11 @@ def prefill(self, reqs: List[Tuple]):
         req_ids = self._init_reqs(reqs, init_req_obj=True)
 
         req_objs = self._trans_req_ids_to_req_objs(req_ids)
-        kwargs, run_reqs = prepare_prefill_inputs(req_objs, is_chuncked_mode=False, is_multimodal=self.is_multimodal)
+        model_input, run_reqs = prepare_prefill_inputs(
+            req_objs, is_chuncked_mode=False, is_multimodal=self.is_multimodal
+        )
 
-        scores: torch.Tensor = self.model.forward(**kwargs)
+        scores: torch.Tensor = self.model.forward(model_input)
         scores = scores.unsqueeze(1).detach().cpu().float().numpy()
 
         next_token_id = 1
diff --git a/lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py b/lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py
@@ -72,11 +72,11 @@ def decode(self):
             self._filter_reqs(ok_finished_reqs)
 
         if prefill_reqs:
-            kwargs, run_reqs = prepare_prefill_inputs(
+            model_input, run_reqs = prepare_prefill_inputs(
                 prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
             )
 
-            logits = self.model.forward(**kwargs)
+            logits = self.model.forward(model_input)
             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()
diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py
@@ -69,10 +69,10 @@ def decode(self):
                 for req in prefill_reqs
                 if convert_sub_id_to_group_id(req.req_id) == req.req_id
             ]
-            kwargs, group_run_reqs = prepare_prefill_inputs(
+            model_input, group_run_reqs = prepare_prefill_inputs(
                 group_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
             )
-            logits = self.model.forward(**kwargs)
+            logits = self.model.forward(model_input)
 
             uninit_req_ids = [req.req_id for req in uninit_reqs]
             self._overlap_req_init_and_filter(
@@ -90,8 +90,8 @@ def decode(self):
             )
 
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
+            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
+            logits = self.model.forward(model_input)
             uninit_req_ids = [req.req_id for req in uninit_reqs]
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True

Original file line number	Diff line number	Diff line change
`@@ -111,6 +111,7 @@ def init_model(self, kvargs):`
`111`	`111`	`"quant_type": kvargs.get("quant_type", None),`
`112`	`112`	`"quant_cfg": kvargs.get("quant_cfg", None),`
`113`	`113`	`"spec_algo": kvargs.get("spec_algo", "NONE"),`
	`114`	`+ "spec_step": kvargs.get("spec_step", 1),`
`114`	`115`	`"run_mode": self.run_mode,`
`115`	`116`	`}`
`116`	`117`	`self.model, self.is_multimodal = get_model(model_cfg, model_kvargs)`