CacheMessager is used to send the cache data between the engine worker and the cache server.
| 117 | |
| 118 | |
| 119 | class CacheMessager: |
| 120 | """ |
| 121 | CacheMessager is used to send the cache data between the engine worker and the cache server. |
| 122 | """ |
| 123 | |
| 124 | def __init__( |
| 125 | self, |
| 126 | splitwise_role, |
| 127 | transfer_protocol, |
| 128 | pod_ip, |
| 129 | engine_worker_queue_port, |
| 130 | local_data_parallel_id, |
| 131 | gpu_cache_kvs, |
| 132 | rank, |
| 133 | nranks, |
| 134 | num_layers, |
| 135 | gpu_id=0, |
| 136 | rdma_port=None, |
| 137 | cache_dtype="bfloat16", |
| 138 | ): |
| 139 | """ |
| 140 | Initialize the CacheMessager object. |
| 141 | |
| 142 | Args: |
| 143 | splitwise_role (str): splitwise_role only can be 'prefill' or 'decode'. |
| 144 | transfer_protocol (str): support ipc and rdma |
| 145 | engine_worker_queue_port (int): engine_worker_queue port |
| 146 | gpu_cache_kvs (dict): GPU kv cache |
| 147 | rank (int): current rank |
| 148 | nranks (int): global rank number |
| 149 | num_layers (int): model layer number |
| 150 | gpu_id (int, optional): GPU ID |
| 151 | rdma_port (int, optional): RDMA port |
| 152 | |
| 153 | Returns: |
| 154 | None |
| 155 | """ |
| 156 | self.splitwise_role = splitwise_role |
| 157 | self.gpu_cache_kvs = gpu_cache_kvs |
| 158 | self.rank = rank |
| 159 | self.nranks = nranks |
| 160 | self.cache_dtype = cache_dtype |
| 161 | if not envs.FD_ENGINE_TASK_QUEUE_WITH_SHM: |
| 162 | address = (pod_ip, engine_worker_queue_port) |
| 163 | else: |
| 164 | address = f"/dev/shm/fd_task_queue_{engine_worker_queue_port}.sock" |
| 165 | self.engine_worker_queue = EngineWorkerQueue( |
| 166 | address=address, |
| 167 | is_server=False, |
| 168 | num_client=self.nranks, |
| 169 | client_id=self.rank, |
| 170 | local_data_parallel_id=local_data_parallel_id, |
| 171 | ) |
| 172 | transfer_protocol = transfer_protocol.split(",") |
| 173 | |
| 174 | logger.info(f"splitwise role: {splitwise_role}, {transfer_protocol}" f"rank: {rank}") |
| 175 | |
| 176 | # 1. initialize the cache_k_ptr_list and cache_v_ptr_list |