Set dummy prefill inputs to model_inputs
(self, num_tokens: int, batch_size: int, expected_decode_len: int)
| 160 | self.model = model_loader.load_model(fd_config=self.fd_config) |
| 161 | |
| 162 | def dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decode_len: int): |
| 163 | """Set dummy prefill inputs to model_inputs""" |
| 164 | max_dec_len = expected_decode_len + 1 |
| 165 | |
| 166 | input_length = min( |
| 167 | num_tokens // batch_size, |
| 168 | self.model_config.max_model_len - max_dec_len, |
| 169 | ) |
| 170 | |
| 171 | # TODO(wanglongzhi): Figure out the accurate buffer size of DeepEP. |
| 172 | if self.fd_config.parallel_config.enable_expert_parallel: |
| 173 | input_length = min(input_length, 32) |
| 174 | |
| 175 | block_num = ( |
| 176 | input_length + self.cache_config.block_size - 1 |
| 177 | ) // self.cache_config.block_size + self.cache_config.enc_dec_block_num |
| 178 | |
| 179 | for i in range(batch_size): |
| 180 | idx = i |
| 181 | self.model_inputs["input_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length) |
| 182 | self.model_inputs["eos_token_id"][:] = np.array([2], dtype="int64").reshape(-1, 1) |
| 183 | self.model_inputs["seq_lens_this_time_buffer"][idx : idx + 1] = input_length |
| 184 | self.model_inputs["seq_lens_encoder"][idx : idx + 1] = input_length |
| 185 | self.model_inputs["seq_lens_decoder"][idx : idx + 1] = 0 |
| 186 | self.model_inputs["step_idx"][idx : idx + 1] = 0 |
| 187 | self.model_inputs["max_dec_len"][idx : idx + 1] = max_dec_len |
| 188 | self.model_inputs["stop_flags"][idx : idx + 1] = False |
| 189 | |
| 190 | self.model_inputs["encoder_block_lens"][idx : idx + 1] = block_num |
| 191 | self.model_inputs["block_tables"][idx : idx + 1, :block_num] = np.arange( |
| 192 | idx * block_num, (idx + 1) * block_num, 1 |
| 193 | ) |
| 194 | self.model_inputs.seq_lens_this_time = self.model_inputs["seq_lens_this_time_buffer"] |
| 195 | |
| 196 | def initialize_kv_cache(self, main_model_num_blocks, profile: bool = False): |
| 197 | """ |
no outgoing calls