Repository: huggingface/transformers
File: src/transformers/generation/continuous_batching/cache_manager.py (Python)

Code before:
    @abstractmethod
    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
        """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""

    def fork_blocks(
        self, parent_request_id: str, children_request_ids: list[str], block_manager: BlockManager
    ) -> tuple[list[int], list[int]]:
        """Forks the cache blocks of a (parent_request_id) to a list of (children_request_ids). To manage the blocks,
        the (block_manager) is used. When forking, the child's block are either shared with the parent, or they need to
        be copied from the parent. Hence we return two lists of blocks that need to be copied: one for the source and
        one for the destination."""

        # Sanity checks
        if parent_request_id not in self.block_table:
            raise ValueError(f"No block table found for request {parent_request_id}")
        # TODO: this check is good in the current context but it might be too much + slow things down
        for children_request_id in children_request_ids:
            if children_request_id in self.block_table:
                raise ValueError(f"Block table already exists for request {children_request_id}")

        # Actual forking
        parent_blocks = self.block_table[parent_request_id]
        list_forked_blocks, copy_src, copy_dst = block_manager.fork_blocks(
            parent_blocks=parent_blocks,
            num_forks=len(children_request_ids),
            shareable=self.uses_block_sharing,
            group_id=self._index,
        )
        if list_forked_blocks is None:
            raise ValueError(f"Failed to fork blocks for request {parent_request_id}")

        # Update the block table for all children requests
        for children_request_id, forked_blocks in zip(children_request_ids, list_forked_blocks):
            self.block_table[children_request_id] = forked_blocks
        return copy_src, copy_dst


class FullAttentionCacheAllocator(CacheAllocator):
    """Cache manager for a group of full attention layers."""

    def __init__(self, index: int, block_size: int, allow_block_sharing: bool) -> None:
        """Initializes the cache manager for a group of full attention layers.
        Args:
            - index: the index of the associated layer group
            - block_size: the size of the blocks in the cache
        """
        self._index = index
        self.uses_block_sharing = allow_block_sharing
        self.block_size = block_size
        self.block_table = {}


Review comment (suggestion):
```suggestion
        the block_manager is used. When forking, the child's block are either shared with the parent, or they need to
```

Code after:
        """Frees all blocks associated with a (request_id) using the (block_manager)."""
        if request_id in self.block_table:
            blocks_to_free = self.block_table.pop(request_id)
            block_manager.free_blocks(blocks_to_free, shareable=self.uses_block_sharing)
        else:
            logger.warning(
                f"CacheAllocator {self._index} attempted to free blocks for non-existent request_id: {request_id}"
            )

    @abstractmethod
    def get_read_indices(self, request_id: str, past_length: int, query_length: int) -> list[int]:
        """Returns the physical indices of where to read request_id's cache in the cache tensor."""

    @abstractmethod
    def get_write_indices(self, request_id: str, past_length: int, query_length: int) -> list[int]:
        """Returns the physical indices of where to write request_id's cache in the cache tensor."""

    @abstractmethod
    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
        """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""

    def fork_blocks(
        self, parent_request_id: str, children_request_ids: list[str], block_manager: BlockManager
    ) -> tuple[list[int], list[int]]:
        """Forks the cache blocks of a (parent_request_id) to a list of (children_request_ids). To manage the blocks,
        the (block_manager) is used. When forking, the child's block are either shared with the parent, or they need to
        be copied from the parent. Hence we return two lists of blocks that need to be copied: one for the source and
        one for the destination."""

        # Sanity checks
        if parent_request_id not in self.block_table:
            raise ValueError(f"No block table found for request {parent_request_id}")

        # Actual forking
        parent_blocks = self.block_table[parent_request_id]
        list_forked_blocks, copy_src, copy_dst = block_manager.fork_blocks(
            parent_blocks=parent_blocks,
            num_forks=len(children_request_ids),
            shareable=self.uses_block_sharing,
            group_id=self._index,
        )
        if list_forked_blocks is None:
            raise ValueError(f"Failed to fork blocks for request {parent_request_id}")

        # Update the block table for all children requests
        for children_request_id, forked_blocks in zip(children_request_ids, list_forked_blocks):
            if children_request_id in self.block_table:
                raise ValueError(f"Block table already exists for request {children_request_id}")
            self.block_table[children_request_id] = forked_blocks
        return copy_src, copy_dst