Initialize kv cache
(self, main_model_num_blocks, profile: bool = False)
| 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 | """ |
| 198 | Initialize kv cache |
| 199 | """ |
| 200 | self.num_gpu_blocks = int(main_model_num_blocks * self.speculative_config.num_gpu_block_expand_ratio) |
| 201 | self.cache_kvs = {} |
| 202 | |
| 203 | # Get kv cache dtype |
| 204 | cache_type = self.model_config.dtype |
| 205 | kv_cache_quant_type = None |
| 206 | if ( |
| 207 | self.quant_config |
| 208 | and hasattr(self.quant_config, "kv_cache_quant_type") |
| 209 | and self.quant_config.kv_cache_quant_type is not None |
| 210 | ): |
| 211 | cache_type = self._get_cache_type() |
| 212 | kv_cache_quant_type = self.quant_config.kv_cache_quant_type |
| 213 | |
| 214 | # Get kv cache shape |
| 215 | key_cache_shape, value_cache_shape = self.attn_backends[0].get_kv_cache_shape( |
| 216 | max_num_blocks=self.num_gpu_blocks, kv_cache_quant_type=kv_cache_quant_type |
| 217 | ) |
| 218 | if kv_cache_quant_type == "block_wise_fp8": |
| 219 | kv_cache_scale_shape = [key_cache_shape[0], key_cache_shape[1], key_cache_shape[2]] |
| 220 | local_rank = self.local_rank % self.parallel_config.tensor_parallel_size |
| 221 | |
| 222 | cache_ready_signal_data = np.zeros(shape=[self.parallel_config.tensor_parallel_size], dtype=np.int32) |
| 223 | cache_ready_signal = IPCSignal( |
| 224 | name="cache_ready_signal", |
| 225 | array=cache_ready_signal_data, |
| 226 | dtype=np.int32, |
| 227 | suffix=self.parallel_config.local_engine_worker_queue_port, |
| 228 | create=False, |
| 229 | ) |
| 230 | |
| 231 | # Check if gpu runner needs to create kv cache |
| 232 | # 1. During profiling, it creates its own kv cache. |
| 233 | # 2. If no need to profile, create kv cache if cache managers do not exist. |
| 234 | create_cache_tensor = profile or not ( |
| 235 | self.fd_config.cache_config.num_cpu_blocks > 0 |
| 236 | or self.fd_config.cache_config.kvcache_storage_backend |
| 237 | or self.fd_config.scheduler_config.splitwise_role != "mixed" |
| 238 | ) |
| 239 | |
| 240 | if not create_cache_tensor: |
| 241 | logger.info(f"Waiting for cache managers to create kv cache.. {cache_ready_signal.value}") |
| 242 | while cache_ready_signal.value[local_rank] != 1: |
| 243 | time.sleep(1) |
| 244 | logger.info(f"OK! Stop waiting. {cache_ready_signal.value}") |
| 245 | |
| 246 | logger.info(f"Initializing kv cache for all layers. {cache_ready_signal.value}") |
| 247 | |
| 248 | if not create_cache_tensor: |
| 249 | cache_kvs_list = [] |
| 250 | for i in range( |
| 251 | self.num_main_model_layers, |
| 252 | self.num_main_model_layers + self.model_config.num_hidden_layers, |
| 253 | ): |