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

Code before:
                    or len(reordered_keys) != len(reordered_ends)
                    or len(reordered_keys) != len(reordered_memory_objs)
                ):
                    logger.warning(
                        f"reordered_keys,reordered_starts,reordered_ends,"
                        f"reordered_memory_objs have different lengths with"
                        f" {len(reordered_keys)}, {len(reordered_starts)}, "
                        f"{len(reordered_ends)}, {len(reordered_memory_objs)}"
                    )

                # broadcast starts and ends
                self.broadcast_object_fn(
                    {"starts": reordered_starts, "ends": reordered_ends},
                    self.metadata.first_rank,
                )

                # broadcast memory_objs
                for i in range(len(reordered_starts)):
                    memory_obj = reordered_memory_objs[i]

                    # broadcast metadata
                    metadata_dict = memory_obj.metadata.to_dict()
                    self.broadcast_object_fn(metadata_dict, self.metadata.first_rank)

                    # broadcast tensor
                    tensor_to_broadcast = memory_obj.tensor.to(
                        f"cuda:{self.metadata.worker_id}"
                    )
                    self.broadcast_fn(tensor_to_broadcast, self.metadata.first_rank)
            else:
                # receive key lists
                data = self.broadcast_object_fn(None, self.metadata.first_rank)
                if data is None:
                    logger.warning(f"rank={self.metadata.worker_id} received None data")
                else:
                    reordered_starts = data["starts"]
                    reordered_ends = data["ends"]

                    # receive tensors
                    for start, end in zip(
                        reordered_starts, reordered_ends, strict=False
                    ):
                        ret_mask[start:end] = True  # set the mask

                        # receive metadata
                        metadata_dict = self.broadcast_object_fn(
                            None, self.metadata.first_rank
                        )
                        if metadata_dict is None:
                            logger.warning(
                                f"rank={self.metadata.worker_id} received None "

Review comment (performance):
I wonder if it’s possible to broadcast after calling LMCache’s kernel to transfer `memory_obj` from CPU to GPU, instead of using `tensor.to`, since it’s faster.

Code after:

        :param **kwargs: The additional arguments for the storage backend which
            will be passed into the gpu_connector.

        return: A generator that yields Optional[torch.Tensor]. The tensor will
            be the boolean mask indicating which tokens are retrieved and will
            only be returned in the last iteration. In the first iteration,
            the generator retrieve the memory objects of the first layer from
            the storage backends. In the next iterations, it moves the KV cache
            of layer i from the memory objects (on CPU) to GPU and retrieves
            the memory objects of layer i+1 from the storage backends. In the
            last iteration, it moves the memory objects of the last layer to
            the GPU.
        """

        if mask is not None:
            num_required_tokens = torch.sum(mask).item()
        else:
            num_required_tokens = len(tokens)
        monitor_req_id = self.stats_monitor.on_retrieve_request(num_required_tokens)

        ret_mask = torch.zeros_like(tokens, dtype=torch.bool, device="cpu")

        starts = []
        ends = []
        keys = []

        tags = kwargs.get("tags")
        if tags is not None and len(tags) != 0:
            assert isinstance(tags, OrderedDict)
        for start, end, key in self.token_database.process_tokens(
            tokens=tokens,
            mask=mask,
            tags=tags,
        ):
            assert isinstance(key, CacheEngineKey)

            keys_multi_layer = key.split_layers(self.num_layers)

            # NOTE: Only check the first layer
            if not self.storage_manager.contains(keys_multi_layer[0]):
                break

            starts.append(start)
            ends.append(end)
            keys.append(keys_multi_layer)

            ret_mask[start:end] = True

        if keys:
            # Transpose the keys into layer major format