Method step_

colossalai/inference/core/rpc_engine.py:257–271

(self, input_token_ids, input_meta_data: InputMetaData)

        return input_ids.tolist(), input_meta_data

    async def step_(self, input_token_ids, input_meta_data: InputMetaData):
        assert len(self.workers) == self.tp_size, "init workers first"

        init_tasks = [
            self.async_parallel_wrapper(
                worker.execute_model_forward,
                input_token_ids,
                input_meta_data.to_rpc_param(),
                self.generation_config_dict,
            )
            for worker in self.workers
        ]
        ret = await asyncio.gather(*init_tasks)

        return ret[0]

    def step(self) -> List[str]:
        batch = self.request_handler.schedule()