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

Code before:
                    location = self.storage_manager.contains(key)
                    if location is None:
                        # TODO(Jiayi): Need to refactor P2P as a storage backend to
                        # clean up the following code.
                        if self.enable_p2p:
                            future_memory_obj = asyncio.run_coroutine_threadsafe(
                                self.distributed_server.issue_get(key),
                                self.distributed_loop,
                            )
                            memory_obj = future_memory_obj.result()
                            if memory_obj:
                                reordered_keys.append(key)
                                reordered_memory_objs.append(memory_obj)
                                reordered_starts.append(start)
                                reordered_ends.append(end)
                                ret_mask[start:end] = True
                            else:
                                # NOTE: break for P2P retrieve KV because of no required
                                # memory obj
                                break
                            continue
                        break

                    # NOTE: Here we make the assumption that the underlying
                    # storage backend support pin operation, and the memory
                    # object is already pinned in the storage backend.
                    ret_mask[start:end] = True

                    if location not in key_mapping:
                        key_mapping[location] = [key]
                        start_mapping[location] = [start]
                        end_mapping[location] = [end]
                        continue

                assert location is not None

                key_mapping[location].append(key)
                start_mapping[location].append(start)
                end_mapping[location].append(end)

            # TODO(Jiayi): We can parallelize the retrieval from
            # different storage backends.
            for location, keys in key_mapping.items():
                memory_objs = self.storage_manager.batched_get(
                    keys=keys,
                    storage_backend_name=location,
                )
                reordered_memory_objs.extend(memory_objs)
                reordered_keys.extend(keys)
                reordered_starts.extend(start_mapping[location])
                reordered_ends.extend(end_mapping[location])

Review comment (suggestion):
The git diff is not so smart, the true diff for this code block is that I just only add a `if not self.save_only_first_rank or self.metadata.is_first_rank():` outside, so i have to add a tab for each line.

Code after:
        reordered_chunks: List[Tuple[CacheEngineKey, MemoryObj, int, int]] = []
        if not self._is_passive():
            reordered_chunks = self._process_tokens_internal(
                tokens,
                mask,
                ret_mask,
                **kwargs,
            )
        if self.save_only_first_rank:
            self._broadcast_or_receive_memory_objs(
                reordered_chunks,
                ret_mask,
            )

        # NOTE(Jiayi): memory_obj doesn't have to be a pinned
        # cpu tensor for the sake of performance.
        # For example, disk->gpu is faster than disk->cpu->gpu.
        # RDMA is another example.
        if len(reordered_chunks) > 0:
            _, memory_objs, starts, ends = zip(*reordered_chunks, strict=False)
            self.gpu_connector.batched_to_gpu(
                list(memory_objs), list(starts), list(ends), **kwargs
            )

        # TODO(Jiayi): Remove the following for loop with batched operations
        for key, memory_obj, _, _ in reordered_chunks:
            if self.remove_after_retrieve and not self._is_passive():
                self.storage_manager.remove(key)
            memory_obj.ref_count_down()

        retrieved_tokens = torch.sum(ret_mask)
        self.stats_monitor.on_retrieve_finished(monitor_req_id, retrieved_tokens)
        logger.info(
            f"Retrieved {retrieved_tokens} "
            f"out of {num_required_tokens} "
            f"out of total {len(tokens)} tokens"
        )
        return ret_mask

    @_lmcache_nvtx_annotate
    @torch.inference_mode()
    def retrieve_layer(
        self,
        tokens: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Generator[Optional[torch.Tensor], None, None]:
        """
        Retrieve the KV cache in a layerwise manner.

        :param torch.Tensor tokens: The tokens of the corresponding KV caches.