Repository: LMCache/LMCache
File: lmcache/v1/storage_backend/connector/base_connector.py (Python)

Code before:

    def __init__(
        self, config: LMCacheEngineConfig, metadata: Optional[LMCacheEngineMetadata]
    ):
        # TODO(chunxiaozheng): support layerwise here
        assert metadata is not None
        self.save_chunk_meta: bool = (
            config.extra_config is None
            or config.extra_config.get("save_chunk_meta", True)
            or config.use_layerwise
        )
        self.meta_shapes: list[torch.Size] = metadata.get_shapes()
        self.meta_dtypes: list[torch.dtype] = metadata.get_dtypes()
        self.meta_fmt: MemoryFormat = (
            MemoryFormat.KV_MLA_FMT if metadata.use_mla else MemoryFormat.KV_2LTD
        )
        self.full_chunk_size: int = get_size_bytes(self.meta_shapes, self.meta_dtypes)
        assert self.full_chunk_size % metadata.chunk_size == 0
        self.single_token_size = self.full_chunk_size // metadata.chunk_size
        logger.info(
            "init remote connector metadata info, shapes: %s, dtypes: %s, "
            "fmt: %s, full chunk size: %s, single token size: %s",
            self.meta_shapes,
            self.meta_dtypes,
            self.meta_fmt,
            self.full_chunk_size,
            self.single_token_size,
        )

    @NotAudit
    def reshape_partial_chunk(
        self,
        memory_obj: MemoryObj,
        bytes_read: int,
    ) -> MemoryObj:
        assert self.full_chunk_size is not None
        assert self.single_token_size is not None
        if (
            bytes_read == 0
            or bytes_read % self.single_token_size != 0
            or bytes_read > self.full_chunk_size
        ):
            raise ValueError(
                f"bytes_read: {bytes_read} is illegal, "
                f"single_token_size: {self.single_token_size}, "
                f"full_chunk_size: {self.full_chunk_size}"
            )

        if bytes_read == self.full_chunk_size:
            # full chunk, return directly
            return memory_obj

Review comment (performance):
Right now, since `__init__` is called before `register_kv_caches`, I want to know what the behavior of `get_shapes` and `get_dtypes` here is, and whether it will impact the behavior when using normal models

Code after:

    def __init__(
        self, config: LMCacheEngineConfig, metadata: Optional[LMCacheEngineMetadata]
    ):
        """
        Initialize some common attributes, which will be used in the subclasses.
        - `save_chunk_meta` is a flag to indicate whether to save the chunk meta info.
        - `meta_shapes` is a list of shapes of lmcache full chunk.
        - `meta_dtypes` is a list of dtypes of lmcache chunk.
        - `meta_fmt` is the memory format of the lmcache chunk.
        - `full_chunk_size` is the size of the lmcache full chunk.
        - `single_token_size` is the size of a single token.`

        Input:
            config: the lmcache engine config
            metadata: the lmcache engine metadata
        """
        # TODO(chunxiaozheng): support layerwise here
        assert metadata is not None
        self.save_chunk_meta: bool = (
            config.extra_config is None
            or config.extra_config.get("save_chunk_meta", True)
            or config.use_layerwise
        )
        self.meta_shapes: list[torch.Size] = metadata.get_shapes()
        self.meta_dtypes: list[torch.dtype] = metadata.get_dtypes()
        self.meta_fmt: MemoryFormat = (
            MemoryFormat.KV_MLA_FMT if metadata.use_mla else MemoryFormat.KV_2LTD
        )
        self.full_chunk_size: int = get_size_bytes(self.meta_shapes, self.meta_dtypes)
        assert self.full_chunk_size % metadata.chunk_size == 0
        self.single_token_size = self.full_chunk_size // metadata.chunk_size
        logger.info(
            "init remote connector metadata info, shapes: %s, dtypes: %s, "
            "fmt: %s, full chunk size: %s, single token size: %s",
            self.meta_shapes,
            self.meta_dtypes,
            self.meta_fmt,
            self.full_chunk_size,
            self.single_token_size,
        )

    @NotAudit
    def reshape_partial_chunk(
        self,
        memory_obj: MemoryObj,
        bytes_read: int,
    ) -> MemoryObj:
        assert self.full_chunk_size is not None
        assert self.single_token_size is not None
        if (