The Entrance of model execute. Args: model_forward_batch: 'Request' contains information related to prompt and is an abstract class at the server level, which is too granular for ModelRunner. We plan to replace it with 'ModelForwardBatch'.
(
self,
model_forward_batch: Optional[List[Request]] = None,
num_running_requests: int = None,
)
| 956 | return prefill_done_idxs |
| 957 | |
| 958 | def execute_model( |
| 959 | self, |
| 960 | model_forward_batch: Optional[List[Request]] = None, |
| 961 | num_running_requests: int = None, |
| 962 | ) -> Optional[ModelRunnerOutput]: |
| 963 | """ |
| 964 | The Entrance of model execute. |
| 965 | Args: |
| 966 | model_forward_batch: 'Request' contains information related to prompt and is an abstract |
| 967 | class at the server level, which is too granular for ModelRunner. |
| 968 | We plan to replace it with 'ModelForwardBatch'. |
| 969 | num_running_requests: batch_size |
| 970 | intermediate_tensors: |
| 971 | """ |
| 972 | # If `not_need_stop`` is False, it means the current worker is in an idle state. |
| 973 | # This logic is not used in TP (Tensor Parallelism) mode. However, in EP (Expert Parallelism) mode, |
| 974 | # when there is data on other runner, the current runner is required to execute part of the model. |
| 975 | if not self.not_need_stop(): |
| 976 | self._execute_empty_input(self.forward_meta) |
| 977 | return None |
| 978 | |
| 979 | # 1. Prepare inputs of model and sampler. |
| 980 | p_done_idxs = self._get_p_done_idxs_gd(model_forward_batch, num_running_requests) |
| 981 | self._prepare_inputs() |
| 982 | self.sampler.pre_process(p_done_idxs) |
| 983 | |
| 984 | # 2. Padding inputs for cuda graph |
| 985 | |
| 986 | # 3. Execute model |
| 987 | model_output = self.model( |
| 988 | ids_remove_padding=self.share_inputs["ids_remove_padding"], |
| 989 | forward_meta=self.forward_meta, |
| 990 | ) |
| 991 | |
| 992 | hidden_states = rebuild_padding( |
| 993 | model_output, |
| 994 | self.share_inputs["cu_seqlens_q"], |
| 995 | self.share_inputs["seq_lens_this_time"], |
| 996 | self.share_inputs["seq_lens_decoder"], |
| 997 | self.share_inputs["seq_lens_encoder"], |
| 998 | (self.share_inputs["output_padding_offset"] if self.speculative_decoding else None), |
| 999 | self.model_config.max_model_len, |
| 1000 | ) |
| 1001 | |
| 1002 | # 4. Compute logits, Sample |
| 1003 | logits = self.model.compute_logits(hidden_states) |
| 1004 | |
| 1005 | if not self.speculative_decoding: |
| 1006 | set_value_by_flags_and_idx( |
| 1007 | self.share_inputs["pre_ids"], |
| 1008 | self.share_inputs["input_ids"], |
| 1009 | self.share_inputs["seq_lens_this_time"], |
| 1010 | self.share_inputs["seq_lens_encoder"], |
| 1011 | self.share_inputs["seq_lens_decoder"], |
| 1012 | self.share_inputs["step_idx"], |
| 1013 | self.share_inputs["stop_flags"], |
| 1014 | ) |
| 1015 | sampler_output = self.sampler( |
nothing calls this directly
no test coverage detected