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

Code before:
            out = torch.empty([self.current_batch_size], dtype=torch.long)
            for seq_id, index_in_b in self._sequences_indexes.items():
                seq: Sequence = self._sequences_dict[seq_id]
                out[index_in_b] = seq.output_token_id[-1]
            return out.to(device=self.device)

    # For compatibility
    def get_block_table_tensor(self) -> torch.Tensor:
        assert self.is_compact  # Debug usage
        block_table = self.block_tables[: self.current_batch_size]
        return block_table.to(device=self.device)

    # For compatibility
    def get_sequence_lengths(self) -> torch.Tensor:
        assert self.is_compact  # Debug usage
        sequence_lengths = self.seq_lengths[: self.current_batch_size]
        return sequence_lengths.to(device=self.device)

    def get_batch_token_ids(self) -> List[torch.LongTensor]:
        assert self.is_compact  # Debug usage
        out = []
        for seq_id, _ in self._sequences_indexes.items():
            seq: Sequence = self._sequences_dict[seq_id]
            out.append(torch.tensor(seq.input_token_id + seq.output_token_id, device=self.device))
        return out

    # For compatibility
    @property
    def fd_inter_tensor(self) -> None:
        assert self.fd_interm_tensor is not None, "fd_interm_tensor is not provided"
        return self.fd_interm_tensor

    def __repr__(self) -> str:
        return f"(sequences_dict={self._sequences_dict}, is_prompts={self.is_prompts})"


Review comment (performance):
You could pass input tokens ids to logit processor, instead of generating from batch bucket every time, which is slow. 

How to get input token ids from a flattened token ids with different lengths of sequences: Grab data which costs less to be retrieved, for example, you could provide sequence lengths to logit processor as well, **or**, pass the sequence list `BatchBucket.seqs_li` to logit processor. Either one costs less than the method and the getter method you added above

Code after:
                seq.output_len > 0 for seq in self._sequences_dict.values()
            ), "Sequence stage (Prefill/Decoding) must be the same in the batch"
            assert self.is_compact, "BatchBucket is not compact"
            out = torch.empty([self.current_batch_size], dtype=torch.long)
            for seq_id, index_in_b in self._sequences_indexes.items():
                seq: Sequence = self._sequences_dict[seq_id]
                out[index_in_b] = seq.output_token_id[-1]
            return out.to(device=self.device)

    # For compatibility
    def get_block_table_tensor(self) -> torch.Tensor:
        assert self.is_compact  # Debug usage
        block_table = self.block_tables[: self.current_batch_size]
        return block_table.to(device=self.device)

    # For compatibility
    def get_sequence_lengths(self) -> torch.Tensor:
        assert self.is_compact  # Debug usage
        sequence_lengths = self.seq_lengths[: self.current_batch_size]
        return sequence_lengths.to(device=self.device)

    # For compatibility
    @property
    def fd_inter_tensor(self) -> None:
        assert self.fd_interm_tensor is not None, "fd_interm_tensor is not provided"
        return self.fd_interm_tensor

    def __repr__(self) -> str:
        return f"(sequences_dict={self._sequences_dict}, is_prompts={self.is_prompts})"