nano-engine-step

nano-vllm/nanovllm/engine/llm_engine.py · bb823b3 · L49–90 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

49    def step(self):
50        seqs, is_prefill = self.scheduler.schedule()
51        num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs)
52        token_ids = self.model_runner.call("run", seqs, is_prefill)
53        self.scheduler.postprocess(seqs, token_ids, is_prefill)
54        outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]
55        return outputs, num_tokens
56
57    def is_finished(self):
58        return self.scheduler.is_finished()
59
60    def generate(
61        self,
62        prompts: list[str] | list[list[int]],
63        sampling_params: SamplingParams | list[SamplingParams],
64        use_tqdm: bool = True,
65    ) -> list[str]:
66        pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True, disable=not use_tqdm)
67        if not isinstance(sampling_params, list):
68            sampling_params = [sampling_params] * len(prompts)
69        for prompt, sp in zip(prompts, sampling_params):
70            self.add_request(prompt, sp)
71        outputs = {}
72        prefill_throughput = decode_throughput = 0.
73        while not self.is_finished():
74            t = perf_counter()
75            output, num_tokens = self.step()
76            if num_tokens > 0:
77                prefill_throughput = num_tokens / (perf_counter() - t)
78            else:
79                decode_throughput = -num_tokens / (perf_counter() - t)
80            pbar.set_postfix({
81                "Prefill": f"{int(prefill_throughput)}tok/s",
82                "Decode": f"{int(decode_throughput)}tok/s",
83            })
84            for seq_id, token_ids in output:
85                outputs[seq_id] = token_ids
86                pbar.update(1)
87        pbar.close()
88        outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())]
89        outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
90        return outputs