Main process for MTP inference. Args: step_use_cudagraph: bool Whether to use cuda graph. Use the target model flag to avoid hanging problems with EP.
(self, step_use_cudagraph: bool = False, is_dummy_run: bool = False)
| 801 | ) |
| 802 | |
| 803 | def _propose_cuda(self, step_use_cudagraph: bool = False, is_dummy_run: bool = False): |
| 804 | """ |
| 805 | Main process for MTP inference. |
| 806 | Args: |
| 807 | step_use_cudagraph: bool |
| 808 | Whether to use cuda graph. Use the target model flag to avoid hanging problems with EP. |
| 809 | """ |
| 810 | for substep in range(self.num_model_steps): |
| 811 | if self.model_inputs["not_need_stop"]: |
| 812 | self.model_inputs["substep"] = substep |
| 813 | # Remove padding |
| 814 | token_num_cpu = self.model_inputs["seq_lens_this_time"].numpy().sum().item() |
| 815 | ( |
| 816 | ids_remove_padding, |
| 817 | batch_id_per_token, |
| 818 | cu_seqlens_q, |
| 819 | cu_seqlens_k, |
| 820 | cu_seqlens_q_output, |
| 821 | batch_id_per_token_output, |
| 822 | ) = pre_process( |
| 823 | token_num_cpu, |
| 824 | self.model_inputs["input_ids"], |
| 825 | self.model_inputs["seq_lens_this_time"], |
| 826 | True, |
| 827 | self.model_inputs["draft_tokens"], |
| 828 | self.model_inputs["seq_lens_encoder"], |
| 829 | self.model_inputs["seq_lens_decoder"], |
| 830 | ) |
| 831 | |
| 832 | if self.enable_mm: |
| 833 | attn_mask_offsets = update_attn_mask_offsets( |
| 834 | ids_remove_padding, |
| 835 | getattr( |
| 836 | self.model_inputs, "seq_lens_this_time", self.model_inputs["seq_lens_this_time_buffer"] |
| 837 | ), |
| 838 | self.model_inputs["seq_lens_encoder"], |
| 839 | self.model_inputs["seq_lens_decoder"], |
| 840 | cu_seqlens_q, |
| 841 | self.model_inputs["attn_mask_offsets_full"], |
| 842 | self.model_inputs["attn_mask_offsets_decoder"], |
| 843 | self.model_inputs["is_block_step"], |
| 844 | self.model_inputs["decode_states"], |
| 845 | self.model_inputs["mask_rollback"], |
| 846 | ) |
| 847 | self.model_inputs["attn_mask_offsets"].copy_(attn_mask_offsets, False) |
| 848 | |
| 849 | # Initialize forward meta data |
| 850 | self.model_inputs["ids_remove_padding"].copy_(ids_remove_padding, False) |
| 851 | self.model_inputs["batch_id_per_token"][:] = -1 |
| 852 | self.model_inputs["cu_seqlens_q"].copy_(cu_seqlens_q, False) |
| 853 | self.model_inputs["cu_seqlens_k"].copy_(cu_seqlens_k, False) |
| 854 | |
| 855 | # For speculative decoding |
| 856 | self.model_inputs["cu_seqlens_q_output"].copy_(cu_seqlens_q_output, False) |
| 857 | self.model_inputs["batch_id_per_token_output"].copy_(batch_id_per_token_output, False) |
| 858 | |
| 859 | # Initialize forward meta data |
| 860 | self._initialize_forward_meta( |
nothing calls this directly
no test coverage detected