(
self,
fd_config: FDConfig,
device: str, # logic device
device_id: int, # physical device id
rank: int,
local_rank: int,
)
| 91 | |
| 92 | class MetaxModelRunner(ModelRunnerBase): |
| 93 | def __init__( |
| 94 | self, |
| 95 | fd_config: FDConfig, |
| 96 | device: str, # logic device |
| 97 | device_id: int, # physical device id |
| 98 | rank: int, |
| 99 | local_rank: int, |
| 100 | ): |
| 101 | super().__init__(fd_config=fd_config, device=device) |
| 102 | self.MAX_INFER_SEED = 9223372036854775806 |
| 103 | self.enable_mm = self.model_config.enable_mm |
| 104 | self.rank = rank |
| 105 | self.local_rank = local_rank |
| 106 | self.device_id = device_id |
| 107 | self.speculative_method = self.fd_config.speculative_config.method |
| 108 | self.speculative_decoding = self.speculative_method is not None |
| 109 | self.enable_logprob = fd_config.model_config.enable_logprob |
| 110 | self.enable_early_stop = self.fd_config.early_stop_config.enable_early_stop |
| 111 | self.is_pooling_model = self.fd_config.model_config.runner_type == "pooling" |
| 112 | self.ori_vocab_size = self.fd_config.model_config.ori_vocab_size |
| 113 | self.max_logprobs = None |
| 114 | if self.enable_logprob: |
| 115 | self.max_logprobs = ( |
| 116 | self.ori_vocab_size |
| 117 | if fd_config.model_config.max_logprobs == -1 |
| 118 | else fd_config.model_config.max_logprobs |
| 119 | ) |
| 120 | self.temp_scaled_logprobs = True |
| 121 | self.top_p_normalized_logprobs = True |
| 122 | self.prompt_logprobs_reqs: dict[str, Request] = {} |
| 123 | self.in_progress_prompt_logprobs: dict[str, LogprobsTensors] = {} |
| 124 | self.forward_batch_reqs_list: list[Request] = [None for _ in range(self.scheduler_config.max_num_seqs)] |
| 125 | self.cache_kvs_map: dict = {} |
| 126 | self.exist_prefill_flag = False |
| 127 | |
| 128 | # VL model config: |
| 129 | if self.enable_mm: |
| 130 | if "ernie" in self.fd_config.model_config.model_type: |
| 131 | self._init_image_preprocess() |
| 132 | |
| 133 | self.amp_black = [ |
| 134 | "reduce_sum", |
| 135 | "c_softmax_with_cross_entropy", |
| 136 | "elementwise_div", |
| 137 | "sin", |
| 138 | "cos", |
| 139 | "sort", |
| 140 | "multinomial", |
| 141 | ] |
| 142 | self.amp_white = [ |
| 143 | "lookup_table", |
| 144 | "lookup_table_v2", |
| 145 | "flash_attn", |
| 146 | "matmul", |
| 147 | "matmul_v2", |
| 148 | "fused_gemm_epilogue", |
| 149 | ] |
| 150 |
nothing calls this directly
no test coverage detected