(
self,
fd_config: FDConfig,
device: str, # logic device
device_id: int, # physical device id
rank: int,
local_rank: int,
)
| 112 | |
| 113 | class GPUModelRunner(ModelRunnerBase): |
| 114 | def __init__( |
| 115 | self, |
| 116 | fd_config: FDConfig, |
| 117 | device: str, # logic device |
| 118 | device_id: int, # physical device id |
| 119 | rank: int, |
| 120 | local_rank: int, |
| 121 | ): |
| 122 | super().__init__(fd_config=fd_config, device=device) |
| 123 | self.MAX_INFER_SEED = 9223372036854775806 |
| 124 | self.enable_mm = self.model_config.enable_mm |
| 125 | self.rank = rank |
| 126 | self.local_rank = local_rank |
| 127 | self.device_id = device_id |
| 128 | self.speculative_method = self.fd_config.speculative_config.method |
| 129 | self.speculative_decoding = self.speculative_method is not None |
| 130 | self.enable_logprob = fd_config.model_config.enable_logprob |
| 131 | self.enable_keep_sampling_mask = fd_config.model_config.enable_keep_sampling_mask |
| 132 | self.enable_early_stop = self.fd_config.early_stop_config.enable_early_stop |
| 133 | self.is_pooling_model = self.fd_config.model_config.runner_type == "pooling" |
| 134 | self.ori_vocab_size = self.fd_config.model_config.ori_vocab_size |
| 135 | self.max_logprobs = None |
| 136 | if self.enable_logprob: |
| 137 | self.max_logprobs = ( |
| 138 | self.ori_vocab_size |
| 139 | if fd_config.model_config.max_logprobs == -1 |
| 140 | else fd_config.model_config.max_logprobs |
| 141 | ) |
| 142 | self.temp_scaled_logprobs = True |
| 143 | self.top_p_normalized_logprobs = True |
| 144 | self.prompt_logprobs_reqs: dict[str, Request] = {} |
| 145 | self.in_progress_prompt_logprobs: dict[str, LogprobsTensors] = {} |
| 146 | self.forward_batch_reqs_list: list[Request] = [None for _ in range(self.scheduler_config.max_num_seqs)] |
| 147 | self.cache_kvs_map: dict = {} |
| 148 | self.exist_prefill_flag = False |
| 149 | |
| 150 | self.is_kvcache_sleeping = False |
| 151 | self.is_weight_sleeping = False |
| 152 | |
| 153 | # VL model config: |
| 154 | if self.enable_mm: |
| 155 | if "ernie" in self.fd_config.model_config.model_type: |
| 156 | self._init_image_preprocess() |
| 157 | |
| 158 | self.amp_black = [ |
| 159 | "reduce_sum", |
| 160 | "c_softmax_with_cross_entropy", |
| 161 | "elementwise_div", |
| 162 | "sin", |
| 163 | "cos", |
| 164 | "sort", |
| 165 | "multinomial", |
| 166 | ] |
| 167 | self.amp_white = [ |
| 168 | "lookup_table", |
| 169 | "lookup_table_v2", |
| 170 | "flash_attn", |
| 171 | "matmul", |
nothing calls this directly
no test coverage detected