PostProcess for generation
(self, sampled_token_ids)
| 744 | self.model_inputs["target_hidden_states"].copy_(target_hidden_states, False) |
| 745 | |
| 746 | def _post_process(self, sampled_token_ids): |
| 747 | """ |
| 748 | PostProcess for generation |
| 749 | """ |
| 750 | draft_model_update( |
| 751 | sampled_token_ids, |
| 752 | self.model_inputs["draft_tokens"], |
| 753 | self.model_inputs["pre_ids"], |
| 754 | self.model_inputs["seq_lens_this_time"], |
| 755 | self.model_inputs["seq_lens_encoder"], |
| 756 | self.model_inputs["seq_lens_decoder"], |
| 757 | self.model_inputs["step_idx"], |
| 758 | # Note(ZKK): |
| 759 | # I strongly advise xpu student delete the fuck `output_cum_offsets` name in XPU backend |
| 760 | # like my pr https://github.com/PaddlePaddle/FastDeploy/pull/6358 |
| 761 | ( |
| 762 | self.model_inputs["cu_seqlens_q_output"] |
| 763 | if current_platform.is_cuda() |
| 764 | else self.model_inputs["output_cum_offsets"] |
| 765 | ), |
| 766 | self.model_inputs["stop_flags"], |
| 767 | self.model_inputs["not_need_stop"], |
| 768 | self.model_inputs["max_dec_len"], |
| 769 | self.model_inputs["eos_token_id"], |
| 770 | self.model_inputs["base_model_draft_tokens"], |
| 771 | self.max_model_len, |
| 772 | self.model_inputs["substep"], |
| 773 | ) |
| 774 | |
| 775 | if self.role == "prefill" and self.parallel_config.tensor_parallel_rank == 0: |
| 776 | skip_save = bool(int(envs.ENABLE_V1_KVCACHE_SCHEDULER)) |
| 777 | recover_model_output_map = recover_batch_index_for_output( |
| 778 | self.model_inputs, |
| 779 | self.model_inputs.index_to_batch_id, |
| 780 | self.model_inputs.enable_pd_reorder, |
| 781 | ["base_model_draft_tokens", "seq_lens_decoder", "prompt_lens", "step_idx"], |
| 782 | ) |
| 783 | mtp_save_first_token( |
| 784 | recover_model_output_map["base_model_draft_tokens"], |
| 785 | self.model_inputs["not_need_stop"], |
| 786 | recover_model_output_map["seq_lens_decoder"], |
| 787 | recover_model_output_map["prompt_lens"], |
| 788 | recover_model_output_map["step_idx"], |
| 789 | self.local_rank, |
| 790 | self.parallel_config.use_ep, |
| 791 | skip_save, |
| 792 | ) |
| 793 | # Ensure only save first token once. |
| 794 | paddle.assign( |
| 795 | paddle.where( |
| 796 | self.model_inputs["stop_flags"], |
| 797 | paddle.zeros_like(self.model_inputs["step_idx"]), |
| 798 | self.model_inputs["step_idx"], |
| 799 | ), |
| 800 | self.model_inputs["step_idx"], |
| 801 | ) |
| 802 | |
| 803 | def _propose_cuda(self, step_use_cudagraph: bool = False, is_dummy_run: bool = False): |