| 657 | |
| 658 | |
| 659 | class ProposerInputBatch(InputBatch): |
| 660 | def __init__(self, fd_config: FDConfig, target_model_input_batch: InputBatch) -> None: |
| 661 | self.enable_mm = fd_config.model_config.enable_mm |
| 662 | self.num_model_steps = fd_config.speculative_config.num_model_steps |
| 663 | self.index_to_batch_id = {} |
| 664 | self.target_model_input_batch = target_model_input_batch |
| 665 | self.fd_config: FDConfig = fd_config |
| 666 | self.scheduler_config = fd_config.scheduler_config |
| 667 | self.model_config: ModelConfig = fd_config.model_config |
| 668 | self.cache_config: CacheConfig = fd_config.cache_config |
| 669 | self.speculative_config: SpeculativeConfig = fd_config.speculative_config |
| 670 | self.enable_pd_reorder: bool = False |
| 671 | |
| 672 | def init_share_inputs(self): |
| 673 | # share with targe model |
| 674 | self.enable_pd_reorder = getattr(self.target_model_input_batch, "enable_pd_reorder", False) |
| 675 | |
| 676 | self.block_tables = paddle.clone(self.target_model_input_batch["block_tables"]) |
| 677 | self.input_ids = paddle.clone(self.target_model_input_batch["input_ids"]) |
| 678 | self.input_ids_cpu = paddle.full( |
| 679 | shape=[self.scheduler_config.max_num_seqs, self.model_config.max_model_len], |
| 680 | fill_value=-1, |
| 681 | dtype="int64", |
| 682 | ).cpu() |
| 683 | self.seq_lens_this_time_buffer = paddle.clone(self.target_model_input_batch["seq_lens_this_time"]) |
| 684 | |
| 685 | self.seq_lens_encoder = paddle.clone(self.target_model_input_batch["seq_lens_encoder"]) |
| 686 | self.seq_lens_decoder = paddle.clone(self.target_model_input_batch["seq_lens_decoder"]) |
| 687 | self.step_idx = paddle.clone(self.target_model_input_batch["step_idx"]) |
| 688 | self.stop_flags = paddle.clone(self.target_model_input_batch["stop_flags"]) |
| 689 | self.not_need_stop = paddle.to_tensor([False], dtype="bool", place="cpu") |
| 690 | self.pre_ids = paddle.clone(self.target_model_input_batch["pre_ids"]) |
| 691 | if current_platform.is_cuda(): |
| 692 | self.cu_seqlens_q_output = paddle.clone(self.target_model_input_batch["cu_seqlens_q_output"]) |
| 693 | self.batch_id_per_token_output = paddle.clone(self.target_model_input_batch["batch_id_per_token_output"]) |
| 694 | else: |
| 695 | self.output_cum_offsets = paddle.clone(self.target_model_input_batch["output_cum_offsets"]) |
| 696 | self.output_padding_offset = paddle.clone(self.target_model_input_batch["output_padding_offset"]) |
| 697 | self.ids_remove_padding = paddle.clone(self.target_model_input_batch["ids_remove_padding"]) |
| 698 | self.batch_id_per_token = paddle.clone(self.target_model_input_batch["batch_id_per_token"]) |
| 699 | self.cu_seqlens_q = paddle.clone(self.target_model_input_batch["cu_seqlens_q"]) |
| 700 | self.cu_seqlens_k = paddle.clone(self.target_model_input_batch["cu_seqlens_k"]) |
| 701 | |
| 702 | self.target_hidden_states = paddle.full( |
| 703 | [ |
| 704 | self.scheduler_config.max_num_batched_tokens + self.scheduler_config.max_extra_num_batched_tokens, |
| 705 | self.model_config.hidden_size, |
| 706 | ], |
| 707 | 0, |
| 708 | dtype="bfloat16", |
| 709 | ) |
| 710 | |
| 711 | tmp_position_ids = paddle.arange(self.model_config.max_model_len).reshape((1, -1)) |
| 712 | |
| 713 | self.rope_emb = get_rope( |
| 714 | rotary_dim=self.model_config.head_dim, |
| 715 | position_ids=tmp_position_ids, |
| 716 | base=self.model_config.rope_theta, |