Repository: LMCache/LMCache
File: lmcache/v1/cache_engine.py (Python)

Code before:

        if self.save_only_first_rank and self.gpu_connector is not None:
            self.broadcast_stream = (
                self.gpu_connector.load_stream
                if hasattr(self.gpu_connector, "load_stream")
                else torch.cuda.Stream()
            )

        self.enable_controller = config.enable_controller

        # NOTE: Unix systems use fork by default
        multiprocessing.set_start_method("spawn", force=True)

        # avoid circular import
        # First Party
        from lmcache.v1.cache_controller import LMCacheWorker

        self.lmcache_worker: Optional[LMCacheWorker] = None
        lmcache_worker_ids = config.get_lmcache_worker_ids(
            metadata.use_mla, metadata.world_size
        )
        if (
            self.enable_controller
            and self.metadata.role != "scheduler"
            and (
                len(lmcache_worker_ids) == 0 or metadata.worker_id in lmcache_worker_ids
            )
        ):
            self.lmcache_worker = LMCacheWorker(config, metadata, self)

        self.async_loading = config.enable_async_loading
        self.event_manager = EventManager()

        self.storage_manager = StorageManager(
            config,
            metadata,
            # self.memory_allocator,
            event_manager=self.event_manager,
            lmcache_worker=self.lmcache_worker,
        )

        # HACK: remove this in the future
        # NOTE (Jiayi): This is currently used to support
        # dropping the kv cache from the buffer in PD backend
        # at decoder.
        self.remove_after_retrieve = config.enable_pd and config.pd_role == "receiver"

        self.use_layerwise = config.use_layerwise
        self.num_layers = metadata.kv_shape[0]
        self.fmt = None
        if self.use_layerwise:

Review comment (performance):
The condition `len(lmcache_worker_ids) == 0` is confusing. Should we set `len(lmcache_worker_ids) == number of total ranks` by default?

Code after:

        if self.save_only_first_rank and self.gpu_connector is not None:
            self.broadcast_stream = (
                self.gpu_connector.load_stream
                if hasattr(self.gpu_connector, "load_stream")
                else torch.cuda.Stream()
            )

        self.enable_controller = config.enable_controller

        # NOTE: Unix systems use fork by default
        multiprocessing.set_start_method("spawn", force=True)

        # avoid circular import
        # First Party
        from lmcache.v1.cache_controller import LMCacheWorker

        self.lmcache_worker: Optional[LMCacheWorker] = None
        lmcache_worker_ids = config.get_lmcache_worker_ids(
            metadata.use_mla, metadata.world_size
        )
        # lmcache_worker_ids is empty means start on all workers
        if (
            self.enable_controller
            and self.metadata.role != "scheduler"
            and (not lmcache_worker_ids or metadata.worker_id in lmcache_worker_ids)
        ):
            self.lmcache_worker = LMCacheWorker(config, metadata, self)

        self.async_loading = config.enable_async_loading
        self.event_manager = EventManager()

        self.use_layerwise = config.use_layerwise

        # TODO: support save_only_first_rank when use layerwise
        # if use_layerwise is True, all ranks will initialize the storage_manager
        # if save_only_first_rank is False, all ranks will initialize
        # the storage_manager
        # if save_only_first_rank is True, only the first rank and
        # lookup server workers will initialize the storage_manager
        self.storage_manager = None
        lookup_server_worker_ids = self.config.get_lookup_server_worker_ids(
            metadata.use_mla, metadata.world_size
        )
        if (
            self.lmcache_worker is not None
            or self.use_layerwise
            or not self.save_only_first_rank
            or self.metadata.is_first_rank()
            or len(lookup_server_worker_ids) == 0
            or self.metadata.worker_id in lookup_server_worker_ids