(
self,
fd_config: FDConfig,
main_model: ModelForCasualLM,
local_rank: int,
device_id: int, # physical device id
target_model_inputs, # main model share inputs
)
| 88 | """ |
| 89 | |
| 90 | def __init__( |
| 91 | self, |
| 92 | fd_config: FDConfig, |
| 93 | main_model: ModelForCasualLM, |
| 94 | local_rank: int, |
| 95 | device_id: int, # physical device id |
| 96 | target_model_inputs, # main model share inputs |
| 97 | ): |
| 98 | super().__init__(fd_config) |
| 99 | self.num_main_model_layers = self.model_config.num_hidden_layers |
| 100 | self.local_rank = local_rank |
| 101 | self.device_id = device_id |
| 102 | self._update_mtp_config(main_model) |
| 103 | self._load_model() |
| 104 | self.target_model_inputs = target_model_inputs |
| 105 | self.mtp_strategy = self.speculative_config.mtp_strategy |
| 106 | self.hybrid_mode = self.mtp_strategy == "with_ngram" and self.max_draft_token_num > self.num_model_steps |
| 107 | self.enable_logprob = self.model_config.enable_logprob |
| 108 | self.enable_draft_logprob = self.speculative_config.enable_draft_logprob |
| 109 | self.cache_kvs_map = {} |
| 110 | |
| 111 | # [mixed, prefill, decoder] |
| 112 | self.role = self.scheduler_config.splitwise_role |
| 113 | self.pd_disaggregation_mode = fd_config.parallel_config.pd_disaggregation_mode |
| 114 | |
| 115 | if current_platform.is_xpu(): |
| 116 | self._propose = self._propose_xpu |
| 117 | elif current_platform.is_cuda() or current_platform.is_maca(): |
| 118 | self._propose = self._propose_cuda |
| 119 | else: |
| 120 | raise RuntimeError("Unsupported platform.") |
| 121 | |
| 122 | self.sampler = MTPSampler(fd_config) |
| 123 | self.model_inputs = ProposerInputBatch(self.fd_config, self.target_model_inputs) |
| 124 | self.model_inputs.init_share_inputs() |
| 125 | |
| 126 | # CUDA Graph |
| 127 | self.draft_model_use_cudagraph = self.graph_opt_config.draft_model_use_cudagraph |
| 128 | self.cudagraph_capture_sizes = list(reversed(self.graph_opt_config.cudagraph_capture_sizes)) |
| 129 | self.sot_warmup_sizes = self.graph_opt_config.sot_warmup_sizes |
| 130 | |
| 131 | self.attn_backends: list[AttentionBackend] = [] |
| 132 | self._initialize_attn_backend() |
| 133 | |
| 134 | # Forward meta store the global meta information of the forward |
| 135 | self.forward_meta = None |
| 136 | |
| 137 | def _update_mtp_config(self, main_model): |
| 138 | """ |
nothing calls this directly
no test coverage detected