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

Code before:
        if self.device is None:
            self.device = torch.cuda.current_device()
        if self.sequences_set is None:
            self.sequences_set = OrderedSet()
        if self.fd_inter_tensor is None:
            self.fd_inter_tensor = FDIntermTensors()

    def init_batch(self, seqs: List["Sequence"] = None):
        """
        Initializes inference batches by input sentence list.

        Args:
            seqs (List["Sequence"]): List of input sequence.
        """

        assert len(self.sequences_set) == 0, "Sequences set has been initialized."

        if seqs is not None:
            if not isinstance(seqs, list):
                seqs = [seqs]
            for seq in seqs:
                if seq in self.sequences_set:
                    logger.warning(f"The sequence(request_id {seq.request_id}) is already in sequences_set.")
                    continue

                self.sequences_set.add(seq)

    def init_fd_tensors(self):
        if not self.fd_inter_tensor.is_initialized:
            self.fd_inter_tensor.initialize(
                max_batch_size=self.max_batch_size,
                num_attn_heads=self.num_heads,
                kv_max_split_num=self.kv_max_split_num,
                head_dim=self.head_dim,
                dtype=self.dtype,
                device=self.device,
            )

    def get_block_table_tensor(self) -> None:
        tesnor_list = []
        block_table = None

        assert len(self.sequences_set) > 0, "Batch has not been initialized yet. Please initialize batch first."

        for seq in self.sequences_set:
            block_table = seq.block_table
            assert (
                block_table is not None
            ), f"The sequence(request_id {seq.request_id}) has not initialized the block_table."
            tesnor_list.append(seq.block_table)


Review comment (performance):
This function seems to be redundant as it is called multiple times to try to initialize the singleton. You might want to just initialize the singleton in cache manager, and then just left `fd_inter_tensor: FDIntermTensors = FDIntermTensors()` to point to the singleton. 

Code after:
        if self.device is None:
            self.device = torch.cuda.current_device()
        if self.sequences_set is None:
            self.sequences_set = OrderedSet()
        if self.fd_inter_tensor is None:
            self.fd_inter_tensor = FDIntermTensors()

    def init_batch(self, seqs: List["Sequence"] = None):
        """
        Initializes inference batches by input sentence list.

        Args:
            seqs (List["Sequence"]): List of input sequence.
        """

        if seqs is not None:
            if not isinstance(seqs, list):
                seqs = [seqs]
            for seq in seqs:
                if seq in self.sequences_set:
                    logger.warning(f"The sequence(request_id {seq.request_id}) is already in sequences_set.")
                    continue

                self.sequences_set.add(seq)

    def init_fd_tensors(self):
        if not self.fd_inter_tensor.is_initialized:
            self.fd_inter_tensor.initialize(
                max_batch_size=self.max_batch_size,
                num_attn_heads=self.num_heads,
                kv_max_split_num=self.kv_max_split_num,
                head_dim=self.head_dim,
                dtype=self.dtype,
                device=self.device,
            )

    def get_block_table_tensor(self) -> None:
        tesnor_list = []
        block_table = None

        assert len(self.sequences_set) > 0, "Batch has not been initialized yet. Please initialize batch first."

        for seq in self.sequences_set:
            block_table = seq.block_table
            assert (
                block_table is not None
            ), f"The sequence(request_id {seq.request_id}) has not initialized the block_table."
            tesnor_list.append(seq.block_table)

        block_table = torch.stack(tesnor_list)
        return block_table