Method set_use_spec_dec

colossalai/inference/batch_bucket.py:149–155

Set batch bucket to use speculatvie decoding. This will notify the adjust the lengths of inputs during modeling, and let the main model verifies tokens in parallel.

(self, num_tokens_to_verify: int = 5)

        return updated_block_ids

    def set_use_spec_dec(self, num_tokens_to_verify: int = 5) -> None:
        """Set batch bucket to use speculatvie decoding.
        This will notify the adjust the lengths of inputs during modeling,
        and let the main model verifies tokens in parallel.
        """
        self._use_spec_dec = True
        self._num_tokens_to_verify = num_tokens_to_verify

    def reset_use_spec_dec(self) -> None:
        """Reset the usage of speculative decoding for the batch bucket"""