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

Code before:
            **kwargs: Additional keyword arguments
        """
        assert "req_id" in kwargs, "req_id is required for async loading"
        request_configs = kwargs.get("request_configs")
        if request_configs is not None and len(request_configs) != 0:
            assert isinstance(request_configs, dict)

        tot_kv_size = 0
        chunks: list[tuple[CacheEngineKey, MemoryObj, int, int]] = []
        future = self.event_manager.pop_event(EventType.LOADING, kwargs["req_id"])

        memory_objs = future.result()
        memory_objs = [mm for m in memory_objs for mm in m]

        # NOTE(Jiayi): here we assume the retrieved memory_objs have
        # the same order as the lookup order.
        # TODO(Jiayi): hashing inside `process_tokens` can be skipped.

        # First, generate a list indicating which chunks are needed
        num_chunks = len(tokens) // self.config.chunk_size
        if mask is not None:
            chunk_mask_list = []
            for i in range(num_chunks):
                start_idx = i * self.config.chunk_size
                end_idx = start_idx + self.config.chunk_size
                # Check if all values in this chunk are True
                chunk_is_valid = all(mask[start_idx:end_idx])
                chunk_mask_list.append(chunk_is_valid)
        else:
            # If no mask, all chunks are needed
            chunk_mask_list = [True] * num_chunks

        # Now chunk_mask_list is like [F, F, T, ..., T]
        # where False means the chunk is in upper layer's cache and
        # doesn't need to be retrieved.
        # So we can use this list to filter out the chunks that are not needed.
        used_indices = set()
        for idx, (start, end, key) in enumerate(
            self.token_database.process_tokens(
                tokens=tokens,
                mask=None,
                request_configs=request_configs,
            )
        ):
            if not chunk_mask_list[idx]:
                continue

            assert isinstance(key, CacheEngineKey)
            memory_obj = memory_objs[idx]
            chunks.append((key, memory_obj, start, end))
            tot_kv_size += memory_obj.get_size()

Review comment (suggestion):
I think the main problem is here. In the example, we should get the 3rd element from the memory objs list. But the idx is from enumerate and will always be 0th. Let's find a way to use the correct indices then.

Code after:
            **kwargs: Additional keyword arguments
        """
        assert "req_id" in kwargs, "req_id is required for async loading"
        request_configs = kwargs.get("request_configs")
        if request_configs is not None and len(request_configs) != 0:
            assert isinstance(request_configs, dict)

        tot_kv_size = 0
        chunks: list[tuple[CacheEngineKey, MemoryObj, int, int]] = []
        future = self.event_manager.pop_event(EventType.LOADING, kwargs["req_id"])

        memory_objs = future.result()
        memory_objs = [mm for m in memory_objs for mm in m]

        # NOTE(Jiayi): here we assume the retrieved memory_objs have
        # the same order as the lookup order.
        # TODO(Jiayi): hashing inside `process_tokens` can be skipped.
        used_indices = set()
        for start, end, key in self.token_database.process_tokens(
            tokens=tokens,
            mask=mask,
            request_configs=request_configs,
        ):
            assert isinstance(key, CacheEngineKey)
            idx = start // self.config.chunk_size
            memory_obj = memory_objs[idx]
            chunks.append((key, memory_obj, start, end))
            tot_kv_size += memory_obj.get_size()
            ret_mask[start:end] = True
            used_indices.add(idx)

        # NOTE: free the memory objects that are not hit.
        for idx, unused_mem_obj in enumerate(memory_objs):
            if idx not in used_indices:
                unused_mem_obj.ref_count_down()

        return chunks, tot_kv_size

    def _process_tokens_internal(
        self,
        tokens,
        mask,
        ret_mask,
        **kwargs,
    ) -> tuple[list[tuple[CacheEngineKey, MemoryObj, int, int]], int]:
        """Process tokens and populate the reordered lists.

        This function is used to process tokens and populate the reordered lists.

        Args:
            tokens: Input tokens to process