CacheMessager is used to send the cache data between the engine worker and the cache server.
| 453 | |
| 454 | |
| 455 | class CacheMessagerV1: |
| 456 | """ |
| 457 | CacheMessager is used to send the cache data between the engine worker and the cache server. |
| 458 | """ |
| 459 | |
| 460 | def __init__( |
| 461 | self, |
| 462 | splitwise_role, |
| 463 | transfer_protocol, |
| 464 | pod_ip, |
| 465 | engine_worker_queue_port, |
| 466 | local_data_parallel_id, |
| 467 | gpu_cache_kvs, |
| 468 | rank, |
| 469 | nranks, |
| 470 | num_layers, |
| 471 | gpu_id=0, |
| 472 | block_size=64, |
| 473 | rdma_port=None, |
| 474 | cache_dtype="bfloat16", |
| 475 | ): |
| 476 | """ |
| 477 | Initialize the CacheMessager object. |
| 478 | |
| 479 | Args: |
| 480 | splitwise_role (str): splitwise_role only can be 'prefill' or 'decode'. |
| 481 | transfer_protocol (str): support ipc and rdma |
| 482 | engine_worker_queue_port (int): engine_worker_queue port |
| 483 | gpu_cache_kvs (dict): GPU kv cache |
| 484 | rank (int): current rank |
| 485 | nranks (int): global rank number |
| 486 | num_layers (int): model layer number |
| 487 | gpu_id (int, optional): GPU ID |
| 488 | rdma_port (int, optional): RDMA port |
| 489 | |
| 490 | Returns: |
| 491 | None |
| 492 | """ |
| 493 | self.splitwise_role = splitwise_role |
| 494 | self.gpu_cache_kvs = gpu_cache_kvs |
| 495 | self.rank = rank |
| 496 | self.nranks = nranks |
| 497 | self.cache_dtype = cache_dtype |
| 498 | if not envs.FD_ENGINE_TASK_QUEUE_WITH_SHM: |
| 499 | address = (pod_ip, engine_worker_queue_port) |
| 500 | else: |
| 501 | address = f"/dev/shm/fd_task_queue_{engine_worker_queue_port}.sock" |
| 502 | self.engine_worker_queue = EngineWorkerQueue( |
| 503 | address=address, |
| 504 | is_server=False, |
| 505 | num_client=self.nranks, |
| 506 | client_id=self.rank, |
| 507 | local_data_parallel_id=local_data_parallel_id, |
| 508 | ) |
| 509 | self.block_size = block_size |
| 510 | transfer_protocol = transfer_protocol.split(",") |
| 511 | |
| 512 | logger.info( |