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

Code before:
            max_new_tokens = generation_config.max_new_tokens
            if max_length is not None:
                max_new_tokens = max_length - seq.input_len
            if seq.output_token_id[-1] == generation_config.eos_token_id or seq.output_len >= max_new_tokens:
                seq.mark_finished()

    def check_unfinished_seqs(self) -> bool:
        return self._has_waiting() or not self.running_list.is_empty()

    def total_requests_in_batch_bucket(self) -> int:
        return self.prefill_bb.current_batch_size + self.running_bb.current_batch_size

    def search_tokens(self, generation_config: GenerationConfig, logits):
        """
        Sample tokens for finished requests.
        """

        # NOTE: need to decide the granularity to process logits (sequence or batch)
        config_dict = generation_config.to_dict()
        # process repetition_penalty, no_repeat_ngram_size
        for type in ["repetition_penalty", "no_repeat_ngram_size"]:
            if type in config_dict and config_dict[type] is not None:
                if not self.prefill_bb.is_empty:
                    batch = self.prefill_bb
                else:
                    batch = self.running_bb
                logits = logit_processor(type, logits, config_dict[type], batch)

        # do logit processor
        if generation_config.do_sample:
            # process temperature, top_k, top_p
            for type in ["temperature", "top_k", "top_p"]:
                if type in config_dict and config_dict[type] is not None:
                    logits = logit_processor(type, logits, config_dict[type])

        # calculate probs
        probs = torch.softmax(logits, dim=-1, dtype=torch.float)
        logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)

        # sample the next tokens
        sample_tokens = self._sample(probs, logprobs, generation_config)
        return sample_tokens

    def append_next_tokens(self, sample_tokens: torch.Tensor):
        assert sample_tokens.dim() == 1
        n_elements = sample_tokens.size(0)
        if not self.prefill_bb.is_empty:
            assert (
                self.prefill_bb.current_batch_size == n_elements
            ), f"Incompatible size: {n_elements} tokens to append while prefill batch size {self.prefill_bb.current_batch_size}"
            self.prefill_bb.append_batch_tokens(sample_tokens)

Review comment (performance):
Better passing in the batch bucket, instead of retrieving from its state in the method. We'd like to decouple search-tokens (sample) method with `RequestHandler`.

Code after:
            if seq.output_token_id[-1] == generation_config.eos_token_id or seq.output_len >= max_new_tokens:
                seq.mark_finished()

    def check_unfinished_seqs(self) -> bool:
        return self._has_waiting() or not self.running_list.is_empty()

    def total_requests_in_batch_bucket(self) -> int:
        return self.prefill_bb.current_batch_size + self.running_bb.current_batch_size

    def search_tokens(self, generation_config: GenerationConfig, logits, cur_batch: BatchBucket):
        """
        Sample tokens for finished requests.
        """

        # NOTE: need to decide the granularity to process logits (sequence or batch)
        config_dict = generation_config.to_dict()
        # process repetition_penalty, no_repeat_ngram_size
        for type in ["repetition_penalty", "no_repeat_ngram_size"]:
            if type in config_dict and config_dict[type] is not None:
                logits = logit_processor(type, logits, config_dict[type], cur_batch)

        # do logit processor
        if generation_config.do_sample:
            # process temperature, top_k, top_p
            for type in ["temperature", "top_k", "top_p"]:
                if type in config_dict and config_dict[type] is not None:
                    logits = logit_processor(type, logits, config_dict[type])

        # calculate probs
        probs = torch.softmax(logits, dim=-1, dtype=torch.float)
        logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)

        # sample the next tokens
        sample_tokens = self._sample(probs, logprobs, generation_config)
        return sample_tokens

    def append_next_tokens(self, sample_tokens: torch.Tensor):
        assert sample_tokens.dim() == 1
        n_elements = sample_tokens.size(0)
        if not self.prefill_bb.is_empty:
            assert (
                self.prefill_bb.current_batch_size == n_elements
            ), f"Incompatible size: {n_elements} tokens to append while prefill batch size {self.prefill_bb.current_batch_size}"
            self.prefill_bb.append_batch_tokens(sample_tokens)
        else:
            assert (
                self.running_bb.current_batch_size == n_elements
            ), f"Incompatible size: {n_elements} tokens to append while running batch size {self.running_bb.current_batch_size}"
            self.running_bb.append_batch_tokens(sample_tokens)

    def update(self):