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

Code before:
                reordered_ends,
                ret_mask,
            )

        if self.save_only_first_rank:
            self._broadcast_or_receive_memory_objs(
                reordered_keys,
                reordered_starts,
                reordered_ends,
                reordered_memory_objs,
                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.
        self.gpu_connector.batched_to_gpu(
            reordered_memory_objs, reordered_starts, reordered_ends, **kwargs
        )

        # TODO(Jiayi): Remove the following for loop with batched operations
        for key, memory_obj in zip(reordered_keys, reordered_memory_objs, strict=False):
            if self.remove_after_retrieve and (
                not self.save_only_first_rank or self.metadata.is_first_rank()
            ):
                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.


Review comment (suggestion):
```suggestion
            if self.remove_after_retrieve and not self.is_passive()
```

Code after:
            Should include KV cache specific information (e.g., paged KV buffer
            and the page tables).

        :return: the boolean mask indicating which tokens are retrieved. The
            length of the mask should be the same as the tokens. On CPU.

        :raises: ValueError if the number of Falses in the mask is not a
            multiple of the chunk size.
        """
        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")

        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} "