Repository: hpcaitech/ColossalAI
File: colossalai/inference/kv_cache/kvcache_manager.py (Python)

Code before:
        Currently, allocations and updates are done at granularity of a single sequence.
        That is, the block table should be a 1D tensor of shape [max_blocks_per_sequence].
        And it's possible to have a batch of sequences with different lengths of block tables.
    """

    def __init__(self, config: InferenceConfig, model_config: PretrainedConfig) -> None:
        self.logger = get_dist_logger(__name__)
        self.device = get_current_device()

        # Parallel settings
        self.tp_size = config.tp_size
        # Model settings
        self.dtype = config.dtype
        self.elem_size_in_bytes = torch.tensor([], dtype=self.dtype).element_size()
        self.num_layers = model_config.num_hidden_layers
        self.head_num = model_config.num_attention_heads
        self.head_size = model_config.hidden_size // self.head_num
        if hasattr(model_config, "num_key_value_heads"):
            self.kv_head_num = model_config.num_key_value_heads
        else:
            self.kv_head_num = self.head_num

        assert (
            self.kv_head_num % self.tp_size == 0
        ), f"Cannot shard {self.kv_head_num} heads with tp size {self.tp_size}"
        self.kv_head_num //= self.tp_size
        self.beam_width = config.beam_width
        self.max_batch_size = config.max_batch_size
        self.max_input_length = config.max_input_len
        self.max_output_length = config.max_output_len
        # Cache block settings
        self.block_size = config.block_size
        # NOTE: `num_blocks` is not prompted, but evaluated from the maximum input/output length, and the maximum batch size
        self.max_blocks_per_sequence = (
            self.max_input_length + self.max_output_length + self.block_size - 1
        ) // self.block_size
        self.num_blocks = self.max_blocks_per_sequence * self.max_batch_size * self.beam_width

        # Physical cache allocation
        if config.use_cuda_kernel:
            x = 16 // torch.tensor([], dtype=config.dtype).element_size()
            kalloc_shape = (self.num_blocks, self.kv_head_num, self.head_size // x, self.block_size, x)
            valloc_shape = (self.num_blocks, self.kv_head_num, self.block_size, self.head_size)
            self.logger.info(
                f"Allocating K cache with shape: {kalloc_shape}, V cache with shape: {valloc_shape} consisting of {self.num_blocks} blocks."
            )
            self._kv_caches = self._init_device_caches(kalloc_shape, valloc_shape)
        else:
            alloc_shape = (self.num_blocks, self.kv_head_num, self.block_size, self.head_size)
            self.logger.info(f"Allocating KV cache with shape: {alloc_shape} consisting of {self.num_blocks} blocks.")
            self._kv_caches = self._init_device_caches(alloc_shape, alloc_shape)

Review comment (performance):
Branching can be skipped: `self.kv_cache_dtype = config.kv_cache_dtype`, as the following comment indicates

Code after:

        For a batch of sequences, the block tables after allocation might be:
        |  0 |  1 |  2 | -1 | -1 | -1 |
        |  3 |  4 |  5 |  6 |  7 | -1 |
        |  8 |  9 | 10 | 11 | -1 | -1 |
        | 12 | 13 | 14 | 15 | -1 | -1 |
        where 16 logical cache blocks are allocated and the same number of physical cache blocks will be used in kernels.

        Currently, allocations and updates are done at granularity of a single sequence.
        That is, the block table should be a 1D tensor of shape [max_blocks_per_sequence].
        And it's possible to have a batch of sequences with different lengths of block tables.
    """

    def __init__(self, config: InferenceConfig, model_config: PretrainedConfig, verbose: bool = False) -> None:
        self.logger = get_dist_logger(__name__)
        self.device = get_current_device()

        # Parallel settings
        self.tp_size = config.tp_size
        # Model settings
        self.dtype = config.dtype

        if config.kv_cache_dtype is None:
            self.kv_cache_dtype = config.dtype
        else:
            self.kv_cache_dtype = config.kv_cache_dtype

        self.elem_size_in_bytes = torch.tensor([], dtype=self.dtype).element_size()
        self.num_layers = get_model_config_attr(model_config, "num_hidden_layers")
        self.head_num = get_model_config_attr(model_config, "num_attention_heads")
        self.head_size = get_model_config_attr(model_config, "hidden_size") // self.head_num

        if hasattr(config, "num_key_value_heads"):
            self.kv_head_num = getattr(config, "num_key_value_heads")
        elif hasattr(config, "attribute_map") and hasattr(config, config.attribute_map["num_key_value_heads"]):
            self.kv_head_num = getattr(config, config.attribute_map["num_key_value_heads"])
        else:
            self.kv_head_num = self.head_num

        assert (
            self.kv_head_num % self.tp_size == 0
        ), f"Cannot shard {self.kv_head_num} heads with tp size {self.tp_size}"
        self.kv_head_num //= self.tp_size
        self.beam_width = config.beam_width
        self.max_batch_size = config.max_batch_size
        self.max_input_length = config.max_input_len
        self.max_output_length = config.max_output_len
        # Cache block settings
        self.block_size = config.block_size
        # NOTE: `num_blocks` is not prompted, but evaluated from the maximum input/output length, and the maximum batch size
        self.max_blocks_per_sequence = (